872ed24f0c
Inspired by zeroclaw's lightweight patterns for slow hardware: - Response cache (SQLite + SHA-256 keyed) to skip redundant LLM calls - History compaction — LLM-summarize old messages when history exceeds 50 - Query classifier routes simple/research queries to cheaper models - Credential scrubbing removes secrets from tool output before sending to LLM - Cost tracker with daily/monthly budget enforcement (SQLite) - Resilient provider with retry + exponential backoff + fallback provider - Approval engine gains session "always allow" and audit log Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
252 lines
9.7 KiB
Python
252 lines
9.7 KiB
Python
"""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
|