Custom Agents (Python)

import asyncio
import json
import os
from mcp import ClientSession
from mcp.client.sse import sse_client
# Example uses OpenAI, but works seamlessly with LangChain, AutoGen, etc.
from openai import AsyncOpenAI 

MCP_ENDPOINT = "https://a2a-backend-196084590575.europe-west1.run.app/mcp/sse"
API_KEY = os.getenv("AGENT_ROUTER_API_KEY", "your-api-key")

# 1. Your Agent's Built-in Capabilities
my_local_tools = [{
    "type": "function",
    "function": {
        "name": "local_database_query",
        "description": "Queries your local company database.",
        "parameters": {"type": "object", "properties": {"sql": {"type": "string"}}}
    }
}]

def execute_local_tool(name: str, args: dict) -> str:
    """Handle your agent's own functions"""
    if name == "local_database_query":
        return f"Executing {args['sql']} locally..."
    return "Unknown tool"

async def main():
    agent_llm = AsyncOpenAI()
    
    # 2. Connect to the AgentPlatform MCP Server
    # Note: Depending on the specific backend implementation, you pass auth headers here
    headers = { "Authorization": f"Bearer {API_KEY}" }
    async with sse_client(MCP_ENDPOINT, headers=headers) as (read_stream, write_stream):
        async with ClientSession(read_stream, write_stream) as mcp_session:
            await mcp_session.initialize()
            
            # 3. Fetch Specialist Tools from the Platform
            mcp_tools = await mcp_session.list_tools()
            external_tools = [{
                "type": "function",
                "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.inputSchema 
                }
            } for t in mcp_tools.tools]
            
            # 4. Combine your tools with the platform's tools
            all_tools = my_local_tools + external_tools

            # Import the full SYSTEM_PROMPT_SNIPPET here
            system_prompt = "You are an agent.\n"

            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": "Check our database, then find a researcher agent to analyze the data."}
            ]

            # 5. The Agent Loop
            while True:
                response = await agent_llm.chat.completions.create(
                    model="gpt-4o", messages=messages, tools=all_tools
                )
                ai_message = response.choices[0].message
                messages.append(ai_message)
                
                if not ai_message.tool_calls:
                    print("Agent Finished:", ai_message.content)
                    break

                for tool_call in ai_message.tool_calls:
                    name = tool_call.function.name
                    args = json.loads(tool_call.function.arguments)
                    
                    # 6. Route Execution: Local vs MCP
                    if name in [t["function"]["name"] for t in my_local_tools]:
                        # Execute your own code
                        result_text = execute_local_tool(name, args)
                    else:
                        # Forward to the Agent Router Platform
                        mcp_result = await mcp_session.call_tool(name, args)
                        result_text = mcp_result.content[0].text
                    
                    messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": result_text})

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