agentpool 2.1.9__py3-none-any.whl

agentpool_server/a2a_server/server.py

@@ -0,0 +1,190 @@
+"""A2A server implementation for agentpool pool.
+
+This module provides a server that exposes all agents in an AgentPool
+via the A2A (Agent-to-Agent) protocol, with each agent accessible at its own route.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from agentpool.log import get_logger
+from agentpool_server.http_server import HTTPServer
+
+
+if TYPE_CHECKING:
+    from starlette.requests import Request
+    from starlette.responses import Response
+    from starlette.routing import Route
+
+    from agentpool import AgentPool
+
+
+logger = get_logger(__name__)
+
+DEFAULT_PORT = 8001
+
+
+class A2AServer(HTTPServer):
+    """A2A server for exposing pool agents on separate routes.
+
+    Provides a unified HTTP server that exposes all agents in the pool via
+    the A2A protocol. Each agent is accessible at `/{agent_name}` route.
+
+    Example:
+        ```python
+        pool = AgentPool(manifest)
+
+        server = A2AServer(pool, host="localhost", port=8001)
+
+        async with server:
+            async with server.run_context():
+                # Agents accessible at:
+                # POST http://localhost:8001/agent1
+                # POST http://localhost:8001/agent2
+                await do_other_work()
+        ```
+    """
+
+    def __init__(
+        self,
+        pool: AgentPool,
+        *,
+        name: str | None = None,
+        host: str = "localhost",
+        port: int = DEFAULT_PORT,
+        raise_exceptions: bool = False,
+    ) -> None:
+        """Initialize A2A server.
+
+        Args:
+            pool: AgentPool containing available agents
+            name: Optional server name (auto-generated if None)
+            host: Host to bind server to
+            port: Port to bind server to
+            raise_exceptions: Whether to raise exceptions during server start
+        """
+        super().__init__(pool, name=name, host=host, port=port, raise_exceptions=raise_exceptions)
+
+    async def get_routes(self) -> list[Route]:
+        """Get Starlette routes for A2A protocol.
+
+        Returns:
+            List of Route objects for each agent plus root listing endpoint
+        """
+        from starlette.routing import Route
+
+        routes: list[Route] = []
+
+        # Create route for each agent in the pool
+        for agent_name in self.pool.agents:
+
+            async def agent_handler(request: Request, agent_name: str = agent_name) -> Response:
+                """Handle A2A requests for a specific agent."""
+                from starlette.responses import JSONResponse
+
+                try:
+                    # Get the agent from pool
+                    pool_agent = self.pool.agents.get(agent_name)
+                    if pool_agent is None:
+                        return JSONResponse(
+                            {"error": f"Agent '{agent_name}' not found"},
+                            status_code=404,
+                        )
+
+                    # Get the underlying pydantic-ai agentlet and convert to A2A app
+                    agentlet = await pool_agent.get_agentlet(None, pool_agent.model_name, str)
+                    a2a_app = agentlet.to_a2a()
+
+                    # ASGI apps don't return a value, they write to send()
+                    await a2a_app(request.scope, request.receive, request._send)
+                    from starlette.responses import Response
+
+                    return Response()
+
+                except Exception as e:
+                    self.log.exception("Error handling A2A request", agent=agent_name)
+                    return JSONResponse(
+                        {"error": str(e)},
+                        status_code=500,
+                    )
+
+            # A2A protocol routes
+            routes.append(Route(f"/{agent_name}", agent_handler, methods=["POST"]))
+            routes.append(
+                Route(
+                    f"/{agent_name}/.well-known/agent-card.json",
+                    agent_handler,
+                    methods=["HEAD", "GET", "OPTIONS"],
+                )
+            )
+            routes.append(Route(f"/{agent_name}/docs", agent_handler, methods=["GET"]))
+
+            self.log.debug("Registered A2A routes", agent=agent_name, route=f"/{agent_name}")
+
+        # Add root endpoint that lists available agents
+        async def list_agents(request: Request) -> Response:
+            """List all available agents."""
+            from starlette.responses import JSONResponse
+
+            agent_list = [
+                {
+                    "name": name,
+                    "route": f"/{name}",
+                    "agent_card": f"/{name}/.well-known/agent-card.json",
+                    "docs": f"/{name}/docs",
+                    "model": agent.model_name,
+                }
+                for name, agent in self.pool.agents.items()
+            ]
+            return JSONResponse({
+                "agents": agent_list,
+                "count": len(agent_list),
+                "protocol": "A2A",
+                "version": "0.3.0",
+            })
+
+        routes.append(Route("/", list_agents, methods=["GET"]))
+
+        self.log.info(
+            "Created A2A routes",
+            agent_count=len(self.pool.agents),
+        )
+        return routes
+
+    def get_agent_url(self, agent_name: str) -> str:
+        """Get the URL for a specific agent.
+
+        Args:
+            agent_name: Name of the agent
+
+        Returns:
+            Full URL for the agent's A2A endpoint
+        """
+        return f"{self.base_url}/{agent_name}"
+
+    def get_agent_card_url(self, agent_name: str) -> str:
+        """Get the agent card URL for a specific agent.
+
+        Args:
+            agent_name: Name of the agent
+
+        Returns:
+            Full URL for the agent's A2A card endpoint
+        """
+        return f"{self.base_url}/{agent_name}/.well-known/agent-card.json"
+
+    def list_agent_routes(self) -> dict[str, dict[str, str]]:
+        """List all agent routes.
+
+        Returns:
+            Dictionary mapping agent names to their URLs
+        """
+        return {
+            name: {
+                "endpoint": self.get_agent_url(name),
+                "agent_card": self.get_agent_card_url(name),
+                "docs": f"{self.base_url}/{name}/docs",
+            }
+            for name in self.pool.agents
+        }

agentpool_server/a2a_server/storage.py

@@ -0,0 +1,81 @@
+"""A2A server implementation following BaseServer pattern."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from agentpool_server.a2a_server.a2a_types import TaskData
+
+
+if TYPE_CHECKING:
+    from agentpool.messaging.messages import ChatMessage
+    from agentpool_server.a2a_server.a2a_types import A2ARequest, TaskStatus
+
+
+@dataclass
+class SimpleStorage:
+    """Simple in-memory storage for A2A tasks."""
+
+    tasks: dict[str, TaskData] = field(default_factory=dict)
+    """Dictionary to store task data."""
+    contexts: dict[str, list[ChatMessage[Any]]] = field(default_factory=dict)
+    """Dictionary to store conversation history."""
+
+    async def store_task(self, task_id: str, task_data: A2ARequest) -> None:
+        """Store a task."""
+        self.tasks[task_id] = TaskData(
+            id=task_id,
+            status="submitted",
+            data=task_data,
+            result=None,
+            error=None,
+            context_id=task_data.get("params", {}).get("context_id"),
+        )
+
+    async def get_task(self, task_id: str) -> TaskData | None:
+        """Get task data."""
+        return self.tasks.get(task_id)
+
+    async def get_task_status(self, task_id: str) -> TaskStatus:
+        """Get task status."""
+        task = self.tasks.get(task_id)
+        if not task:
+            return {"status": "not_found", "result": None, "error": None}
+        return {
+            "status": task["status"],
+            "result": task["result"],
+            "error": task["error"],
+        }
+
+    async def update_task_status(
+        self,
+        task_id: str,
+        status: str,
+        result: dict[str, Any] | None = None,
+        error: str | None = None,
+    ) -> None:
+        """Update task status."""
+        if task_id in self.tasks:
+            self.tasks[task_id]["status"] = status
+            if result is not None:
+                self.tasks[task_id]["result"] = result
+            if error is not None:
+                self.tasks[task_id]["error"] = error
+
+    async def cancel_task(self, task_id: str) -> None:
+        """Cancel a task."""
+        await self.update_task_status(task_id, "cancelled")
+
+    async def load_context(self, context_id: str) -> list[ChatMessage[Any]]:
+        """Load conversation history for a context."""
+        return self.contexts.get(context_id, [])
+
+    async def update_context(
+        self,
+        context_id: str,
+        messages: list[ChatMessage[Any]],
+    ) -> None:
+        """Update conversation history for a context."""
+        if context_id:
+            self.contexts[context_id] = messages

agentpool_server/acp_server/__init__.py

@@ -0,0 +1,22 @@
+"""ACP (Agent Client Protocol) integration for agentpool."""
+
+from __future__ import annotations
+
+from agentpool_server.acp_server.server import ACPServer
+from agentpool_server.acp_server.acp_agent import AgentPoolACPAgent
+from agentpool_server.acp_server.session import ACPSession
+from agentpool_server.acp_server.session_manager import ACPSessionManager
+from agentpool_server.acp_server.converters import (
+    convert_acp_mcp_server_to_config,
+    from_acp_content,
+)
+
+
+__all__ = [
+    "ACPServer",
+    "ACPSession",
+    "ACPSessionManager",
+    "AgentPoolACPAgent",
+    "convert_acp_mcp_server_to_config",
+    "from_acp_content",
+]