agentpool 2.1.9__py3-none-any.whl

agentpool_server/aggregating_server.py

@@ -0,0 +1,309 @@
+"""AggregatingServer for managing multiple servers with shared pool."""
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import AsyncExitStack
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal, Self
+
+from agentpool.log import get_logger
+from agentpool_server.base import BaseServer
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from types import TracebackType
+
+    from agentpool import AgentPool
+
+# Type-safe server status literals
+ServerStatus = Literal["not_initialized", "initialized", "running", "failed", "stopped"]
+
+
+@dataclass(frozen=True)
+class ServerInfo:
+    """Type-safe server information."""
+
+    name: str
+    server_type: type[BaseServer]
+    status: ServerStatus
+
+
+logger = get_logger(__name__)
+
+
+class AggregatingServer(BaseServer):
+    """Server that manages multiple protocol servers sharing one AgentPool.
+
+    Coordinates multiple server instances (AG-UI, A2A, MCP, etc.) as a single unit.
+    All servers share the same AgentPool and are started/stopped together.
+
+    This is useful when you want to expose the same pool of agents via multiple
+    protocols simultaneously (e.g., AG-UI for web clients, A2A for agent-to-agent
+    communication).
+
+    Example:
+        ```python
+        async with AgentPool(config) as pool:
+            servers = [
+                AGUIServer(pool, port=8002),
+                A2AServer(pool, port=8001),
+            ]
+            async with AggregatingServer(pool, servers) as agg:
+                async with agg.run_context():
+                    # Both servers running, sharing the same pool
+                    await asyncio.Event().wait()
+        ```
+    """
+
+    def __init__(
+        self,
+        pool: AgentPool[Any],
+        servers: Sequence[BaseServer],
+        *,
+        name: str | None = None,
+        raise_exceptions: bool = False,
+    ) -> None:
+        """Initialize aggregating server.
+
+        Args:
+            pool: AgentPool to be shared by all servers
+            servers: Sequence of servers to aggregate
+            name: Server name for logging (auto-generated if None)
+            raise_exceptions: Whether to raise exceptions during server start
+        """
+        if not servers:
+            msg = "At least one server must be provided"
+            raise ValueError(msg)
+
+        super().__init__(pool, name=name, raise_exceptions=raise_exceptions)
+
+        self.servers = list(servers)
+        self.exit_stack = AsyncExitStack()
+        self._initialized_servers: list[BaseServer] = []
+
+    async def __aenter__(self) -> Self:
+        """Initialize aggregating server and all child servers."""
+        await super().__aenter__()
+
+        self.log.info("Initializing aggregated servers", count=len(self.servers))
+
+        try:
+            for server in self.servers:
+                try:
+                    initialized_server = await self.exit_stack.enter_async_context(server)
+                    self._initialized_servers.append(initialized_server)
+                    self.log.info("Initialized server", server_name=server.name)
+                except Exception:
+                    self.log.exception("Failed to initialize server", server_name=server.name)
+                    if self.raise_exceptions:
+                        raise
+
+            if not self._initialized_servers:
+                msg = "No servers were successfully initialized"
+                raise RuntimeError(msg)  # noqa: TRY301
+
+            self.log.info(
+                "All servers initialized",
+                successful=len(self._initialized_servers),
+                failed=len(self.servers) - len(self._initialized_servers),
+            )
+
+        except Exception:
+            await self.exit_stack.aclose()
+            raise
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Cleanup all servers and base server resources."""
+        self.log.info("Shutting down aggregated servers")
+
+        await self.exit_stack.aclose()
+        self._initialized_servers.clear()
+
+        await super().__aexit__(exc_type, exc_val, exc_tb)
+
+        self.log.info("Aggregated servers shutdown complete")
+
+    async def _start_async(self) -> None:
+        """Start all initialized servers concurrently."""
+        if not self._initialized_servers:
+            self.log.warning("No initialized servers to start")
+            return
+
+        self.log.info("Starting aggregated servers", count=len(self._initialized_servers))
+
+        server_tasks: list[tuple[BaseServer, asyncio.Task[None] | None]] = []
+
+        for server in self._initialized_servers:
+            try:
+                server.start_background()
+                server_tasks.append((server, server._server_task))
+                self.log.info("Started server in background", server=server.name)
+            except Exception:
+                self.log.exception("Failed to start server", server=server.name)
+                if self.raise_exceptions:
+                    for started_server, _ in server_tasks:
+                        try:
+                            started_server.stop()
+                        except Exception:
+                            self.log.exception(
+                                "Error stopping server", server_name=started_server.name
+                            )
+                    raise
+
+        if not server_tasks:
+            msg = "No servers were successfully started"
+            raise RuntimeError(msg)
+
+        self.log.info(
+            "All servers started",
+            successful=len(server_tasks),
+            failed=len(self._initialized_servers) - len(server_tasks),
+        )
+
+        tasks = [task for _, task in server_tasks if task is not None]
+        if tasks:
+            await self._wait_for_tasks(tasks, server_tasks)
+
+    async def _wait_for_tasks(
+        self,
+        tasks: list[asyncio.Task[None]],
+        server_tasks: Sequence[tuple[BaseServer, asyncio.Task[None] | None]],
+    ) -> None:
+        """Wait for server tasks and handle shutdown."""
+        try:
+            done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+
+            for task in done:
+                if task.exception():
+                    self.log.error("Server task failed", error=task.exception())
+                else:
+                    self.log.info("Server task completed")
+
+            # Stop all other servers gracefully
+            for server, _ in server_tasks:
+                try:
+                    server.stop()
+                except Exception:
+                    self.log.exception("Error stopping server")
+
+            if pending:
+                for task in pending:
+                    task.cancel()
+                await asyncio.gather(*pending, return_exceptions=True)
+
+        except asyncio.CancelledError:
+            for server, _ in server_tasks:
+                try:
+                    server.stop()
+                except Exception:
+                    self.log.exception("Error stopping server during cancellation")
+            raise
+
+    def add_server(self, server: BaseServer) -> None:
+        """Add a server to the aggregation.
+
+        Args:
+            server: Server to add
+
+        Raises:
+            RuntimeError: If aggregating server is currently running
+        """
+        if self.is_running:
+            msg = "Cannot add server while aggregating server is running"
+            raise RuntimeError(msg)
+
+        self.servers.append(server)
+        self.log.info("Added server to aggregation", server=server.name)
+
+    def remove_server(self, server: BaseServer) -> None:
+        """Remove a server from the aggregation.
+
+        Args:
+            server: Server to remove
+
+        Raises:
+            RuntimeError: If aggregating server is currently running
+            ValueError: If server is not in aggregation
+        """
+        if self.is_running:
+            msg = "Cannot remove server while aggregating server is running"
+            raise RuntimeError(msg)
+
+        try:
+            self.servers.remove(server)
+            self.log.info("Removed server from aggregation", server=server.name)
+        except ValueError as e:
+            msg = f"Server {server.name} not found in aggregation"
+            raise ValueError(msg) from e
+
+    def get_server(self, name: str) -> BaseServer | None:
+        """Get a server by name from the aggregation.
+
+        Args:
+            name: Server name to find
+
+        Returns:
+            Server instance or None if not found
+        """
+        all_servers = self.servers + self._initialized_servers
+        for server in all_servers:
+            if server.name == name:
+                return server
+        return None
+
+    def list_servers(self) -> list[ServerInfo]:
+        """List all servers in the aggregation with their status.
+
+        Returns:
+            List of type-safe ServerInfo objects
+        """
+        return [
+            ServerInfo(
+                name=server.name,
+                server_type=type(server),
+                status=self._get_server_status(server),
+            )
+            for server in self.servers
+        ]
+
+    def get_server_status(self) -> dict[str, ServerStatus]:
+        """Get status of all servers.
+
+        Returns:
+            Dict mapping server names to their type-safe status
+        """
+        return {server.name: self._get_server_status(server) for server in self.servers}
+
+    def _get_server_status(self, server: BaseServer) -> ServerStatus:
+        """Get type-safe status for a specific server."""
+        if server in self._initialized_servers:
+            return "running" if server.is_running else "initialized"
+        return "not_initialized"
+
+    @property
+    def initialized_server_count(self) -> int:
+        """Number of successfully initialized servers."""
+        return len(self._initialized_servers)
+
+    @property
+    def running_server_count(self) -> int:
+        """Number of currently running servers."""
+        return sum(1 for server in self._initialized_servers if server.is_running)
+
+    def __repr__(self) -> str:
+        """String representation of aggregating server."""
+        return (
+            f"AggregatingServer(name={self.name}, "
+            f"servers={len(self.servers)}, "
+            f"initialized={len(self._initialized_servers)}, "
+            f"running={self.running_server_count})"
+        )

agentpool_server/agui_server/__init__.py

@@ -0,0 +1,11 @@
+"""AG-UI server module for agentpool.
+
+This module provides server implementation for exposing AgentPool agents
+via the AG-UI protocol with each agent on its own route.
+"""
+
+from __future__ import annotations
+
+from agentpool_server.agui_server.server import AGUIServer
+
+__all__ = ["AGUIServer"]

agentpool_server/agui_server/server.py

@@ -0,0 +1,128 @@
+"""AG-UI server implementation for agentpool pool.
+
+This module provides a server that exposes all agents in an AgentPool
+via the AG-UI 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 = 8002
+
+
+class AGUIServer(HTTPServer):
+    """AG-UI server for exposing pool agents on separate routes.
+
+    Provides a unified HTTP server that exposes all agents in the pool via
+    the AG-UI protocol. Each agent is accessible at `/{agent_name}` route.
+
+    Example:
+        ```python
+        pool = AgentPool(manifest)
+
+        server = AGUIServer(pool, host="localhost", port=8002)
+
+        async with server:
+            async with server.run_context():
+                # Agents accessible at:
+                # POST http://localhost:8002/agent1
+                # POST http://localhost:8002/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 AG-UI 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 AG-UI protocol.
+
+        Returns:
+            List of Route objects for each agent plus root listing endpoint
+        """
+        from pydantic_ai.ui.ag_ui import AGUIAdapter
+        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 AG-UI requests for a specific agent."""
+                from starlette.responses import JSONResponse
+
+                pool_agent = self.pool.agents.get(agent_name)
+                if pool_agent is None:
+                    msg = f"Agent {agent_name!r} not found"
+                    return JSONResponse({"error": msg}, status_code=404)
+                agentlet = await pool_agent.get_agentlet(None, pool_agent.model_name, str)
+                try:
+                    # Use AGUIAdapter.dispatch_request() which handles the full
+                    # AG-UI protocol: parsing request, running agent, streaming response
+                    return await AGUIAdapter.dispatch_request(request, agent=agentlet)
+                except Exception as e:
+                    self.log.exception("Error handling AG-UI request", agent=agent_name)
+                    return JSONResponse({"error": str(e)}, status_code=500)
+
+            routes.append(Route(f"/{agent_name}", agent_handler, methods=["POST"]))
+            self.log.debug("Registered AG-UI route", 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}", "model": agent.model_name}
+                for name, agent in self.pool.agents.items()
+            ]
+            return JSONResponse({"agents": agent_list, "count": len(agent_list)})
+
+        routes.append(Route("/", list_agents, methods=["GET"]))
+        self.log.info("Created AG-UI routes", agent_count=len(self.pool.agents))
+        return routes
+
+    def get_agent_url(self, agent_name: str) -> str:
+        """Get the endpoint URL for a specific agent."""
+        return f"{self.base_url}/{agent_name}"
+
+    def list_agent_routes(self) -> dict[str, str]:
+        """List all agent routes.
+
+        Returns:
+            Dictionary mapping agent names to their URLs
+        """
+        return {name: self.get_agent_url(name) for name in self.pool.agents}

agentpool_server/base.py

@@ -0,0 +1,189 @@
+"""Base server class for AgentPool servers."""
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, Self
+
+from agentpool.log import get_logger
+from agentpool.utils.tasks import TaskManager
+
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+    from types import TracebackType
+
+    from agentpool import AgentPool
+
+logger = get_logger(__name__)
+
+
+class BaseServer:
+    """Base class for all AgentPool servers.
+
+    Provides standardized interface for server lifecycle management:
+    - async def _start_async() - blocking server execution (implemented by subclasses)
+    - def start_background() - non-blocking server start via background task
+    - def stop() - stop background server task
+    - async with run_context() - automatic server start/stop management
+
+    Features:
+    - Centralized task management via TaskManager
+    - Configurable exception handling via raise_exceptions parameter
+    - Automatic pool lifecycle management
+    - Background task lifecycle management
+    """
+
+    def __init__(
+        self,
+        pool: AgentPool[Any],
+        *,
+        name: str | None = None,
+        raise_exceptions: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize base server with agent pool.
+
+        Args:
+            pool: AgentPool containing available agents
+            name: Optional Server name (auto-generated from class name if None)
+            raise_exceptions: Whether to raise exceptions during server start
+            **kwargs: Additional arguments (for subclass compatibility)
+        """
+        self.pool = pool
+        self.name = name or f"{self.__class__.__name__}-{id(self):x}"
+        self.raise_exceptions = raise_exceptions
+        self.task_manager = TaskManager()
+        self._server_task: asyncio.Task[None] | None = None
+        self._shutdown_event = asyncio.Event()
+        self.log = logger.bind(server_name=self.name)
+
+    async def __aenter__(self) -> Self:
+        """Enter async context and initialize server resources (pool, etc.)."""
+        await self.pool.__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit async context and cleanup server resources."""
+        # Cleanup any pending tasks
+        await self.task_manager.cleanup_tasks()
+        # Cleanup pool
+        await self.pool.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def _start_async(self) -> None:
+        """Start the server (blocking async - runs until stopped).
+
+        This method must be implemented by subclasses and should run
+        the server until it's stopped or encounters an error.
+
+        This is the internal implementation that subclasses override.
+        The public start() method handles exception management.
+        """
+        msg = "Subclasses must implement _start_async()"
+        raise NotImplementedError(msg)
+
+    async def start(self) -> None:
+        """Start the server (blocking async - runs until stopped).
+
+        This is the public interface that handles exception management
+        based on the raise_exceptions setting. Subclasses should implement
+        _start_async() instead of overriding this method.
+        """
+        try:
+            self.log.info("Starting server")
+            await self._start_async()
+        except Exception as e:
+            if self.raise_exceptions:
+                raise
+            self.log.exception("Server error", exc_info=e)
+        finally:
+            await self.shutdown()
+
+    async def shutdown(self) -> None:
+        """Shutdown server resources.
+
+        This method can be overridden by subclasses to add specific
+        cleanup logic. The base implementation handles task cleanup.
+        """
+        try:
+            await self.task_manager.cleanup_tasks()
+        except Exception:
+            self.log.exception("Error during server shutdown")
+        finally:
+            self._shutdown_event.set()
+            self.log.info("Server shutdown complete")
+
+    def start_background(self) -> None:
+        """Start server in background task (non-blocking).
+
+        Creates a background task that runs start() method.
+        Server will run in the background until stop() is called.
+        """
+        if self._server_task is not None and not self._server_task.done():
+            msg = "Server is already running in background"
+            raise RuntimeError(msg)
+
+        self._shutdown_event.clear()
+        self._server_task = self.task_manager.create_task(
+            self._run_with_shutdown(), name=f"{self.name}-task"
+        )
+
+    async def _run_with_shutdown(self) -> None:
+        """Internal wrapper that handles shutdown signaling."""
+        try:
+            await self.start()
+        finally:
+            self._shutdown_event.set()
+
+    def stop(self) -> None:
+        """Stop the background server task (non-blocking)."""
+        if self._server_task is not None and not self._server_task.done():
+            self._server_task.cancel()
+
+    async def wait_until_stopped(self) -> None:
+        """Wait until the server stops (either by stop() or natural completion)."""
+        if self._server_task is None:
+            return
+
+        # Wait for either task completion or shutdown event
+        await self._shutdown_event.wait()
+
+        # Ensure task is cleaned up
+        if self._server_task and not self._server_task.done():
+            self._server_task.cancel()
+
+        if self._server_task:
+            await asyncio.gather(self._server_task, return_exceptions=True)
+            self._server_task = None
+
+    @property
+    def is_running(self) -> bool:
+        """Check if server is currently running in background."""
+        return self._server_task is not None and not self._server_task.done()
+
+    @asynccontextmanager
+    async def run_context(self) -> AsyncIterator[None]:
+        """Async context manager for automatic server start/stop.
+
+        Starts the server in background when entering context,
+        automatically stops it when exiting context.
+
+        Example:
+            async with server:  # Initialize resources
+                async with server.run_context():  # Start server
+                    # Server is running in background
+                    await do_other_work()
+                # Server automatically stopped
+        """
+        self.start_background()
+        try:
+            yield
+        finally:
+            self.stop()
+            await self.wait_until_stopped()