agentpool 2.1.9__py3-none-any.whl

agentpool/hooks/prompt.py

@@ -0,0 +1,122 @@
+"""Prompt-based hook implementation."""
+
+from __future__ import annotations
+
+import json
+from string import Template
+from typing import TYPE_CHECKING, Literal
+
+from pydantic import Field
+from schemez import Schema
+
+from agentpool.hooks.base import Hook, HookResult
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from agentpool.hooks.base import HookEvent, HookInput
+
+
+logger = get_logger(__name__)
+
+# Default fast model for hook evaluation
+DEFAULT_HOOK_MODEL = "openai:gpt-4o-mini"
+
+
+class HookDecision(Schema):
+    """Structured response from LLM hook evaluation."""
+
+    decision: Literal["allow", "deny"]
+    """Whether to allow or deny the action"""
+
+    reason: str = Field(default="")
+    """Explanation for the decision"""
+
+
+class PromptHook(Hook):
+    """Hook that uses an LLM to evaluate the action.
+
+    The prompt is sent to a fast LLM which returns a structured decision.
+    Supports placeholders: $INPUT, $TOOL_NAME, $TOOL_INPUT, $AGENT_NAME, etc.
+    """
+
+    def __init__(
+        self,
+        event: HookEvent,
+        prompt: str,
+        matcher: str | None = None,
+        timeout: float = 30.0,
+        enabled: bool = True,
+        model: str | None = None,
+    ):
+        """Initialize prompt hook.
+
+        Args:
+            event: The lifecycle event this hook handles.
+            prompt: Prompt template for LLM evaluation.
+            matcher: Regex pattern for matching.
+            timeout: Maximum execution time in seconds.
+            enabled: Whether this hook is active.
+            model: Model to use for evaluation. Defaults to a fast model.
+        """
+        super().__init__(event=event, matcher=matcher, timeout=timeout, enabled=enabled)
+        self.prompt_template = prompt
+        self.model = model or DEFAULT_HOOK_MODEL
+
+    async def execute(self, input_data: HookInput) -> HookResult:
+        """Execute the LLM evaluation.
+
+        Args:
+            input_data: The hook input data.
+
+        Returns:
+            Hook result from LLM.
+        """
+        try:
+            prompt = self._build_prompt(input_data)
+            decision = await self._query_llm(prompt)
+            return HookResult(decision=decision.decision, reason=decision.reason)
+        except Exception as e:
+            logger.exception("Prompt hook failed")
+            return HookResult(decision="allow", reason=str(e))
+
+    def _build_prompt(self, input_data: HookInput) -> str:
+        """Build prompt with placeholder substitutions.
+
+        Args:
+            input_data: The hook input data.
+
+        Returns:
+            Formatted prompt string.
+        """
+        variables = {
+            "INPUT": json.dumps(dict(input_data), indent=2),
+            "TOOL_NAME": input_data.get("tool_name", ""),
+            "TOOL_INPUT": json.dumps(input_data.get("tool_input", {}), indent=2),
+            "TOOL_OUTPUT": json.dumps(input_data.get("tool_output", ""), indent=2),
+            "AGENT_NAME": input_data.get("agent_name", ""),
+            "PROMPT": input_data.get("prompt", ""),
+            "EVENT": input_data.get("event", ""),
+        }
+        template = Template(self.prompt_template)
+        return template.safe_substitute(variables)
+
+    async def _query_llm(self, prompt: str) -> HookDecision:
+        """Query the LLM for evaluation.
+
+        Args:
+            prompt: The evaluation prompt.
+
+        Returns:
+            Structured hook decision.
+        """
+        from pydantic_ai import Agent
+
+        agent: Agent[None, HookDecision] = Agent(
+            model=self.model,
+            system_prompt="You are a security/policy evaluation assistant. "
+            "Analyze the request and decide whether to allow or deny it.",
+            output_type=HookDecision,
+        )
+        result = await agent.run(prompt)
+        return result.output

agentpool/jinja_filters.py

@@ -0,0 +1,132 @@
+"""Jinja filters for agentpool documentation."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import jinja2
+
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+    import os
+
+    from mknodes.jinja.nodeenvironment import NodeEnvironment
+
+
+@jinja2.pass_environment
+def pydantic_playground_url(
+    env: NodeEnvironment,
+    files: Mapping[str, str] | Sequence[str | os.PathLike[str]],
+    active_index: int = 0,
+) -> str:
+    """Generate a Pydantic Playground URL from files.
+
+    Args:
+        env: The jinja environment (passed automatically)
+        files: Either a mapping of filenames to content, or a sequence of file paths
+        active_index: Index of the file to show as active (default: 0)
+
+    Returns:
+        URL to Pydantic Playground with files pre-loaded
+    """
+    import mknodes as mk
+
+    link = mk.MkLink.for_pydantic_playground(files, active_index=active_index)
+    return str(link.target)
+
+
+@jinja2.pass_environment
+def pydantic_playground_iframe(
+    env: NodeEnvironment,
+    files: Mapping[str, str] | Sequence[str | os.PathLike[str]],
+    width: int = 1200,
+    height: int = 900,
+    active_index: int = 0,
+) -> str:
+    """Generate an MkIFrame for Pydantic Playground.
+
+    Args:
+        env: The jinja environment (passed automatically)
+        files: Either a mapping of filenames to content, or a sequence of file paths
+        width: Width of the iframe
+        height: Height of the iframe
+        active_index: Index of the file to show as active
+
+    Returns:
+        MkIFrame node
+    """
+    import mknodes as mk
+
+    link = mk.MkLink.for_pydantic_playground(files, active_index=active_index)
+    return str(mk.MkIFrame(str(link.target), width=width, height=height, parent=env.node))
+
+
+@jinja2.pass_environment
+def pydantic_playground_link(
+    env: NodeEnvironment,
+    files: Mapping[str, str] | Sequence[str | os.PathLike[str]],
+    title: str = "Open in Pydantic Playground",
+    active_index: int = 0,
+    as_button: bool = True,
+) -> str:
+    """Generate an MkLink to Pydantic Playground.
+
+    Args:
+        env: The jinja environment (passed automatically)
+        files: Either a mapping of filenames to content, or a sequence of file paths
+        title: Link text
+        active_index: Index of the file to show as active
+        as_button: Whether to style as a button
+
+    Returns:
+        MkLink node
+    """
+    import mknodes as mk
+
+    link = mk.MkLink.for_pydantic_playground(
+        files, title=title, active_index=active_index, parent=env.node
+    )
+    if as_button:
+        link.as_button = True
+    return str(link)
+
+
+@jinja2.pass_environment
+def pydantic_playground(
+    env: NodeEnvironment,
+    files: Mapping[str, str] | Sequence[str | os.PathLike[str]],
+    width: int = 1200,
+    height: int = 900,
+    active_index: int = 0,
+    show_link: bool = True,
+    link_title: str = "Open in Pydantic Playground",
+) -> str:
+    """Generate both iframe and link for Pydantic Playground.
+
+    Args:
+        env: The jinja environment (passed automatically)
+        files: Either a mapping of filenames to content, or a sequence of file paths
+        width: Width of the iframe
+        height: Height of the iframe
+        active_index: Index of the file to show as active
+        show_link: Whether to show a link below the iframe
+        link_title: Text for the link
+
+    Returns:
+        MkContainer with iframe and optional link
+    """
+    import mknodes as mk
+
+    link = mk.MkLink.for_pydantic_playground(files, active_index=active_index)
+    iframe = mk.MkIFrame(str(link.target), width=width, height=height)
+
+    if show_link:
+        button_link = mk.MkLink.for_pydantic_playground(
+            files, title=link_title, active_index=active_index
+        )
+        button_link.as_button = True
+        container = mk.MkContainer([iframe, button_link], parent=env.node)
+        return str(container)
+    container = mk.MkContainer([iframe], parent=env.node)
+    return str(container)

agentpool/log.py

@@ -0,0 +1,224 @@
+"""Logging configuration for agentpool with structlog support."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+import logging
+from logging.handlers import RotatingFileHandler
+import sys
+from typing import TYPE_CHECKING, Any
+
+import logfire
+import structlog
+
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator, Sequence
+
+    from slashed import OutputWriter
+
+
+LogLevel = int | str
+
+
+_LOGGING_CONFIGURED = False
+
+
+def configure_logging(
+    level: LogLevel = "INFO",
+    *,
+    use_colors: bool | None = None,
+    json_logs: bool = False,
+    force: bool = False,
+    log_file: str | None = None,
+) -> None:
+    """Configure structlog and standard logging.
+
+    Args:
+        level: Logging level
+        use_colors: Whether to use colored output (auto-detected if None)
+        json_logs: Force JSON output regardless of TTY detection
+        force: Force reconfiguration even if already configured
+        log_file: Optional file path to write logs to instead of stderr
+    """
+    global _LOGGING_CONFIGURED  # noqa: PLW0603
+
+    if _LOGGING_CONFIGURED and not force:
+        return
+
+    if isinstance(level, str):
+        level = getattr(logging, level.upper())
+
+    if log_file:
+        _configure_file_logging(level, log_file)
+    else:
+        _configure_console_logging(level, use_colors=use_colors, json_logs=json_logs)
+
+    _LOGGING_CONFIGURED = True
+
+
+def _configure_file_logging(level: LogLevel, log_file: str, max_lines: int = 5000) -> None:
+    """Configure logging to write to a file with human-readable format and line-based rotation.
+
+    Args:
+        level: Logging level
+        log_file: Path to log file
+        max_lines: Maximum number of lines before rotation (default: 5000)
+    """
+    if isinstance(level, str):
+        level = getattr(logging, level.upper())
+
+    # Estimate bytes per line (assuming ~200 chars average)
+    max_bytes = max_lines * 200
+
+    # Set up rotating file handler with 3 backup files
+    handler = RotatingFileHandler(log_file, maxBytes=max_bytes, backupCount=3, encoding="utf-8")
+    handler.setFormatter(logging.Formatter("%(message)s"))
+    logging.basicConfig(level=level, handlers=[handler], force=True)
+
+    # Configure structlog processors for file output
+    processors: list[Any] = [
+        structlog.stdlib.filter_by_level,
+        structlog.stdlib.add_log_level,
+        structlog.stdlib.add_logger_name,
+        structlog.stdlib.PositionalArgumentsFormatter(),
+        structlog.processors.TimeStamper(fmt="iso"),
+        structlog.processors.StackInfoRenderer(),
+        structlog.processors.format_exc_info,
+        structlog.processors.UnicodeDecoder(),
+        logfire.StructlogProcessor(),  # Capture logs for observability
+        structlog.dev.ConsoleRenderer(
+            colors=False, exception_formatter=structlog.dev.plain_traceback
+        ),
+    ]
+
+    structlog.configure(
+        processors=processors,
+        context_class=dict,
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        wrapper_class=structlog.stdlib.BoundLogger,
+        cache_logger_on_first_use=True,
+    )
+
+
+def _configure_console_logging(
+    level: LogLevel, *, use_colors: bool | None = None, json_logs: bool = False
+) -> None:
+    """Configure logging to stderr with console or JSON format."""
+    if isinstance(level, str):
+        level = getattr(logging, level.upper())
+
+    # Determine output format
+    colors = sys.stderr.isatty() and not json_logs if use_colors is None else use_colors or False
+    use_console_renderer = not json_logs and (colors or sys.stderr.isatty())
+    # Configure standard logging as backend
+    handler = logging.StreamHandler(sys.stderr)
+    if use_console_renderer:
+        # For console output, don't show level in stdlib logging (structlog handles it)
+        handler.setFormatter(logging.Formatter("%(message)s"))
+        logging.basicConfig(level=level, handlers=[handler], force=True)
+    else:  # For structured output, use minimal formatting
+        logging.basicConfig(level=level, handlers=[handler], force=True, format="%(message)s")
+    processors: list[Any] = [
+        structlog.stdlib.filter_by_level,
+        structlog.stdlib.add_log_level,
+        structlog.stdlib.PositionalArgumentsFormatter(),
+        structlog.processors.TimeStamper(fmt="iso"),
+        structlog.processors.StackInfoRenderer(),
+        structlog.processors.UnicodeDecoder(),
+    ]
+
+    # Add logger name only for non-console renderers (avoid duplication with stdlib)
+    if not use_console_renderer:
+        processors.insert(1, structlog.stdlib.add_logger_name)
+        processors.append(structlog.processors.format_exc_info)
+
+    # Add logfire processor before final renderer to capture logs for observability
+    processors.append(logfire.StructlogProcessor())
+
+    # Add final renderer
+    if use_console_renderer:
+        processors.append(structlog.dev.ConsoleRenderer(colors=colors))
+    else:
+        processors.append(structlog.processors.JSONRenderer())
+
+    structlog.configure(
+        processors=processors,
+        context_class=dict,
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        wrapper_class=structlog.stdlib.BoundLogger,
+        cache_logger_on_first_use=True,
+    )
+
+
+def get_logger(name: str, log_level: LogLevel | None = None) -> structlog.stdlib.BoundLogger:
+    """Get a structlog logger for the given name.
+
+    Args:
+        name: The name of the logger, will be prefixed with 'agentpool.'
+        log_level: The logging level to set for the logger
+
+    Returns:
+        A structlog BoundLogger instance
+    """
+    # Ensure basic structlog configuration exists for tests
+    if not _LOGGING_CONFIGURED and not structlog.is_configured():
+        # Minimal configuration that doesn't interfere with stdio
+        structlog.configure(
+            processors=[
+                structlog.stdlib.filter_by_level,
+                structlog.stdlib.add_log_level,
+                structlog.processors.StackInfoRenderer(),
+                logfire.StructlogProcessor(),
+                structlog.dev.ConsoleRenderer(colors=False),
+            ],
+            wrapper_class=structlog.stdlib.BoundLogger,
+            logger_factory=structlog.stdlib.LoggerFactory(),
+            cache_logger_on_first_use=True,
+        )
+
+    logger = structlog.get_logger(f"agentpool.{name}")
+    if log_level is not None:
+        if isinstance(log_level, str):
+            log_level = getattr(logging, log_level.upper())
+            assert log_level
+        # Set level on underlying stdlib logger
+        stdlib_logger = logging.getLogger(f"agentpool.{name}")
+        stdlib_logger.setLevel(log_level)
+    return logger  # type: ignore[no-any-return]
+
+
+@contextmanager
+def set_handler_level(
+    level: int,
+    logger_names: Sequence[str],
+    *,
+    session_handler: OutputWriter | None = None,
+) -> Iterator[None]:
+    """Temporarily set logging level and optionally add session handler.
+
+    Args:
+        level: Logging level to set
+        logger_names: Names of loggers to configure
+        session_handler: Optional output writer for session logging
+    """
+    loggers = [logging.getLogger(name) for name in logger_names]
+    old_levels = [logger.level for logger in loggers]
+
+    handler = None
+    if session_handler:
+        from slashed.log import SessionLogHandler
+
+        handler = SessionLogHandler(session_handler)
+        for logger in loggers:
+            logger.addHandler(handler)
+
+    try:
+        for logger in loggers:
+            logger.setLevel(level)
+        yield
+    finally:
+        for logger, old_level in zip(loggers, old_levels, strict=True):
+            logger.setLevel(old_level)
+            if handler:
+                logger.removeHandler(handler)

agentpool/mcp_server/__init__.py

@@ -0,0 +1,17 @@
+"""MCP server integration for AgentPool."""
+
+from agentpool.mcp_server.client import MCPClient
+from agentpool.mcp_server.tool_bridge import (
+    BridgeConfig,
+    ToolBridgeRegistry,
+    ToolManagerBridge,
+    create_tool_bridge,
+)
+
+__all__ = [
+    "BridgeConfig",
+    "MCPClient",
+    "ToolBridgeRegistry",
+    "ToolManagerBridge",
+    "create_tool_bridge",
+]