agentpool 2.1.9__py3-none-any.whl

agentpool_server/http_server.py

@@ -0,0 +1,164 @@
+"""Base class for HTTP-based servers.
+
+This module provides a base class for servers that expose HTTP routes via Starlette.
+"""
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Self
+
+from agentpool.log import get_logger
+from agentpool_server.base import BaseServer
+
+
+if TYPE_CHECKING:
+    from starlette.applications import Starlette
+    from starlette.routing import Route
+    from upathtools import JoinablePathLike
+
+    from agentpool import AgentPool
+
+
+logger = get_logger(__name__)
+
+DEFAULT_HTTP_PORT = 8000
+
+
+class HTTPServer(BaseServer):
+    """Base class for HTTP-based servers.
+
+    Provides common infrastructure for servers that expose HTTP routes via Starlette.
+
+    Subclasses must implement:
+    - `get_routes()`: Return list of Starlette Route objects
+
+    Example:
+        ```python
+        class MyHTTPServer(HTTPServer):
+            async def get_routes(self) -> list[Route]:
+                return [
+                    Route("/hello", self.hello_handler, methods=["GET"]),
+                    Route("/world", self.world_handler, methods=["POST"]),
+                ]
+        ```
+    """
+
+    def __init__(
+        self,
+        pool: AgentPool,
+        *,
+        name: str | None = None,
+        host: str = "localhost",
+        port: int = DEFAULT_HTTP_PORT,
+        raise_exceptions: bool = False,
+    ) -> None:
+        """Initialize HTTP 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, raise_exceptions=raise_exceptions)
+        self.host = host
+        self.port = port
+
+    @abstractmethod
+    async def get_routes(self) -> list[Route]:
+        """Get Starlette routes for this server.
+
+        Subclasses must implement this method to provide their HTTP routes.
+
+        Returns:
+            List of Starlette Route objects
+        """
+        ...
+
+    async def create_app(self) -> Starlette:
+        """Create Starlette application with this server's routes.
+
+        Returns:
+            Starlette application instance
+        """
+        from starlette.applications import Starlette
+
+        routes = await self.get_routes()
+        app = Starlette(debug=False, routes=routes)
+        self.log.info(
+            "Created HTTP app",
+            route_count=len(routes),
+            host=self.host,
+            port=self.port,
+        )
+        return app
+
+    async def _start_async(self) -> None:
+        """Start the HTTP server (blocking async - runs until stopped)."""
+        try:
+            import uvicorn
+        except ImportError as e:
+            msg = "Please install 'uvicorn' to use HTTPServer. You can use: uv add uvicorn"
+            raise ImportError(msg) from e
+
+        self.log.info(
+            "Starting HTTP server",
+            host=self.host,
+            port=self.port,
+        )
+
+        app = await self.create_app()
+
+        config = uvicorn.Config(
+            app=app,
+            host=self.host,
+            port=self.port,
+            log_level="info",
+        )
+
+        server = uvicorn.Server(config)
+
+        try:
+            await server.serve()
+        except Exception:
+            self.log.exception("HTTP server error")
+            raise
+        finally:
+            self.log.info("HTTP server stopped")
+
+    @classmethod
+    def from_config(
+        cls,
+        config_path: JoinablePathLike,
+        *,
+        host: str = "localhost",
+        port: int = DEFAULT_HTTP_PORT,
+        raise_exceptions: bool = False,
+    ) -> Self:
+        """Create HTTP server from configuration file.
+
+        Args:
+            config_path: Path to agentpool YAML config file
+            host: Host to bind server to
+            port: Port to bind server to
+            raise_exceptions: Whether to raise exceptions during server start
+
+        Returns:
+            Configured server instance with agent pool from config
+        """
+        from agentpool import AgentPool
+        from agentpool.models.manifest import AgentsManifest
+
+        manifest = AgentsManifest.from_file(config_path)
+        pool = AgentPool(manifest=manifest)
+        server = cls(pool, host=host, port=port, raise_exceptions=raise_exceptions)
+        agent_names = list(server.pool.agents.keys())
+        server.log.info("Created HTTP server from config", agent_names=agent_names)
+        return server
+
+    @property
+    def base_url(self) -> str:
+        """Get the base URL for the server."""
+        return f"http://{self.host}:{self.port}"

agentpool_server/mcp_server/__init__.py

@@ -0,0 +1,6 @@
+"""MCP protocol server implementation for AgentPool."""
+
+from agentpool_server.mcp_server.server import MCPServer
+
+
+__all__ = ["MCPServer"]

agentpool_server/mcp_server/server.py

@@ -0,0 +1,314 @@
+"""MCP protocol server implementation using FastMCP decorators."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any
+
+from fastmcp import FastMCP
+from fastmcp.server.middleware.caching import ResponseCachingMiddleware
+from key_value.aio.stores.disk import DiskStore
+import mcp
+from mcp import types
+import platformdirs
+
+import agentpool
+from agentpool.log import get_logger
+from agentpool.mcp_server import constants
+from agentpool_server import BaseServer
+
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+    from contextlib import AbstractAsyncContextManager
+
+    from mcp.server.lowlevel.server import LifespanResultT
+    from pydantic import AnyUrl
+
+    from agentpool import AgentPool
+    from agentpool.prompts.prompts import BasePrompt
+    from agentpool.tools import Tool
+    from agentpool_config.pool_server import MCPPoolServerConfig
+
+    LifespanHandler = Callable[
+        [FastMCP[LifespanResultT]],
+        AbstractAsyncContextManager[LifespanResultT],
+    ]
+
+logger = get_logger(__name__)
+
+
+agentpool_dir = platformdirs.user_config_dir("agentpool")
+
+store = DiskStore(directory=agentpool_dir)
+middleware = ResponseCachingMiddleware(cache_storage=store)
+
+
+class MCPServer(BaseServer):
+    """MCP protocol server implementation using FastMCP decorators.
+
+    This server uses FastMCP's native @tool and @prompt decorators for
+    cleaner, more maintainable code. Tools and prompts from the pool are
+    registered dynamically using the same decorator pattern.
+    """
+
+    def __init__(
+        self,
+        pool: AgentPool[Any],
+        config: MCPPoolServerConfig,
+        lifespan: LifespanHandler | None = None,
+        instructions: str | None = None,
+        name: str = "agentpool-server",
+        raise_exceptions: bool = False,
+    ) -> None:
+        """Initialize server with agent pool.
+
+        Args:
+            pool: AgentPool to expose through MCP
+            config: Server configuration
+            name: Server name for MCP protocol
+            lifespan: Lifespan context manager
+            instructions: Instructions for server usage
+            raise_exceptions: Whether to raise exceptions during server start
+        """
+        from agentpool.resource_providers.pool import PoolResourceProvider
+
+        super().__init__(pool, name=name, raise_exceptions=raise_exceptions)
+        self.provider = PoolResourceProvider(pool, zed_mode=config.zed_mode)
+        self.config = config
+
+        self._subscriptions: defaultdict[str, set[mcp.ServerSession]] = defaultdict(set)
+        self._tools_registered = False
+        self._prompts_registered = False
+
+        # Create FastMCP instance
+        self.fastmcp = FastMCP(
+            instructions=instructions,
+            lifespan=lifespan,
+            version=agentpool.__version__,
+            middleware=[middleware],
+        )
+        self.server = self.fastmcp._mcp_server
+
+        # Register static handlers
+        self._register_logging_handlers()
+        self._register_resource_handlers()
+
+    async def _register_pool_tools(self) -> None:
+        """Register all pool tools using FastMCP's @tool decorator."""
+        if self._tools_registered:
+            return
+
+        tools = await self.provider.get_tools()
+
+        for tool in tools:
+            # Create handler with closure to capture tool instance
+            def make_handler(tool_instance: Tool) -> Callable[..., Awaitable[str]]:
+                async def handler(**kwargs: Any) -> str:
+                    """Dynamically generated tool handler."""
+                    # Filter out _meta arguments
+                    args = {k: v for k, v in kwargs.items() if not k.startswith("_")}
+                    try:
+                        result = await tool_instance.execute(**args)
+                        return str(result)
+                    except Exception as exc:
+                        logger.exception("Tool execution failed", name=tool_instance.name)
+                        return f"Tool execution failed: {exc}"
+
+                # Set proper function metadata for FastMCP
+                handler.__name__ = tool_instance.name
+                handler.__doc__ = tool_instance.description or f"Execute {tool_instance.name}"
+                return handler
+
+            # Register using FastMCP's tool decorator
+            tool_handler = make_handler(tool)
+            tool_annotations = types.ToolAnnotations(
+                title=tool.name,
+                readOnlyHint=tool.hints.read_only,
+                destructiveHint=tool.hints.destructive,
+                idempotentHint=tool.hints.idempotent,
+                openWorldHint=tool.hints.open_world,
+            )
+            self.fastmcp.tool(annotations=tool_annotations, task=True)(tool_handler)
+
+        self._tools_registered = True
+        logger.info("Registered MCP tools", count=len(tools))
+
+    async def _register_pool_prompts(self) -> None:
+        """Register all pool prompts using FastMCP's @prompt decorator."""
+        if self._prompts_registered:
+            return
+
+        prompts = await self.provider.get_prompts()
+
+        for prompt in prompts:
+            # Create handler with closure to capture prompt instance
+            def make_handler(prompt_instance: BasePrompt) -> Callable[..., Awaitable[str]]:
+                async def handler(**kwargs: Any) -> str:
+                    """Dynamically generated prompt handler."""
+                    from agentpool.mcp_server import conversions
+
+                    try:
+                        assert prompt_instance.name
+                        parts = await self.provider.get_request_parts(
+                            prompt_instance.name, kwargs or {}
+                        )
+                        messages = [msg for p in parts for msg in conversions.to_mcp_messages(p)]
+
+                        # Format messages as text
+                        result_lines = []
+                        for msg in messages:
+                            if isinstance(msg.content, types.TextContent):
+                                result_lines.append(f"{msg.role}: {msg.content.text}")
+                            else:
+                                result_lines.append(f"{msg.role}: {msg.content}")
+
+                        return "\n".join(result_lines)
+                    except Exception as exc:
+                        logger.exception("Prompt execution failed", name=prompt_instance.name)
+                        return f"Prompt execution failed: {exc}"
+
+                # Set proper function metadata for FastMCP
+                assert prompt_instance.name
+                handler.__name__ = prompt_instance.name
+                handler.__doc__ = (
+                    prompt_instance.description or f"Get prompt {prompt_instance.name}"
+                )
+                return handler
+
+            # Register using FastMCP's prompt decorator
+            prompt_handler = make_handler(prompt)
+            self.fastmcp.prompt()(prompt_handler)
+
+        self._prompts_registered = True
+        logger.info("Registered MCP prompts", count=len(prompts))
+
+    def _register_resource_handlers(self) -> None:
+        """Register resource subscription handlers."""
+
+        @self.server.subscribe_resource()  # type: ignore[no-untyped-call, untyped-decorator]
+        async def handle_subscribe(uri: AnyUrl) -> None:
+            """Subscribe to resource updates."""
+            uri_str = str(uri)
+            self._subscriptions[uri_str].add(self.current_session)
+            logger.debug("Added subscription", uri=uri)
+
+        @self.server.unsubscribe_resource()  # type: ignore[no-untyped-call, untyped-decorator]
+        async def handle_unsubscribe(uri: AnyUrl) -> None:
+            """Unsubscribe from resource updates."""
+            if (uri_str := str(uri)) in self._subscriptions:
+                self._subscriptions[uri_str].discard(self.current_session)
+                if not self._subscriptions[uri_str]:
+                    del self._subscriptions[uri_str]
+                logger.debug("Removed subscription", uri=uri)
+
+    def _register_logging_handlers(self) -> None:
+        """Register logging-related handlers."""
+
+        @self.server.set_logging_level()  # type: ignore[no-untyped-call, untyped-decorator]
+        async def handle_set_level(level: mcp.LoggingLevel) -> None:
+            """Handle logging level changes."""
+            try:
+                python_level = constants.MCP_TO_LOGGING[level]
+                logger.setLevel(python_level)
+                data = f"Log level set to {level}"
+                await self.current_session.send_log_message(level, data, logger=self.name)
+            except Exception as exc:
+                error_data = mcp.ErrorData(
+                    message="Error setting log level",
+                    code=types.INTERNAL_ERROR,
+                    data=str(exc),
+                )
+                raise mcp.McpError(error_data) from exc
+
+        @self.server.progress_notification()  # type: ignore[no-untyped-call, untyped-decorator]
+        async def handle_progress(
+            token: str | int,
+            progress: float,
+            total: float | None,
+            message: str | None = None,
+        ) -> None:
+            """Handle progress notifications from client."""
+            logger.debug(
+                "Progress notification",
+                token=token,
+                progress=progress,
+                total=total,
+                message=message,
+            )
+
+    async def _start_async(self) -> None:
+        """Start the server (blocking async - runs until stopped)."""
+        # Register pool tools and prompts before starting
+        await self._register_pool_tools()
+        await self._register_pool_prompts()
+
+        # Start FastMCP server
+        if self.config.transport == "stdio":
+            await self.fastmcp.run_async(transport=self.config.transport)
+        else:
+            await self.fastmcp.run_async(
+                transport=self.config.transport,
+                host=self.config.host,
+                port=self.config.port,
+            )
+
+    @property
+    def current_session(self) -> mcp.ServerSession:
+        """Get current session from request context."""
+        try:
+            return self.server.request_context.session
+        except LookupError as exc:
+            msg = "No active request context"
+            raise RuntimeError(msg) from exc
+
+    async def report_progress(
+        self,
+        progress: float,
+        total: float | None = None,
+        message: str | None = None,
+        related_request_id: str | None = None,
+    ) -> None:
+        """Report progress for the current operation."""
+        progress_token = (
+            self.server.request_context.meta.progressToken
+            if self.server.request_context.meta
+            else None
+        )
+
+        if progress_token is None:
+            return
+
+        await self.server.request_context.session.send_progress_notification(
+            progress_token=progress_token,
+            progress=progress,
+            total=total,
+            message=message,
+            related_request_id=related_request_id,
+        )
+
+    @property
+    def client_info(self) -> mcp.Implementation | None:
+        """Get client info from current session."""
+        session = self.current_session
+        if not session.client_params:
+            return None
+        return session.client_params.clientInfo
+
+    async def notify_tool_list_changed(self) -> None:
+        """Notify clients about tool list changes."""
+        try:
+            self.task_manager.create_task(self.current_session.send_tool_list_changed())
+        except RuntimeError:
+            self.log.debug("No active session for notification")
+        except Exception:
+            self.log.exception("Failed to send tool list change notification")
+
+    async def notify_prompt_list_changed(self) -> None:
+        """Notify clients about prompt list changes."""
+        try:
+            self.task_manager.create_task(self.current_session.send_prompt_list_changed())
+        except RuntimeError:
+            self.log.debug("No active session for notification")
+        except Exception:
+            self.log.exception("Failed to send prompt list change notification")

agentpool_server/mcp_server/zed_wrapper.py

@@ -0,0 +1,110 @@
+"""Function parameter encoding utils for Zed compatibility."""
+
+from __future__ import annotations
+
+from functools import wraps
+from typing import TYPE_CHECKING, Any
+
+from agentpool.log import get_logger
+from agentpool.prompts.prompts import DynamicPrompt
+from agentpool.utils.importing import import_callable
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from agentpool.prompts.prompts import BasePrompt
+
+
+logger = get_logger(__name__)
+
+
+def decode_zed_args(input_str: str) -> dict[str, Any]:
+    """Decode Zed-style input string into kwargs dict."""
+    if " :: " not in input_str:
+        # Single argument case - use as first argument
+        return {"main_arg": input_str}
+
+    # Multiple arguments case
+    parts = input_str.split(" :: ", 1)
+    result: dict[str, Any] = {"main_arg": parts[0]}
+
+    if len(parts) > 1:
+        for pair in parts[1].split(" | "):
+            if not pair:
+                continue
+            key, value = pair.split("=", 1)
+            # Convert value types
+            match value.lower():
+                case "true":
+                    result[key] = True
+                case "false":
+                    result[key] = False
+                case "null":
+                    result[key] = None
+                case _:
+                    try:
+                        if "." in value:
+                            result[key] = float(value)
+                        else:
+                            result[key] = int(value)
+                    except ValueError:
+                        result[key] = value
+
+    return result
+
+
+def create_zed_wrapper(func: Any) -> Any:
+    """Create a wrapper that accepts Zed-style input."""
+
+    @wraps(func)
+    def wrapper(input: str, **_kwargs: Any) -> Any:  # noqa: A002
+        kwargs = decode_zed_args(input)
+        return func(**kwargs)
+
+    # Copy over metadata
+    wrapper.__module__ = func.__module__
+    wrapper.__qualname__ = f"zed_wrapped_{func.__qualname__}"
+    return wrapper
+
+
+def prepare_prompts_for_zed(prompts: Sequence[BasePrompt]) -> list[BasePrompt]:
+    """Prepare prompts for Zed compatibility."""
+    result: list[BasePrompt] = []
+
+    for prompt in prompts:
+        # Only process DynamicPrompts
+        if not isinstance(prompt, DynamicPrompt):
+            result.append(prompt)
+            continue
+
+        # Skip if only one or zero parameters
+        # if len(prompt.arguments) <= 1:
+        #     result.append(prompt)
+        #     continue
+
+        try:
+            # Get and wrap the original function
+            original_func = (
+                import_callable(prompt.import_path)
+                if isinstance(prompt.import_path, str)
+                else prompt.import_path
+            )
+            wrapped_func = create_zed_wrapper(original_func)
+
+            # Create new prompt with single input parameter
+            new_prompt = DynamicPrompt(
+                name=prompt.name,
+                description=prompt.description,
+                import_path=f"{wrapped_func.__module__}.{wrapped_func.__qualname__}",
+                template=prompt.template,
+                completions=prompt.completions,
+            )
+            result.append(new_prompt)
+            logger.debug("Wrapped prompt for Zed mode", name=prompt.name)
+
+        except Exception:
+            logger.exception("Failed to wrap function for prompt", name=prompt.name)
+            result.append(prompt)
+
+    return result

agentpool_server/openai_api_server/__init__.py

@@ -0,0 +1,5 @@
+"""OpenAI-compatible API server for AgentPool."""
+
+from .server import OpenAIAPIServer
+
+__all__ = ["OpenAIAPIServer"]

agentpool_server/openai_api_server/completions/__init__.py

@@ -0,0 +1 @@
+"""Completions API."""

agentpool_server/openai_api_server/completions/helpers.py

@@ -0,0 +1,81 @@
+"""Helpers for OpenAI-compatible API server."""
+
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, Any
+
+import anyenv
+from pydantic_ai import PartDeltaEvent, TextPartDelta
+
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
+    from agentpool.common_types import SupportsRunStream
+    from agentpool_server.openai_api_server.completions.models import ChatCompletionRequest
+
+logger = get_logger(__name__)
+
+
+async def stream_response(
+    agent: SupportsRunStream[Any],
+    content: str,
+    request: ChatCompletionRequest,
+) -> AsyncGenerator[str]:
+    """Generate streaming response chunks."""
+    response_id = f"chatcmpl-{int(time.time() * 1000)}"
+    created = int(time.time())
+
+    try:
+        # First chunk with role
+        choice = {"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}
+        first_chunk = {
+            "id": response_id,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": request.model,
+            "choices": [choice],
+        }
+        yield f"data: {anyenv.dump_json(first_chunk)}\n\n"
+        async for event in agent.run_stream(content):
+            match event:
+                case PartDeltaEvent(delta=TextPartDelta(content_delta=chunk)):
+                    # Skip empty chunks
+                    if not chunk:
+                        continue
+                    delta = {"content": chunk}
+                    choice = {"index": 0, "delta": delta, "finish_reason": None}
+                    chunk_data = {
+                        "id": response_id,
+                        "object": "chat.completion.chunk",
+                        "created": created,
+                        "model": request.model,
+                        "choices": [choice],
+                    }
+                    yield f"data: {anyenv.dump_json(chunk_data)}\n\n"
+        final_chunk = {
+            "id": response_id,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": request.model,
+            "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+        }
+        yield f"data: {anyenv.dump_json(final_chunk)}\n\n"
+        yield "data: [DONE]\n\n"
+
+    except Exception as e:
+        logger.exception("Error during streaming response")
+        delta = {"content": f"Error: {e!s}"}
+        choice = {"index": 0, "delta": delta, "finish_reason": "error"}
+        error_chunk = {
+            "id": response_id,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": request.model,
+            "choices": [choice],
+        }
+        yield f"data: {anyenv.dump_json(error_chunk)}\n\n"
+        yield "data: [DONE]\n\n"