"""Orchestrator — manages multiple agent engines and delegation.""" from __future__ import annotations import asyncio from contextlib import AsyncExitStack from pathlib import Path from typing import Any from loguru import logger from xtrm_agent.bus import AgentMessage, InboundMessage, MessageBus, OutboundMessage from xtrm_agent.cache import ResponseCache from xtrm_agent.classifier import QueryClassifier from xtrm_agent.config import Config, AgentFileConfig, parse_agent_file from xtrm_agent.cost import BudgetConfig, CostTracker from xtrm_agent.engine import Engine from xtrm_agent.llm.anthropic import AnthropicProvider from xtrm_agent.llm.litellm import LiteLLMProvider from xtrm_agent.llm.provider import LLMProvider from xtrm_agent.router import Router from xtrm_agent.tools.approval import ApprovalEngine from xtrm_agent.tools.builtin import register_builtin_tools from xtrm_agent.tools.delegate import DelegateTool from xtrm_agent.tools.mcp_client import connect_mcp_servers from xtrm_agent.tools.registry import ToolRegistry class Orchestrator: """Creates and manages multiple agent engines.""" def __init__(self, config: Config, interactive: bool = True) -> None: self.config = config self.bus = MessageBus() self.interactive = interactive self._engines: dict[str, Engine] = {} self._delegate_tools: dict[str, DelegateTool] = {} self._agent_configs: dict[str, AgentFileConfig] = {} self._mcp_stack = AsyncExitStack() self._running = False self._cache: ResponseCache | None = None self._cost_tracker: CostTracker | None = None # Channel defaults for routing channel_defaults = {} if config.channels.cli.default_agent: channel_defaults["cli"] = config.channels.cli.default_agent if config.channels.discord.default_agent: channel_defaults["discord"] = config.channels.discord.default_agent self.router = Router( agent_names=list(config.agents.keys()), channel_defaults=channel_defaults, ) async def setup(self) -> None: """Load agent definitions and create engines.""" workspace = Path(self.config.tools.workspace).resolve() workspace.mkdir(parents=True, exist_ok=True) # Initialize shared response cache self._cache = ResponseCache( db_path=workspace / "cache.db", ttl=self.config.performance.cache_ttl, ) await self._cache.setup() # Initialize cost tracker budget = BudgetConfig( daily_limit_usd=self.config.performance.daily_budget_usd, monthly_limit_usd=self.config.performance.monthly_budget_usd, ) self._cost_tracker = CostTracker( db_path=workspace / "costs.db", budget=budget, ) await self._cost_tracker.setup() # Initialize query classifier classifier = QueryClassifier(model_map=self.config.performance.model_routing) # Parse all agent definitions for agent_name, agent_path in self.config.agents.items(): p = Path(agent_path) if not p.is_absolute(): p = Path.cwd() / p if p.exists(): agent_cfg = parse_agent_file(p) else: logger.warning(f"Agent file not found: {p} — using defaults") agent_cfg = AgentFileConfig() agent_cfg.name = agent_cfg.name or agent_name self._agent_configs[agent_name] = agent_cfg # Build shared tool registry, then create per-agent registries global_registry = ToolRegistry() register_builtin_tools(global_registry, workspace) # Connect MCP servers await self._mcp_stack.__aenter__() await connect_mcp_servers(self.config.mcp_servers, global_registry, self._mcp_stack) # Create fallback provider (LiteLLM with a cheap model) fallback_model = self.config.performance.fallback_model fallback_provider = LiteLLMProvider(model=fallback_model) if fallback_model else None # Create one engine per agent agent_names = list(self._agent_configs.keys()) for agent_name, agent_cfg in self._agent_configs.items(): provider = self._create_provider(agent_cfg) approval = ApprovalEngine( auto_approve=self.config.tools.auto_approve, require_approval=self.config.tools.require_approval, interactive=self.interactive, ) # Filter tools for this agent if agent_cfg.tools: agent_registry = global_registry.filtered(agent_cfg.tools) else: agent_registry = global_registry # Add delegate tool if agent has "delegate" in its tool list other_agents = [n for n in agent_names if n != agent_name] if not agent_cfg.tools or "delegate" in agent_cfg.tools: delegate_tool = DelegateTool( bus=self.bus, from_agent=agent_name, available_agents=other_agents, timeout=self.config.orchestrator.delegation_timeout, ) agent_registry.register(delegate_tool) self._delegate_tools[agent_name] = delegate_tool engine = Engine( agent_config=agent_cfg, provider=provider, tools=agent_registry, approval=approval, cache=self._cache, cost_tracker=self._cost_tracker, classifier=classifier, fallback_provider=fallback_provider, ) self._engines[agent_name] = engine logger.info(f"Orchestrator ready: {len(self._engines)} agent(s)") def _create_provider(self, agent_cfg: AgentFileConfig) -> LLMProvider: """Create the appropriate LLM provider for an agent.""" provider_name = agent_cfg.provider if provider_name == "anthropic": model = agent_cfg.model or "claude-sonnet-4-5-20250929" return AnthropicProvider(model=model) # LiteLLM for everything else model = agent_cfg.model if not model: # Look up from config prov_cfg = self.config.llm.providers.get(provider_name) model = prov_cfg.model if prov_cfg else "deepseek/deepseek-chat-v3.1" return LiteLLMProvider(model=model) async def handle_message(self, msg: InboundMessage) -> str: """Route and process an inbound message.""" agent_name = self.router.resolve(msg) engine = self._engines.get(agent_name) if not engine: return f"Error: Agent '{agent_name}' not found" content = self.router.strip_mention(msg.content) if msg.content.startswith("@") else msg.content # Prepend attachment content so the LLM can see it if msg.attachments: parts: list[str] = [] for att in msg.attachments: parts.append(f"[Attached file: {att.filename}]\n{att.content}") parts.append(content) content = "\n\n".join(parts) logger.info(f"[{agent_name}] Processing: {content[:80]}") return await engine.run(content) async def handle_delegation(self, agent_msg: AgentMessage) -> None: """Handle an inter-agent delegation request.""" engine = self._engines.get(agent_msg.to_agent) if not engine: response = f"Error: Agent '{agent_msg.to_agent}' not found" else: logger.info( f"[{agent_msg.to_agent}] Delegation from {agent_msg.from_agent}: " f"{agent_msg.task[:80]}" ) response = await engine.run_delegation(agent_msg.task) # Resolve the delegation future in the delegate tool delegate_tool = self._delegate_tools.get(agent_msg.from_agent) if delegate_tool: delegate_tool.resolve(agent_msg.request_id, response) async def run_loop(self) -> None: """Main orchestrator loop — process inbound and agent messages.""" self._running = True logger.info("Orchestrator loop started") while self._running: # Check for inbound messages msg = await self.bus.consume_inbound(timeout=0.1) if msg: response = await self.handle_message(msg) await self.bus.publish_outbound( OutboundMessage( channel=msg.channel, chat_id=msg.chat_id, content=response, ) ) # Check for agent-to-agent messages agent_msg = await self.bus.consume_agent_message(timeout=0.1) if agent_msg: asyncio.create_task(self.handle_delegation(agent_msg)) async def stop(self) -> None: self._running = False if self._cache: await self._cache.close() if self._cost_tracker: await self._cost_tracker.close() await self._mcp_stack.aclose() logger.info("Orchestrator stopped") def get_agent_names(self) -> list[str]: return list(self._engines.keys()) def get_agent_info(self) -> list[dict[str, Any]]: """Get info about all registered agents.""" info = [] for name, cfg in self._agent_configs.items(): engine = self._engines.get(name) info.append( { "name": name, "provider": cfg.provider, "model": cfg.model or "(default)", "tools": engine.tools.names() if engine else [], "max_iterations": cfg.max_iterations, } ) return info