agentic_rag/runtime/orchestrator.py

"""ChatOrchestrator — central coordinator for all agent interactions."""

import uuid
from typing import AsyncIterator, Optional

from agentic_rag.data.models import AgentEvent, AgentInput, AgentOutput, Message
from agentic_rag.runtime.stream_bus import get_stream_bus
from agentic_rag.runtime.turn_manager import get_turn_manager
from agentic_rag.runtime.unified_context import UnifiedContext


class ChatOrchestrator:
    """The central conductor for the Agentic RAG system.

    Orchestrates the full lifecycle of each user request:
    1. Create/set up UnifiedContext
    2. Route to the appropriate ReActEngine via AgentRouter
    3. Execute the agent (streaming or non-streaming)
    4. Persist results to memory
    5. Emit events via StreamBus
    """

    def __init__(self):
        self.stream_bus = get_stream_bus()
        self.turn_manager = get_turn_manager()

    async def process(
        self,
        query: str,
        session_id: str = "",
        mode: str | None = None,
        has_media: bool = False,
        context: UnifiedContext | None = None,
    ) -> AgentOutput:
        """Process a query through the full pipeline (non-streaming)."""
        if context is None:
            context = UnifiedContext.create(session_id=session_id)

        # Start turn
        turn = self.turn_manager.start_turn(context.session_id)

        # Route and execute
        engine = await self._route(query, has_media, mode)
        input_data = AgentInput(query=query)

        try:
            output = await engine.run(input_data, turn_id=turn.turn_id)

            # Update turn stats
            self.turn_manager.update_usage(
                turn.turn_id,
                prompt_tokens=output.usage.get("prompt_tokens", 0),
                completion_tokens=output.usage.get("completion_tokens", 0),
            )
            self.turn_manager.finish_turn(turn.turn_id)

            # Persist to memory
            await self._persist_memory(context, query, output.final_answer)

            return output
        except Exception as e:
            self.turn_manager.finish_turn(turn.turn_id, error=str(e))
            raise

    async def process_stream(
        self,
        query: str,
        session_id: str = "",
        mode: str | None = None,
        has_media: bool = False,
        context: UnifiedContext | None = None,
    ) -> AsyncIterator[AgentEvent]:
        """Process a query with streaming output."""
        if context is None:
            context = UnifiedContext.create(session_id=session_id)

        turn = self.turn_manager.start_turn(context.session_id)
        engine = await self._route(query, has_media, mode)
        input_data = AgentInput(query=query)

        full_answer = ""
        try:
            async for event in engine.stream(input_data, turn_id=turn.turn_id):
                self.stream_bus.publish(turn.turn_id, event)
                if event.event_type.value == "text_delta":
                    full_answer += event.data.get("content", "")
                yield event

            self.turn_manager.finish_turn(turn.turn_id)
            await self._persist_memory(context, query, full_answer)

        except Exception as e:
            self.turn_manager.finish_turn(turn.turn_id, error=str(e))
            yield AgentEvent(event_type="error", data={"message": str(e)})

    async def _route(self, query: str, has_media: bool, mode: str | None):
        """Route query to the correct ReActEngine via AgentRouter."""
        from agentic_rag.services.llm.factory import get_llm
        from agentic_rag.orchestration.l1_tools.registry import get_tool_registry
        from agentic_rag.agent.router import AgentRouter

        llm = get_llm()
        tool_registry = get_tool_registry()
        self._ensure_tools(tool_registry)

        router = AgentRouter(llm, tool_registry)
        return await router.route(query=query, has_media=has_media, preferred_mode=mode)

    async def _persist_memory(self, context: UnifiedContext, query: str, answer: str) -> None:
        """Save the interaction to memory."""
        try:
            memory = context.memory_manager
            if memory:
                memory.add_message(context.session_id, Message.user(query))
                memory.add_message(context.session_id, Message.assistant(answer))
        except Exception:
            pass  # Memory persistence should never fail the main flow

    @staticmethod
    def _ensure_tools(registry) -> None:
        """Ensure built-in tools are registered."""
        if registry.tool_count == 0:
            from agentic_rag.orchestration.l1_tools.rag_tools import RAGSearchTool
            from agentic_rag.orchestration.l1_tools.web_tools import WebFetchTool, WebSearchTool
            from agentic_rag.orchestration.l1_tools.code_tools import CodeExecuteTool

            registry.register(RAGSearchTool())
            registry.register(WebSearchTool())
            registry.register(WebFetchTool())
            registry.register(CodeExecuteTool())


# Global instance
_orchestrator: Optional[ChatOrchestrator] = None


def get_orchestrator() -> ChatOrchestrator:
    global _orchestrator
    if _orchestrator is None:
        _orchestrator = ChatOrchestrator()
    return _orchestrator