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;
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
    }
}