275 lines
8.5 KiB
Rust
275 lines
8.5 KiB
Rust
use super::message_builder::MessageBuilder;
|
|
use super::{AiRequest, StreamCallback};
|
|
use crate::client::AiClientConfig;
|
|
use crate::client::StreamChunk;
|
|
use crate::compact::CompactService;
|
|
use crate::embed::EmbedService;
|
|
use crate::error::Result;
|
|
use crate::perception::PerceptionService;
|
|
use crate::tool::registry::ToolRegistry;
|
|
use queue::MessageProducer;
|
|
|
|
/// Result from streaming AI response.
|
|
pub struct StreamResult {
|
|
pub content: String,
|
|
pub reasoning_content: String,
|
|
pub input_tokens: i64,
|
|
pub output_tokens: i64,
|
|
/// All chunks in arrival order — preserves ReAct multi-cycle ordering.
|
|
pub chunks: Vec<StreamChunk>,
|
|
}
|
|
|
|
/// Result from non-streaming AI response.
|
|
pub struct ProcessResult {
|
|
pub content: String,
|
|
pub input_tokens: i64,
|
|
pub output_tokens: i64,
|
|
}
|
|
|
|
/// Service for handling AI chat requests in rooms.
|
|
pub struct ChatService {
|
|
ai_base_url: Option<String>,
|
|
ai_api_key: Option<String>,
|
|
message_builder: MessageBuilder,
|
|
tool_registry: Option<ToolRegistry>,
|
|
}
|
|
|
|
impl ChatService {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
ai_base_url: None,
|
|
ai_api_key: None,
|
|
message_builder: MessageBuilder::new(),
|
|
tool_registry: None,
|
|
}
|
|
}
|
|
|
|
pub fn with_ai_client_config(mut self, config: AiClientConfig) -> Self {
|
|
self.ai_base_url = config.base_url.clone();
|
|
self.ai_api_key = Some(config.api_key.clone());
|
|
self
|
|
}
|
|
|
|
pub fn with_compact_service(mut self, compact_service: CompactService) -> Self {
|
|
self.message_builder = self.message_builder.with_compact_service(compact_service);
|
|
self
|
|
}
|
|
|
|
pub fn with_embed_service(mut self, embed_service: EmbedService) -> Self {
|
|
self.message_builder = self.message_builder.with_embed_service(embed_service);
|
|
self
|
|
}
|
|
|
|
pub fn with_perception_service(mut self, perception_service: PerceptionService) -> Self {
|
|
self.message_builder = self
|
|
.message_builder
|
|
.with_perception_service(perception_service);
|
|
self
|
|
}
|
|
|
|
pub fn with_tool_registry(mut self, registry: ToolRegistry) -> Self {
|
|
self.tool_registry = Some(registry);
|
|
self
|
|
}
|
|
|
|
/// Returns all registered tools as JSON tool definitions.
|
|
pub fn tools(&self) -> Vec<serde_json::Value> {
|
|
self.tool_registry
|
|
.as_ref()
|
|
.map(|r| r.to_openai_tools())
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
/// Build a RigToolSet from the registered tool registry.
|
|
///
|
|
/// This enables using the same tools with `RigAgentService` via rig's native Agent.
|
|
/// The context (db, cache, config, room_id, sender_id) is passed through to each
|
|
/// tool handler at creation time.
|
|
#[cfg(feature = "rig")]
|
|
pub fn rig_toolset(
|
|
&self,
|
|
db: db::database::AppDatabase,
|
|
cache: db::cache::AppCache,
|
|
config: config::AppConfig,
|
|
room_id: uuid::Uuid,
|
|
sender_id: Option<uuid::Uuid>,
|
|
project_id: uuid::Uuid,
|
|
) -> Option<crate::RigToolSet> {
|
|
self.tool_registry.as_ref().map(|registry| {
|
|
crate::RigToolSet::from_registry(
|
|
registry,
|
|
db,
|
|
cache,
|
|
config,
|
|
room_id,
|
|
sender_id,
|
|
project_id,
|
|
None,
|
|
None,
|
|
None,
|
|
std::sync::Arc::new(std::sync::Mutex::new(Vec::new())),
|
|
)
|
|
})
|
|
}
|
|
|
|
/// Get a reference to the underlying ToolRegistry.
|
|
pub fn tool_registry(&self) -> Option<&ToolRegistry> {
|
|
self.tool_registry.as_ref()
|
|
}
|
|
|
|
pub async fn build_room_optimized_context_text(
|
|
&self,
|
|
request: &AiRequest,
|
|
) -> Result<(String, Option<i64>)> {
|
|
self.message_builder
|
|
.build_room_optimized_context_text(request)
|
|
.await
|
|
}
|
|
|
|
/// Process AI request without streaming (tool-call loop with non-streaming API).
|
|
pub async fn process(&self, request: AiRequest) -> Result<ProcessResult> {
|
|
super::orchestrator::execute_orchestrated_process(
|
|
request,
|
|
&self.message_builder,
|
|
&self.tool_registry,
|
|
self.ai_base_url.clone(),
|
|
self.ai_api_key.clone(),
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Process AI request with streaming (tool-call loop with streaming API, incremental chunks).
|
|
pub async fn process_stream(
|
|
&self,
|
|
request: AiRequest,
|
|
on_chunk: StreamCallback,
|
|
) -> Result<StreamResult> {
|
|
super::orchestrator::execute_orchestrated_stream(
|
|
request,
|
|
on_chunk,
|
|
&self.message_builder,
|
|
&self.tool_registry,
|
|
self.ai_base_url.clone(),
|
|
self.ai_api_key.clone(),
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Process AI request for room context — direct execution path (bypasses orchestrator).
|
|
///
|
|
/// Room AI uses a fast single-agent loop: all tools available, no multi-agent delegation.
|
|
/// Merges `room_tools` (send_message, retract_message) into the base registry,
|
|
/// then runs `execute_process` / `execute_process_stream` directly.
|
|
pub async fn process_room(
|
|
&self,
|
|
request: AiRequest,
|
|
room_tools: ToolRegistry,
|
|
) -> Result<ProcessResult> {
|
|
let mut merged = self
|
|
.tool_registry
|
|
.clone()
|
|
.unwrap_or_default();
|
|
merged.merge(room_tools);
|
|
|
|
super::nonstreaming_execution::execute_process(
|
|
request,
|
|
&self.message_builder,
|
|
&Some(merged),
|
|
self.ai_base_url.clone(),
|
|
self.ai_api_key.clone(),
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Process AI request for room context with streaming — direct execution path.
|
|
///
|
|
/// Same as `process_room` but with streaming response. Bypasses orchestrator,
|
|
/// gives the room AI all tools (base + room) for fast single-agent execution.
|
|
pub async fn process_room_stream(
|
|
&self,
|
|
request: AiRequest,
|
|
on_chunk: StreamCallback,
|
|
room_tools: ToolRegistry,
|
|
) -> Result<StreamResult> {
|
|
let mut merged = self
|
|
.tool_registry
|
|
.clone()
|
|
.unwrap_or_default();
|
|
merged.merge(room_tools);
|
|
|
|
super::streaming_execution::execute_process_stream(
|
|
request,
|
|
on_chunk,
|
|
&self.message_builder,
|
|
&Some(merged),
|
|
self.ai_base_url.clone(),
|
|
self.ai_api_key.clone(),
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Process AI request via rig-based ReAct streaming loop.
|
|
pub async fn process_react<C, Fut>(
|
|
&self,
|
|
request: &AiRequest,
|
|
on_chunk: C,
|
|
) -> Result<(String, i64, i64)>
|
|
where
|
|
C: FnMut(crate::react::ReactStep) -> Fut + Send,
|
|
Fut: std::future::Future<Output = ()> + Send,
|
|
{
|
|
let Some(registry) = &self.tool_registry else {
|
|
return Err(crate::error::AgentError::Internal(
|
|
"no tool registry registered".into(),
|
|
));
|
|
};
|
|
super::react_execution::execute_process_react(
|
|
request,
|
|
on_chunk,
|
|
registry,
|
|
self.ai_base_url.clone(),
|
|
self.ai_api_key.clone(),
|
|
None,
|
|
None,
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Process AI request via rig-based ReAct streaming loop with room-specific tools.
|
|
///
|
|
/// Merges `room_tools` (e.g. `send_message`, `retract_message`) into the base
|
|
/// tool registry on-the-fly. The `room_preamble` is prepended to the default
|
|
/// system prompt to instruct the AI about room communication rules.
|
|
/// `message_producer` enables tools to publish events via the message queue.
|
|
pub async fn process_react_room<C, Fut>(
|
|
&self,
|
|
request: &AiRequest,
|
|
on_chunk: C,
|
|
room_tools: ToolRegistry,
|
|
room_preamble: Option<&str>,
|
|
message_producer: Option<MessageProducer>,
|
|
) -> Result<(String, i64, i64)>
|
|
where
|
|
C: FnMut(crate::react::ReactStep) -> Fut + Send,
|
|
Fut: std::future::Future<Output = ()> + Send,
|
|
{
|
|
let Some(registry) = &self.tool_registry else {
|
|
return Err(crate::error::AgentError::Internal(
|
|
"no tool registry registered".into(),
|
|
));
|
|
};
|
|
let mut merged = registry.clone();
|
|
merged.merge(room_tools);
|
|
super::react_execution::execute_process_react(
|
|
request,
|
|
on_chunk,
|
|
&merged,
|
|
self.ai_base_url.clone(),
|
|
self.ai_api_key.clone(),
|
|
room_preamble,
|
|
message_producer,
|
|
)
|
|
.await
|
|
}
|
|
}
|