234 lines
8.0 KiB
Rust
234 lines
8.0 KiB
Rust
use futures::StreamExt;
|
|
use models::rooms::room_ai;
|
|
use rig::agent::{AgentBuilder, MultiTurnStreamItem};
|
|
use rig::client::CompletionClient;
|
|
use rig::streaming::{StreamedAssistantContent, StreamingPrompt};
|
|
use sea_orm::*;
|
|
use uuid::Uuid;
|
|
|
|
use super::AiRequest;
|
|
use super::session_recording::record_ai_session;
|
|
use crate::client::AiClientConfig;
|
|
use crate::error::{AgentError, Result};
|
|
use crate::react::types::Action as ReactAction;
|
|
use crate::react::{DEFAULT_SYSTEM_PROMPT, ReactStep};
|
|
use crate::tool::{RecordingTool, registry::ToolRegistry};
|
|
|
|
pub async fn execute_process_react<C, Fut>(
|
|
request: &AiRequest,
|
|
mut on_chunk: C,
|
|
tool_registry: &ToolRegistry,
|
|
ai_base_url: Option<String>,
|
|
ai_api_key: Option<String>,
|
|
room_preamble: Option<&str>,
|
|
message_producer: Option<queue::MessageProducer>,
|
|
) -> Result<(String, i64, i64)>
|
|
where
|
|
C: FnMut(ReactStep) -> Fut + Send,
|
|
Fut: std::future::Future<Output = ()> + Send,
|
|
{
|
|
let base_url = ai_base_url.unwrap_or_else(|| "https://api.openai.com".into());
|
|
let api_key = ai_api_key.unwrap_or_default();
|
|
let client_config = AiClientConfig::new(api_key).with_base_url(base_url);
|
|
|
|
let db = request.db.clone();
|
|
let cache = request.cache.clone();
|
|
let cfg = request.config.clone();
|
|
let room_id = request.room.id;
|
|
let sender_uid = request.sender.uid;
|
|
let project_id = request.project.id;
|
|
let ai_model_id = request.model.id;
|
|
let ai_model_name = request.model.name.clone();
|
|
let sent_in_turn = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
|
|
let session_id = Uuid::now_v7();
|
|
let session_start = std::time::Instant::now();
|
|
let version_id = room_ai::Entity::find()
|
|
.filter(room_ai::Column::Room.eq(request.room.id))
|
|
.filter(room_ai::Column::Model.eq(request.model.id))
|
|
.one(&request.db)
|
|
.await
|
|
.ok()
|
|
.flatten()
|
|
.and_then(|r| r.version);
|
|
|
|
let mut tools: Vec<Box<dyn rig::tool::ToolDyn + 'static>> = Vec::new();
|
|
for def in tool_registry.definitions() {
|
|
let name = def.name.clone();
|
|
if let Some(handler) = tool_registry.get(&name) {
|
|
let adapter = crate::tool::RigToolAdapter::new(
|
|
handler.clone(),
|
|
def.clone(),
|
|
db.clone(),
|
|
cache.clone(),
|
|
cfg.clone(),
|
|
room_id,
|
|
Some(sender_uid),
|
|
project_id,
|
|
message_producer.clone(),
|
|
Some(ai_model_id),
|
|
Some(ai_model_name.clone()),
|
|
sent_in_turn.clone(),
|
|
);
|
|
tools.push(Box::new(RecordingTool::new(
|
|
Box::new(adapter),
|
|
db.clone(),
|
|
session_id,
|
|
sender_uid,
|
|
)));
|
|
}
|
|
}
|
|
|
|
let rig_client = client_config.build_rig_client();
|
|
let model = rig_client.completion_model(&request.model.name);
|
|
|
|
// General rules first (strong LLM attention), room context appended after
|
|
// so that output-format rules aren't buried behind long room preamble.
|
|
let preamble = match room_preamble {
|
|
Some(rp) => format!("{}\n{}", DEFAULT_SYSTEM_PROMPT, rp),
|
|
None => DEFAULT_SYSTEM_PROMPT.to_string(),
|
|
};
|
|
|
|
let agent = AgentBuilder::new(model)
|
|
.preamble(&preamble)
|
|
.tools(tools)
|
|
.default_max_turns(request.max_tool_depth)
|
|
.build();
|
|
|
|
let stream = agent
|
|
.stream_prompt(&request.input)
|
|
.with_history(Vec::<rig::completion::Message>::new())
|
|
.multi_turn(request.max_tool_depth)
|
|
.await;
|
|
|
|
tokio::pin!(stream);
|
|
|
|
let mut step_count = 0usize;
|
|
let mut final_content = String::new();
|
|
let mut total_input_tokens: i64 = 0;
|
|
let mut total_output_tokens: i64 = 0;
|
|
|
|
while let Some(item) = stream.next().await {
|
|
match item {
|
|
Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(text))) => {
|
|
step_count += 1;
|
|
let t = text.text;
|
|
on_chunk(ReactStep::Answer {
|
|
step: step_count,
|
|
answer: t.clone(),
|
|
})
|
|
.await;
|
|
final_content.push_str(&t);
|
|
}
|
|
Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Reasoning(
|
|
reasoning,
|
|
))) => {
|
|
let reasoning_text: String = reasoning
|
|
.content
|
|
.iter()
|
|
.filter_map(|c| match c {
|
|
rig::completion::message::ReasoningContent::Text { text, .. } => {
|
|
Some(text.as_str())
|
|
}
|
|
_ => None,
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join("");
|
|
if !reasoning_text.is_empty() {
|
|
step_count += 1;
|
|
on_chunk(ReactStep::Thought {
|
|
step: step_count,
|
|
thought: reasoning_text,
|
|
})
|
|
.await;
|
|
}
|
|
}
|
|
Ok(MultiTurnStreamItem::StreamAssistantItem(
|
|
StreamedAssistantContent::ReasoningDelta { reasoning, .. },
|
|
)) => {
|
|
if !reasoning.is_empty() {
|
|
step_count += 1;
|
|
on_chunk(ReactStep::Thought {
|
|
step: step_count,
|
|
thought: reasoning,
|
|
})
|
|
.await;
|
|
}
|
|
}
|
|
Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::ToolCall {
|
|
tool_call,
|
|
..
|
|
})) => {
|
|
step_count += 1;
|
|
let args: serde_json::Value = match &tool_call.function.arguments {
|
|
serde_json::Value::String(s) => {
|
|
serde_json::from_str(s).unwrap_or(serde_json::Value::Null)
|
|
}
|
|
v => v.clone(),
|
|
};
|
|
on_chunk(ReactStep::Action {
|
|
step: step_count,
|
|
action: ReactAction::new(&tool_call.function.name, args),
|
|
})
|
|
.await;
|
|
}
|
|
Ok(MultiTurnStreamItem::StreamUserItem(
|
|
rig::streaming::StreamedUserContent::ToolResult { tool_result, .. },
|
|
)) => {
|
|
step_count += 1;
|
|
let obs = tool_result_content_to_string(&tool_result.content);
|
|
on_chunk(ReactStep::Observation {
|
|
step: step_count,
|
|
observation: obs,
|
|
})
|
|
.await;
|
|
}
|
|
Ok(MultiTurnStreamItem::FinalResponse(resp)) => {
|
|
let usage = resp.usage();
|
|
total_input_tokens = usage.input_tokens as i64;
|
|
total_output_tokens = usage.output_tokens as i64;
|
|
}
|
|
Err(e) => {
|
|
let err_msg = format!("rig agent stream error: {}", e);
|
|
return Err(AgentError::OpenAi(err_msg));
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
let elapsed_ms = session_start.elapsed().as_millis() as i64;
|
|
record_ai_session(
|
|
&request.cache,
|
|
&request.db,
|
|
request.project.id,
|
|
request.sender.uid,
|
|
session_id,
|
|
request.room.id,
|
|
request.model.id,
|
|
version_id.unwrap_or_default(),
|
|
total_input_tokens,
|
|
total_output_tokens,
|
|
elapsed_ms,
|
|
)
|
|
.await;
|
|
|
|
Ok((final_content, total_input_tokens, total_output_tokens))
|
|
}
|
|
|
|
/// Extract text from rig's ToolResultContent, ignoring images.
|
|
fn tool_result_content_to_string(
|
|
content: &rig::one_or_many::OneOrMany<rig::completion::message::ToolResultContent>,
|
|
) -> String {
|
|
use rig::completion::message::ToolResultContent;
|
|
content
|
|
.iter()
|
|
.filter_map(|item| {
|
|
if let ToolResultContent::Text(t) = item {
|
|
Some(t.text.clone())
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join("\n")
|
|
}
|