#!/usr/bin/env python3
"""05 — Multi-agent orchestration.

Route user requests to specialized agents using an Orchestrator with a RuleRouter.

Requirements:
    pip install koan[llm]
    export OPENAI_API_KEY=sk-...
"""

from __future__ import annotations

import asyncio

from koan.agents import AgentConfig, ReactAgent
from koan.agents.orchestrator import Orchestrator
from koan.agents.router import RuleRouter
from koan.llm.client import LLMClient


async def main() -> None:
    async with LLMClient() as client:
        # Create specialized agents
        weather_agent = ReactAgent(
            "weather-expert",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="You are a weather expert. Answer weather questions concisely.",
            ),
        )

        code_agent = ReactAgent(
            "code-helper",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="You are a Python coding assistant. Write clear, concise code.",
            ),
        )

        general_agent = ReactAgent(
            "general",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="You are a helpful general assistant.",
            ),
        )

        # Create a router: keywords map to agent names
        router = RuleRouter(
            rules={
                "weather": "weather",
                "temperature": "weather",
                "forecast": "weather",
                "code": "code",
                "python": "code",
                "function": "code",
            },
        )

        # Build the orchestrator with a fallback agent
        orchestrator = Orchestrator(
            agents={
                "weather": weather_agent,
                "code": code_agent,
                "general": general_agent,
            },
            router=router,
            fallback="general",
        )

        # Test routing to different agents
        queries = [
            "What's the weather like in London?",
            "Write a Python function to sort a list",
            "Tell me a joke",
        ]

        for query in queries:
            print(f"\nQuery: {query}")
            result = await orchestrator.run(query)
            print(f"Routed to: {result.agent_name}")
            print(f"Response: {result.output[:150]}...")


if __name__ == "__main__":
    asyncio.run(main())
