agent.py

"""Multi-MCP Orchestrator Deep Agent.

Demonstrates connecting to multiple MCP servers and delegating domain-specific
tools to specialized SubAgents.
"""

import argparse
import asyncio
import os
from dotenv import load_dotenv
from rich.console import Console
from rich.panel import Panel

from deepagents import create_deep_agent
from deepagents.middleware.subagents import SubAgent

# Load environment
load_dotenv()
console = Console()


async def build_mcp_orchestrator(model_name: str = "claude-3-7-sonnet-latest"):
    """Initialize MCP servers and assemble the multi-agent orchestrator."""
    from langchain_mcp_adapters.client import MultiServerMCPClient

    base_dir = os.path.dirname(os.path.abspath(__file__))

    # Configure multiple MCP servers (stdio or SSE)
    mcp_config = {
        "sqlite": {
            "command": "uvx",
            "args": ["mcp-server-sqlite", "--db-path", os.path.join(base_dir, "incidents.db")],
            "transport": "stdio",
        },
        "github": {
            "command": "npx",
            "args": ["-y", "@modelcontextprotocol/server-github"],
            "transport": "stdio",
            "env": {"GITHUB_PERSONAL_ACCESS_TOKEN": os.environ.get("GITHUB_TOKEN", "")},
        },
    }

    # Initialize client and retrieve tools per server
    async with MultiServerMCPClient(mcp_config) as client:
        sqlite_tools = client.get_tools(server_name="sqlite")
        github_tools = client.get_tools(server_name="github")

        # Define specialized SubAgents
        subagents = [
            SubAgent(
                name="db-agent",
                description="Specialized in SQL database exploration and incident analytics.",
                system_prompt="You are a Database Specialist. Query tables to find metrics, error logs, and timestamps.",
                tools=sqlite_tools,
                mode="isolated",
            ),
            SubAgent(
                name="github-agent",
                description="Specialized in GitHub repository operations (issues, pull requests, releases).",
                system_prompt="You are a GitHub Ops Specialist. Inspect issues, diffs, and pull requests.",
                tools=github_tools,
                mode="isolated",
            ),
        ]

        # Assemble the top-level orchestrator agent
        orchestrator = create_deep_agent(
            model=model_name,
            subagents=subagents,
            memory=[os.path.join(base_dir, "AGENTS.md")],
            system_prompt=(
                "You are the Lead Operations Orchestrator. Delegate system-level tasks "
                "to the specialized subagents (db-agent, github-agent) and combine their findings."
            ),
        )

        return orchestrator


async def run_interactive(agent):
    """Run an interactive prompt loop."""
    console.print(
        Panel.fit(
            "[bold green]Multi-MCP Orchestrator Agent Initialized![/bold green]\n"
            "Domain subagents available: [cyan]db-agent[/cyan], [cyan]github-agent[/cyan]",
            title="DeepAgents",
        )
    )

    while True:
        try:
            user_input = console.input("\n[bold yellow]You:[/bold yellow] ")
            if user_input.strip().lower() in {"exit", "quit", "q"}:
                break
            if not user_input.strip():
                continue

            console.print("\n[bold blue]Orchestrator working...[/bold blue]")
            result = await agent.ainvoke(
                {"messages": [{"role": "user", "content": user_input}]}
            )

            final_message = result["messages"][-1].content
            console.print(Panel(final_message, title="[bold green]Response[/bold green]"))

        except KeyboardInterrupt:
            break
        except Exception as e:
            console.print(f"[bold red]Error:[/bold red] {e}")


def main():
    parser = argparse.ArgumentParser(description="Multi-MCP Orchestrator Deep Agent")
    parser.add_argument("--query", type=str, help="Single query to run non-interactively")
    parser.add_argument("--model", type=str, default="claude-3-7-sonnet-latest", help="Model to use")
    args = parser.parse_args()

    async def _runner():
        model = args.model
        if "ANTHROPIC_API_KEY" not in os.environ and "OPENAI_API_KEY" in os.environ:
            model = "openai:gpt-4o"

        agent = await build_mcp_orchestrator(model_name=model)

        if args.query:
            result = await agent.ainvoke({"messages": [{"role": "user", "content": args.query}]})
            print(result["messages"][-1].content)
        else:
            await run_interactive(agent)

    asyncio.run(_runner())


if __name__ == "__main__":
    main()