agentpool 2.1.9__py3-none-any.whl → 2.2.3__py3-none-any.whl

agentpool_config/compaction.py

@@ -0,0 +1,258 @@
+"""Compaction configuration for message history management."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Annotated, Any, Literal
+
+from pydantic import Field
+from schemez import Schema
+
+
+if TYPE_CHECKING:
+    from agentpool.messaging.compaction import CompactionPipeline, CompactionStep
+
+
+class FilterThinkingConfig(Schema):
+    """Configuration for FilterThinking step."""
+
+    type: Literal["filter_thinking"] = "filter_thinking"
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import FilterThinking
+
+        return FilterThinking()
+
+
+class FilterRetryPromptsConfig(Schema):
+    """Configuration for FilterRetryPrompts step."""
+
+    type: Literal["filter_retry_prompts"] = "filter_retry_prompts"
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import FilterRetryPrompts
+
+        return FilterRetryPrompts()
+
+
+class FilterBinaryContentConfig(Schema):
+    """Configuration for FilterBinaryContent step."""
+
+    type: Literal["filter_binary"] = "filter_binary"
+    keep_references: bool = False
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import FilterBinaryContent
+
+        return FilterBinaryContent(keep_references=self.keep_references)
+
+
+class FilterToolCallsConfig(Schema):
+    """Configuration for FilterToolCalls step."""
+
+    type: Literal["filter_tools"] = "filter_tools"
+    exclude_tools: list[str] = Field(default_factory=list)
+    include_only: list[str] | None = None
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import FilterToolCalls
+
+        return FilterToolCalls(exclude_tools=self.exclude_tools, include_only=self.include_only)
+
+
+class FilterEmptyMessagesConfig(Schema):
+    """Configuration for FilterEmptyMessages step."""
+
+    type: Literal["filter_empty"] = "filter_empty"
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import FilterEmptyMessages
+
+        return FilterEmptyMessages()
+
+
+class TruncateToolOutputsConfig(Schema):
+    """Configuration for TruncateToolOutputs step."""
+
+    type: Literal["truncate_tool_outputs"] = "truncate_tool_outputs"
+    max_length: int = 2000
+    suffix: str = "\n... [truncated]"
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import TruncateToolOutputs
+
+        return TruncateToolOutputs(max_length=self.max_length, suffix=self.suffix)
+
+
+class TruncateTextPartsConfig(Schema):
+    """Configuration for TruncateTextParts step."""
+
+    type: Literal["truncate_text"] = "truncate_text"
+    max_length: int = 5000
+    suffix: str = "\n... [truncated]"
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import TruncateTextParts
+
+        return TruncateTextParts(max_length=self.max_length, suffix=self.suffix)
+
+
+class KeepLastMessagesConfig(Schema):
+    """Configuration for KeepLastMessages step."""
+
+    type: Literal["keep_last"] = "keep_last"
+    count: int = 10
+    count_pairs: bool = True
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import KeepLastMessages
+
+        return KeepLastMessages(count=self.count, count_pairs=self.count_pairs)
+
+
+class KeepFirstMessagesConfig(Schema):
+    """Configuration for KeepFirstMessages step."""
+
+    type: Literal["keep_first"] = "keep_first"
+    count: int = 2
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import KeepFirstMessages
+
+        return KeepFirstMessages(count=self.count)
+
+
+class KeepFirstAndLastConfig(Schema):
+    """Configuration for KeepFirstAndLast step."""
+
+    type: Literal["keep_first_last"] = "keep_first_last"
+    first_count: int = 2
+    last_count: int = 5
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import KeepFirstAndLast
+
+        return KeepFirstAndLast(first_count=self.first_count, last_count=self.last_count)
+
+
+class TokenBudgetConfig(Schema):
+    """Configuration for TokenBudget step."""
+
+    type: Literal["token_budget"] = "token_budget"
+    max_tokens: int = 4000
+    model: str = "gpt-4o"
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import TokenBudget
+
+        return TokenBudget(max_tokens=self.max_tokens, model=self.model)
+
+
+class SummarizeConfig(Schema):
+    """Configuration for Summarize step."""
+
+    type: Literal["summarize"] = "summarize"
+    model: str = "openai:gpt-4o-mini"
+    threshold: int = 15
+    keep_recent: int = 5
+    summary_prompt: str | None = None
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import Summarize
+
+        kwargs: dict[str, Any] = {
+            "model": self.model,
+            "threshold": self.threshold,
+            "keep_recent": self.keep_recent,
+        }
+        if self.summary_prompt:
+            kwargs["summary_prompt"] = self.summary_prompt
+        return Summarize(**kwargs)
+
+
+class WhenMessageCountExceedsConfig(Schema):
+    """Configuration for WhenMessageCountExceeds wrapper."""
+
+    type: Literal["when_count_exceeds"] = "when_count_exceeds"
+    threshold: int = 20
+    step: "CompactionStepConfig"
+
+    def build(self) -> CompactionStep:
+        from agentpool.messaging.compaction import WhenMessageCountExceeds
+
+        return WhenMessageCountExceeds(step=self.step.build(), threshold=self.threshold)
+
+
+# Union of all config types with discriminator
+CompactionStepConfig = Annotated[
+    FilterThinkingConfig
+    | FilterRetryPromptsConfig
+    | FilterBinaryContentConfig
+    | FilterToolCallsConfig
+    | FilterEmptyMessagesConfig
+    | TruncateToolOutputsConfig
+    | TruncateTextPartsConfig
+    | KeepLastMessagesConfig
+    | KeepFirstMessagesConfig
+    | KeepFirstAndLastConfig
+    | TokenBudgetConfig
+    | SummarizeConfig
+    | WhenMessageCountExceedsConfig,
+    Field(discriminator="type"),
+]
+
+# Update forward reference
+WhenMessageCountExceedsConfig.model_rebuild()
+
+
+class CompactionConfig(Schema):
+    """Configuration for message compaction/summarization.
+
+    Example YAML:
+        ```yaml
+        compaction:
+          steps:
+            - type: filter_thinking
+            - type: truncate_tool_outputs
+              max_length: 1000
+            - type: keep_last
+              count: 10
+        ```
+
+    Or use a preset:
+        ```yaml
+        compaction:
+          preset: balanced
+        ```
+    """
+
+    preset: Literal["minimal", "balanced", "summarizing"] | None = None
+    """Use a predefined compaction pipeline.
+
+    - minimal: Aggressively minimize context (filter thinking, truncate, keep last 10)
+    - balanced: Balance context preservation with size (filter thinking, truncate,
+                keep first 2 + last 8)
+    - summarizing: Summarize old messages (filter thinking, summarize when > 20 messages)
+    """
+
+    steps: list[CompactionStepConfig] = Field(default_factory=list)
+    """Ordered list of compaction steps to apply (ignored if preset is set)."""
+
+    def build(self) -> CompactionPipeline:
+        """Build a CompactionPipeline from this configuration."""
+        from agentpool.messaging.compaction import (
+            CompactionPipeline,
+            balanced_context,
+            minimal_context,
+            summarizing_context,
+        )
+
+        if self.preset:
+            match self.preset:
+                case "minimal":
+                    return minimal_context()
+                case "balanced":
+                    return balanced_context()
+                case "summarizing":
+                    return summarizing_context()
+
+        return CompactionPipeline(steps=[step.build() for step in self.steps])

agentpool_config/mcp_server.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import os
 from typing import TYPE_CHECKING, Annotated, Literal, Self
 
-from pydantic import ConfigDict, Field, HttpUrl
+from pydantic import ConfigDict, Field, HttpUrl, model_validator
 from schemez import Schema
 
 
@@ -73,6 +73,62 @@ class BaseMCPServerConfig(Schema):
     )
     """Timeout for the server process in seconds."""
 
+    enabled_tools: list[str] | None = Field(
+        default=None,
+        examples=[["read_file", "list_directory"], ["search", "fetch"]],
+        title="Enabled tools",
+    )
+    """If set, only these tools will be available (whitelist).
+    Mutually exclusive with disabled_tools."""
+
+    disabled_tools: list[str] | None = Field(
+        default=None,
+        examples=[["delete_file", "write_file"], ["dangerous_tool"]],
+        title="Disabled tools",
+    )
+    """Tools to exclude from this server (blacklist). Mutually exclusive with enabled_tools."""
+
+    @model_validator(mode="after")
+    def _validate_tool_filters(self) -> Self:
+        """Validate that enabled_tools and disabled_tools are mutually exclusive."""
+        if self.enabled_tools is not None and self.disabled_tools is not None:
+            msg = "Cannot specify both 'enabled_tools' and 'disabled_tools'"
+            raise ValueError(msg)
+        return self
+
+    def is_tool_allowed(self, tool_name: str) -> bool:
+        """Check if a tool is allowed based on enabled/disabled lists.
+
+        Args:
+            tool_name: Name of the tool to check
+
+        Returns:
+            True if the tool is allowed, False otherwise
+        """
+        if self.enabled_tools is not None:
+            return tool_name in self.enabled_tools
+        if self.disabled_tools is not None:
+            return tool_name not in self.disabled_tools
+        return True
+
+    def needs_tool_filtering(self) -> bool:
+        """Check if this config has tool filtering configured."""
+        return self.enabled_tools is not None or self.disabled_tools is not None
+
+    def wrap_with_mcp_filter(self) -> StdioMCPServerConfig:
+        """Wrap this MCP server with mcp-filter for tool filtering.
+
+        Creates a new StdioMCPServerConfig that runs mcp-filter as a proxy,
+        applying the configured enabled_tools/disabled_tools filtering.
+
+        Returns:
+            A new StdioMCPServerConfig that wraps the original server with mcp-filter
+
+        Raises:
+            NotImplementedError: Subclasses must implement this method
+        """
+        raise NotImplementedError
+
     def get_env_vars(self) -> dict[str, str]:
         """Get environment variables for the server process."""
         env = os.environ.copy()
@@ -142,6 +198,34 @@ class StdioMCPServerConfig(BaseMCPServerConfig):
         """Generate a unique client ID for this stdio server configuration."""
         return f"{self.command}_{' '.join(self.args)}"
 
+    def wrap_with_mcp_filter(self) -> StdioMCPServerConfig:
+        """Wrap this stdio MCP server with mcp-filter for tool filtering.
+
+        Returns:
+            A new StdioMCPServerConfig that wraps this server with mcp-filter
+        """
+        filter_args = ["mcp-filter", "run", "-t", "stdio", "--stdio-command", self.command]
+
+        # Add original args as a single --stdio-arg
+        if self.args:
+            filter_args.extend(["--stdio-arg", " ".join(self.args)])
+
+        # Add allowlist (exact tool names)
+        if self.enabled_tools:
+            filter_args.extend(["-a", ",".join(self.enabled_tools)])
+
+        # Add denylist (regex patterns)
+        if self.disabled_tools:
+            filter_args.extend(["-d", ",".join(self.disabled_tools)])
+
+        return StdioMCPServerConfig(
+            name=self.name,
+            command="uvx",
+            args=filter_args,
+            env=self.env,
+            timeout=self.timeout,
+        )
+
     def to_pydantic_ai(self) -> MCPServerStdio:
         """Convert to pydantic-ai MCPServerStdio instance."""
         from pydantic_ai.mcp import MCPServerStdio
@@ -186,6 +270,29 @@ class SSEMCPServerConfig(BaseMCPServerConfig):
         """Generate a unique client ID for this SSE server configuration."""
         return f"sse_{self.url}"
 
+    def wrap_with_mcp_filter(self) -> StdioMCPServerConfig:
+        """Wrap this SSE MCP server with mcp-filter for tool filtering.
+
+        Returns:
+            A new StdioMCPServerConfig that wraps this server with mcp-filter
+        """
+        filter_args = ["mcp-filter", "run", "-t", "http", "--http-url", str(self.url)]
+
+        # Add allowlist (exact tool names)
+        if self.enabled_tools:
+            filter_args.extend(["-a", ",".join(self.enabled_tools)])
+
+        # Add denylist (regex patterns)
+        if self.disabled_tools:
+            filter_args.extend(["-d", ",".join(self.disabled_tools)])
+
+        return StdioMCPServerConfig(
+            name=self.name,
+            command="uvx",
+            args=filter_args,
+            timeout=self.timeout,
+        )
+
     def to_pydantic_ai(self) -> MCPServerSSE:
         """Convert to pydantic-ai MCPServerSSE instance."""
         from pydantic_ai.mcp import MCPServerSSE
@@ -225,6 +332,29 @@ class StreamableHTTPMCPServerConfig(BaseMCPServerConfig):
         """Generate a unique client ID for this streamable HTTP server configuration."""
         return f"streamable_http_{self.url}"
 
+    def wrap_with_mcp_filter(self) -> StdioMCPServerConfig:
+        """Wrap this HTTP MCP server with mcp-filter for tool filtering.
+
+        Returns:
+            A new StdioMCPServerConfig that wraps this server with mcp-filter
+        """
+        filter_args = ["mcp-filter", "run", "-t", "http", "--http-url", str(self.url)]
+
+        # Add allowlist (exact tool names)
+        if self.enabled_tools:
+            filter_args.extend(["-a", ",".join(self.enabled_tools)])
+
+        # Add denylist (regex patterns)
+        if self.disabled_tools:
+            filter_args.extend(["-d", ",".join(self.disabled_tools)])
+
+        return StdioMCPServerConfig(
+            name=self.name,
+            command="uvx",
+            args=filter_args,
+            timeout=self.timeout,
+        )
+
     def to_pydantic_ai(self) -> MCPServerStreamableHTTP:
         """Convert to pydantic-ai MCPServerStreamableHTTP instance."""
         from pydantic_ai.mcp import MCPServerStreamableHTTP

agentpool_config/storage.py

@@ -170,8 +170,53 @@ class MemoryStorageConfig(BaseStorageProviderConfig):
     """In-memory storage configuration for testing."""
 
 
+class ClaudeStorageConfig(BaseStorageProviderConfig):
+    """Claude Code native storage format configuration.
+
+    Reads/writes to Claude Code's native JSONL format in ~/.claude/projects/.
+    Useful for sharing conversation history between agentpool and Claude Code CLI.
+    """
+
+    model_config = ConfigDict(json_schema_extra={"x-doc-title": "Claude Storage"})
+
+    type: Literal["claude"] = Field("claude", init=False)
+    """Claude Code native storage configuration."""
+
+    path: str = Field(
+        default="~/.claude",
+        examples=["~/.claude", "/home/user/.claude"],
+        title="Claude data directory",
+    )
+    """Path to Claude data directory (default: ~/.claude)"""
+
+
+class OpenCodeStorageConfig(BaseStorageProviderConfig):
+    """OpenCode native storage format configuration.
+
+    Reads from OpenCode's native JSON format in ~/.local/share/opencode/storage/.
+    Useful for importing conversation history from OpenCode.
+    """
+
+    model_config = ConfigDict(json_schema_extra={"x-doc-title": "OpenCode Storage"})
+
+    type: Literal["opencode"] = Field("opencode", init=False)
+    """OpenCode native storage configuration."""
+
+    path: str = Field(
+        default="~/.local/share/opencode/storage",
+        examples=["~/.local/share/opencode/storage"],
+        title="OpenCode storage directory",
+    )
+    """Path to OpenCode storage directory."""
+
+
 StorageProviderConfig = Annotated[
-    SQLStorageConfig | FileStorageConfig | TextLogConfig | MemoryStorageConfig,
+    SQLStorageConfig
+    | FileStorageConfig
+    | TextLogConfig
+    | MemoryStorageConfig
+    | ClaudeStorageConfig
+    | OpenCodeStorageConfig,
     Field(discriminator="type"),
 ]
 

agentpool_config/tools.py

@@ -109,4 +109,10 @@ class ImportToolConfig(BaseToolConfig):
         )
 
 
-ToolConfig = Annotated[ImportToolConfig, Field(discriminator="type")]
+from agentpool_config.builtin_tools import BuiltinToolConfig  # noqa: E402
+
+
+ToolConfig = Annotated[
+    ImportToolConfig | BuiltinToolConfig,
+    Field(discriminator="type"),
+]