gitdataai/libs/agent/chat/service.rs

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;

/// 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,
            )
        })
    }

    /// Get a reference to the underlying ToolRegistry.
    pub fn tool_registry(&self) -> Option<&ToolRegistry> {
        self.tool_registry.as_ref()
    }

    /// Process AI request without streaming (tool-call loop with non-streaming API).
    pub async fn process(&self, request: AiRequest) -> Result<ProcessResult> {
        super::nonstreaming_execution::execute_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::streaming_execution::execute_process_stream(
            request, on_chunk, &self.message_builder, &self.tool_registry,
            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(),
        ).await
    }
}