agentpool 2.1.9__py3-none-any.whl

agentpool_cli/serve_vercel.py

@@ -0,0 +1,233 @@
+"""Command for serving agents via Vercel AI protocol."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import TYPE_CHECKING, Annotated, Any
+
+import typer as t
+
+from agentpool_cli import log, resolve_agent_config
+
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from agentpool import ChatMessage
+
+
+logger = log.get_logger(__name__)
+
+
+def vercel_command(  # noqa: PLR0915
+    ctx: t.Context,
+    config: Annotated[str | None, t.Argument(help="Path to agent configuration")] = None,
+    agent_name: Annotated[
+        str | None, t.Option("--agent", "-a", help="Specific agent to serve")
+    ] = None,
+    host: Annotated[str, t.Option(help="Host to bind server to")] = "localhost",
+    port: Annotated[int, t.Option(help="Port to listen on")] = 8000,
+    cors: Annotated[bool, t.Option(help="Enable CORS")] = True,
+    show_messages: Annotated[
+        bool, t.Option("--show-messages", help="Show message activity")
+    ] = False,
+) -> None:
+    """Serve agents via Vercel AI Data Stream Protocol.
+
+    This creates a server compatible with Vercel AI SDK frontends,
+    allowing you to use your agents with Vercel AI UI components.
+
+    The server exposes a POST /chat endpoint that accepts Vercel AI
+    protocol requests and streams responses back.
+
+    If --agent is specified, only that agent is served. Otherwise,
+    the endpoint accepts an 'agent' field in the request to select
+    which agent to use.
+    """
+    from fastapi import FastAPI
+    from fastapi.middleware.cors import CORSMiddleware
+    from pydantic_ai import PartDeltaEvent, PartStartEvent, TextPart, TextPartDelta
+    from starlette.requests import Request
+    from starlette.responses import JSONResponse, Response, StreamingResponse
+    import uvicorn
+
+    from agentpool import AgentPool, AgentsManifest
+    from agentpool.agents.events import StreamCompleteEvent
+
+    logger.info("Server PID", pid=os.getpid())
+
+    def on_message(message: ChatMessage[Any]) -> None:
+        print(message.format(style="simple"))
+
+    try:
+        config_path = resolve_agent_config(config)
+    except ValueError as e:
+        msg = str(e)
+        raise t.BadParameter(msg) from e
+
+    manifest = AgentsManifest.from_file(config_path)
+    pool = AgentPool(manifest)
+
+    if show_messages:
+        for agent in pool.agents.values():
+            agent.message_sent.connect(on_message)
+
+    # Create FastAPI app
+    app = FastAPI(
+        title="AgentPool - Vercel AI Server",
+        description="Vercel AI Data Stream Protocol server for AgentPool",
+    )
+
+    if cors:
+        app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )
+
+    @app.on_event("startup")
+    async def startup() -> None:
+        await pool.__aenter__()
+        logger.info("Agent pool initialized")
+
+    @app.on_event("shutdown")
+    async def shutdown() -> None:
+        await pool.__aexit__(None, None, None)
+        logger.info("Agent pool shut down")
+
+    @app.post("/chat")
+    async def chat(request: Request) -> Response:
+        """Handle Vercel AI protocol chat requests.
+
+        Implements the Vercel AI Data Stream Protocol:
+        https://sdk.vercel.ai/docs/ai-sdk-ui/stream-protocol#data-stream-protocol
+        """
+        body = await request.body()
+        try:
+            data = json.loads(body)
+        except json.JSONDecodeError as e:
+            return JSONResponse({"error": f"Invalid JSON: {e}"}, status_code=400)
+
+        # Extract messages from the request
+        messages = data.get("messages", [])
+        if not messages:
+            return JSONResponse({"error": "No messages provided"}, status_code=400)
+
+        # Get the last user message
+        last_message = messages[-1]
+        user_text = ""
+        if last_message.get("role") == "user":
+            parts = last_message.get("parts", [])
+            for part in parts:
+                if part.get("type") == "text":
+                    user_text = part.get("text", "")
+                    break
+
+        if not user_text:
+            return JSONResponse({"error": "No user text found"}, status_code=400)
+
+        # Determine which agent to use
+        if agent_name:
+            selected_agent = pool.get_agent(agent_name)
+        else:
+            if not pool.agents:
+                return JSONResponse({"error": "No agents available"}, status_code=500)
+            first_name = next(iter(pool.agents.keys()))
+            selected_agent = pool.get_agent(first_name)
+
+        async def generate_stream() -> AsyncIterator[str]:
+            """Generate Vercel AI Data Stream Protocol events.
+
+            Protocol format:
+            - Text: 0:"text content"
+            - Finish: e:{"finishReason":"stop",...}
+            - Done: d:{"finishReason":"stop"}
+            """
+            try:
+                async for event in selected_agent.run_stream(user_text):
+                    # Handle pydantic-ai streaming events
+                    if isinstance(event, PartStartEvent):
+                        # New part started - if it's text, emit it
+                        if isinstance(event.part, TextPart) and event.part.content:
+                            text = event.part.content
+                            escaped = json.dumps(text)
+                            yield f"0:{escaped}\n"
+
+                    elif isinstance(event, PartDeltaEvent):
+                        # Delta update - emit text deltas
+                        if isinstance(event.delta, TextPartDelta):
+                            text = event.delta.content_delta
+                            if text:
+                                escaped = json.dumps(text)
+                                yield f"0:{escaped}\n"
+
+                    elif isinstance(event, StreamCompleteEvent):
+                        # Stream complete - we've received the final message
+                        # The content has already been streamed via deltas
+                        pass
+
+                # Send finish event
+                finish_data = {
+                    "finishReason": "stop",
+                    "usage": {"promptTokens": 0, "completionTokens": 0},
+                }
+                yield f"e:{json.dumps(finish_data)}\n"
+
+                # Send done marker
+                done_data = {"finishReason": "stop"}
+                yield f"d:{json.dumps(done_data)}\n"
+
+            except Exception as e:
+                logger.exception("Error during streaming")
+                # Send error as text
+                error_msg = f"Error: {e}"
+                escaped = json.dumps(error_msg)
+                yield f"0:{escaped}\n"
+                # Still send finish
+                finish_data = {"finishReason": "error"}
+                yield f"e:{json.dumps(finish_data)}\n"
+                done_data = {"finishReason": "error"}
+                yield f"d:{json.dumps(done_data)}\n"
+
+        return StreamingResponse(
+            generate_stream(),
+            media_type="text/plain; charset=utf-8",
+            headers={
+                "X-Vercel-AI-Data-Stream": "v1",
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+            },
+        )
+
+    @app.get("/agents")
+    async def list_agents() -> dict[str, Any]:
+        """List available agents."""
+        return {
+            "agents": [
+                {"name": name, "description": agent.description}
+                for name, agent in pool.agents.items()
+            ]
+        }
+
+    @app.get("/health")
+    async def health() -> dict[str, str]:
+        """Health check endpoint."""
+        return {"status": "ok"}
+
+    # Get log level from the global context
+    log_level = ctx.obj.get("log_level", "info") if ctx.obj else "info"
+
+    print(f"Starting Vercel AI server on http://{host}:{port}")
+    print(f"Chat endpoint: POST http://{host}:{port}/chat")
+    print(f"Available agents: {list(pool.agents.keys())}")
+
+    uvicorn.run(app, host=host, port=port, log_level=log_level.lower())
+
+
+if __name__ == "__main__":
+    import typer
+
+    typer.run(vercel_command)

agentpool_cli/store.py

@@ -0,0 +1,171 @@
+"""Manages stored config mappings."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypedDict
+
+from agentpool_cli.log import get_logger
+
+
+if TYPE_CHECKING:
+    from upath import UPath
+    from upathtools import JoinablePathLike
+
+
+logger = get_logger(__name__)
+
+
+@dataclass(frozen=True)
+class ConfigFile:
+    """Represents an active configuration."""
+
+    name: str
+    path: str
+
+    def __str__(self) -> str:
+        """Format for display."""
+        return f"{self.name} ({self.path})"
+
+
+class ConfigMapping(TypedDict):
+    """Type for config storage format."""
+
+    configs: dict[str, str]  # name -> uri mapping
+    active: str | None
+
+
+class ConfigStore:
+    """Manages stored config mappings."""
+
+    def __init__(self, filename: str | None = None) -> None:
+        """Initialize store with default paths."""
+        self._filename = filename or "configs.json"
+
+    @property
+    def config_dir(self) -> UPath:
+        """Get config directory."""
+        import platformdirs
+        from upathtools import to_upath
+
+        return to_upath(platformdirs.user_config_dir("agentpool"))
+
+    @property
+    def config_file(self) -> UPath:
+        """Get config file path."""
+        return self.config_dir / self._filename
+
+    def _ensure_config_dir(self) -> None:
+        """Create config directory if needed."""
+        self.config_dir.mkdir(parents=True, exist_ok=True)
+        if not self.config_file.exists():
+            mapping = ConfigMapping(configs={}, active=None)
+            self.save_mapping(mapping)
+
+    def load_mapping(self) -> ConfigMapping:
+        """Load config mapping from disk."""
+        import anyenv
+
+        if not self.config_file.exists():
+            return ConfigMapping(configs={}, active=None)
+        try:
+            text = self.config_file.read_text("utf-8")
+            data = anyenv.load_json(text, return_type=dict)
+            active = data.get("active")
+            configs = data.get("configs", {})
+            return ConfigMapping(configs=configs, active=active)
+        except Exception:
+            logger.exception("Failed to load config mapping")
+            return ConfigMapping(configs={}, active=None)
+
+    def save_mapping(self, mapping: ConfigMapping) -> None:
+        """Save config mapping to disk."""
+        import anyenv
+
+        try:
+            self.config_file.write_text(anyenv.dump_json(mapping, indent=True))
+        except Exception:
+            logger.exception("Failed to save config mapping")
+
+    def add_config(self, name: str, path: JoinablePathLike) -> None:
+        """Add a new named config.
+
+        Args:
+            name: Name to register the config under
+            path: Path to the config file
+
+        Raises:
+            ValueError: If name is invalid
+            FileNotFoundError: If config file doesn't exist
+            PermissionError: If config file can't be read
+            IsADirectoryError: If path points to a directory
+        """
+        # Basic validation
+        from upathtools import to_upath
+
+        if not name.isidentifier():
+            msg = f"Invalid config name: {name} (must be a valid Python identifier)"
+            raise ValueError(msg)
+
+        path_obj = to_upath(path).resolve()
+        if not path_obj.exists():
+            msg = f"Config file not found: {path}"
+            raise FileNotFoundError(msg)
+
+        if not path_obj.is_file():
+            msg = f"Path is not a file: {path}"
+            raise IsADirectoryError(msg)
+
+        try:
+            # Try to read file to verify access
+            path_obj.read_bytes()
+        except PermissionError as exc:
+            msg = f"Cannot read config file: {path}"
+            raise PermissionError(msg) from exc
+        # All good, save the config
+        mapping = self.load_mapping()
+        mapping["configs"][name] = str(path_obj)
+        logger.debug("Adding config %r -> %s", name, path_obj)
+        self.save_mapping(mapping)
+
+    def remove_config(self, name: str) -> None:
+        """Remove a named config."""
+        mapping = self.load_mapping()
+        if name == mapping["active"]:
+            mapping["active"] = None
+        mapping["configs"].pop(name, None)
+        self.save_mapping(mapping)
+
+    def set_active(self, name: str) -> None:
+        """Set the active config."""
+        mapping = self.load_mapping()
+        if name not in mapping["configs"]:
+            msg = f"Config {name} not found"
+            raise KeyError(msg)
+        mapping["active"] = name
+        self.save_mapping(mapping)
+
+    def get_active(self) -> ConfigFile | None:
+        """Get active config if one is set.
+
+        Returns:
+            ConfigFile if an active config is set, None otherwise
+        """
+        mapping = self.load_mapping()
+        if not mapping["active"]:
+            return None
+        name = mapping["active"]
+        return ConfigFile(name=name, path=mapping["configs"][name])
+
+    def list_configs(self) -> list[tuple[str, str]]:
+        """List all configs with their paths."""
+        mapping = self.load_mapping()
+        return list(mapping["configs"].items())
+
+    def get_config(self, name: str) -> str:
+        """Get path for a named config."""
+        mapping = self.load_mapping()
+        if name not in mapping["configs"]:
+            msg = f"Config {name} not found"
+            raise KeyError(msg)
+        return mapping["configs"][name]

agentpool_cli/task.py

@@ -0,0 +1,84 @@
+"""Task execution command."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Annotated
+
+import typer as t
+
+from agentpool_cli import log, resolve_agent_config
+
+
+if TYPE_CHECKING:
+    from agentpool import AgentsManifest
+
+
+TASK_HELP = """
+Execute a defined task with the specified agent.
+
+Example:
+    agentpool task docs write_api_docs --prompt "Include code examples"
+    agentpool task docs write_api_docs --log-level DEBUG
+"""
+
+
+async def execute_job(
+    agent_name: str,
+    task_name: str,
+    config: AgentsManifest,
+    *,
+    prompt: str | None = None,
+) -> str:
+    """Execute task with agent."""
+    from agentpool import AgentPool
+
+    async with AgentPool(config) as pool:
+        # Get both agent and task
+        agent = pool.get_agent(agent_name)
+        task = pool.get_job(task_name)
+
+        # Create final prompt from task and additional input
+        task_prompt = task.prompt
+        if prompt:
+            task_prompt = f"{task_prompt}\n\nAdditional instructions:\n{prompt}"
+
+        result = await agent.run(task_prompt)
+        return result.data
+
+
+def task_command(
+    agent_name: Annotated[str, t.Argument(help="Name of agent to run task with")],
+    task_name: Annotated[str, t.Argument(help="Name of task to execute")],
+    config: Annotated[
+        str | None, t.Option("--config", "-c", help="Agent configuration file")
+    ] = None,
+    prompt: Annotated[str | None, t.Option("--prompt", "-p", help="Additional prompt")] = None,
+) -> None:
+    """Execute a task with the specified agent."""
+    logger = log.get_logger(__name__)
+    try:
+        logger.debug("Starting task execution", name=task_name)
+
+        # Resolve configuration path
+        try:
+            config_path = resolve_agent_config(config)
+            logger.debug("Using config", source=config_path)
+        except ValueError as e:
+            msg = str(e)
+            raise t.BadParameter(msg) from e
+
+        from agentpool import AgentsManifest
+
+        manifest = AgentsManifest.from_file(config_path)
+        result = asyncio.run(execute_job(agent_name, task_name, manifest, prompt=prompt))
+        print(result)
+
+    except Exception as e:
+        t.echo(f"Error: {e}", err=True)
+        logger.debug("Exception details", exc_info=True)
+        raise t.Exit(1) from e
+
+
+if __name__ == "__main__":
+    t.run(task_command)

agentpool_cli/utils.py

@@ -0,0 +1,104 @@
+"""Utilities for creating Typer commands from Pydantic models."""
+
+from __future__ import annotations
+
+import inspect
+import types as pytypes
+from typing import TYPE_CHECKING, Annotated, Any, Union, get_args, get_origin
+
+import typer
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from inspect import Parameter
+
+    from pydantic import BaseModel
+
+
+def resolve_type(field_type: Any) -> Any:
+    """Resolve actual type from Union/Optional types."""
+    origin = get_origin(field_type)
+
+    # Handle Any
+    if field_type is Any:
+        return str
+
+    # Handle both Union and | syntax
+    if origin is Union or origin is pytypes.UnionType:
+        types_ = [t for t in get_args(field_type) if t is not type(None)]
+        return resolve_type(types_[0]) if len(types_) == 1 else str
+
+    # Handle lists/sets
+    if origin in {list, set}:
+        item_type = get_args(field_type)[0]
+        return list[resolve_type(item_type)]  # type: ignore[misc]
+
+    return field_type
+
+
+def create_typer_command[T: BaseModel](
+    config_cls: type[T],
+    callback: Callable[[T], Any] | None = None,
+    *,
+    name: str | None = None,
+) -> typer.models.CommandInfo:
+    """Create a Typer CommandInfo from a Pydantic model."""
+
+    def make_command(**kwargs: Any) -> Any:
+        """Command created from model."""
+        config = config_cls(**kwargs)
+        if callback:
+            return callback(config)
+        return config
+
+    # Create parameters from model fields
+    params: dict[str, Parameter] = {}
+    for field_name, field in config_cls.model_fields.items():
+        if field_name == "type":  # Skip discriminator
+            continue
+
+        field_type = resolve_type(field.annotation)
+        default = field.default if field.default != ... else None
+
+        # First field becomes argument
+        if not params:
+            # For arguments, the param name should be the field name
+            params[field_name] = inspect.Parameter(
+                field_name,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                annotation=Annotated[field_type, typer.Argument(help=field.description)],
+                default=default,
+            )
+            continue
+
+        # For options, we need proper param_decls
+        option_name = field_name.replace("_", "-")
+        params[field_name] = inspect.Parameter(
+            field_name,
+            inspect.Parameter.KEYWORD_ONLY,
+            annotation=Annotated[
+                field_type,
+                typer.Option(
+                    "--" + option_name,  # Explicit param_decls
+                    help=field.description,
+                ),
+            ],
+            default=default,
+        )
+
+    # Create function with signature
+    make_command.__signature__ = inspect.Signature(  # type: ignore
+        list(params.values()), return_annotation=config_cls
+    )
+    make_command.__annotations__ = {k: v.annotation for k, v in params.items()}
+
+    # Create command
+    if name is None:
+        name = config_cls.__name__.lower().replace("command", "")
+
+    return typer.models.CommandInfo(
+        name=name,
+        callback=make_command,
+        help=config_cls.__doc__,
+    )

agentpool_cli/watch.py

@@ -0,0 +1,54 @@
+"""Command for watching agents and displaying messages."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Annotated, Any
+
+import typer as t
+
+from agentpool_cli import log
+from agentpool_cli.cli_types import DetailLevel  # noqa: TC001
+
+
+if TYPE_CHECKING:
+    from agentpool import ChatMessage
+
+
+logger = log.get_logger(__name__)
+
+
+def watch_command(
+    config: Annotated[str, t.Argument(help="Path to agent configuration")],
+    show_messages: Annotated[
+        bool, t.Option("--show-messages", help="Show all messages (not just final responses)")
+    ] = True,
+    detail_level: Annotated[
+        DetailLevel, t.Option("-d", "--detail", help="Output detail level")
+    ] = "simple",
+    show_metadata: Annotated[bool, t.Option("--metadata", help="Show message metadata")] = False,
+    show_costs: Annotated[bool, t.Option("--costs", help="Show token usage and costs")] = False,
+) -> None:
+    """Run agents in event-watching mode."""
+
+    def on_message(chat_message: ChatMessage[Any]) -> None:
+        text = chat_message.format(
+            style=detail_level,
+            show_metadata=show_metadata,
+            show_costs=show_costs,
+        )
+        print(text)
+
+    async def run_watch() -> None:
+        from agentpool import AgentPool, AgentsManifest
+
+        manifest = AgentsManifest.from_file(config)
+        async with AgentPool(manifest) as pool:
+            # Connect message handlers if showing all messages
+            if show_messages:
+                for agent in pool.agents.values():
+                    agent.message_sent.connect(on_message)
+
+            await pool.run_event_loop()
+
+    asyncio.run(run_watch())