headroom-ai 0.2.13__py3-none-any.whl

headroom/integrations/mcp/server.py

@@ -0,0 +1,533 @@
+"""Headroom MCP Integration: Compress tool outputs from MCP servers.
+
+This module provides multiple ways to integrate Headroom with MCP:
+
+1. HeadroomMCPProxy - A proxy server that wraps upstream MCP servers
+2. compress_tool_result() - Standalone function for host applications
+3. HeadroomMCPMiddleware - Transport-level middleware
+
+The key insight: MCP tool outputs are the PERFECT use case for Headroom.
+They're often large (100s-1000s of items), structured (JSON), and contain
+mostly low-relevance data with a few critical items (errors, matches).
+
+Example - Proxy Server:
+    ```python
+    # Configure proxy to wrap your MCP servers
+    proxy = HeadroomMCPProxy(
+        upstream_servers=["slack", "database", "github"],
+        config=HeadroomConfig(),
+    )
+
+    # Run as MCP server - clients connect to this instead
+    proxy.run()
+    ```
+
+Example - Standalone Function:
+    ```python
+    # In your MCP host application
+    result = await mcp_client.call_tool("search_logs", {"service": "api"})
+
+    # Compress before adding to context
+    compressed = compress_tool_result(
+        content=result,
+        tool_name="search_logs",
+        tool_args={"service": "api"},
+        user_query="find errors in api service",
+    )
+    ```
+
+Example - Middleware (for MCP client libraries):
+    ```python
+    # Wrap your MCP client's transport
+    middleware = HeadroomMCPMiddleware(config)
+    client = MCPClient(transport=middleware.wrap(base_transport))
+    ```
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from headroom.config import HeadroomConfig, SmartCrusherConfig
+from headroom.providers import OpenAIProvider
+from headroom.transforms import SmartCrusher
+
+
+@dataclass
+class MCPCompressionResult:
+    """Result of compressing an MCP tool output."""
+
+    original_content: str
+    compressed_content: str
+    original_tokens: int
+    compressed_tokens: int
+    tokens_saved: int
+    compression_ratio: float
+    items_before: int | None
+    items_after: int | None
+    errors_preserved: int
+    was_compressed: bool
+    tool_name: str
+    context_used: str
+
+
+@dataclass
+class MCPToolProfile:
+    """Configuration profile for a specific MCP tool.
+
+    Different tools may need different compression strategies:
+    - Slack search: High error preservation, relevance to query
+    - Database query: Schema detection, anomaly preservation
+    - File listing: Minimal compression (paths are important)
+    """
+
+    tool_name_pattern: str  # Regex pattern to match tool names
+    enabled: bool = True
+    max_items: int = 20
+    min_tokens_to_compress: int = 500
+    preserve_error_keywords: set[str] = field(
+        default_factory=lambda: {"error", "failed", "exception", "critical", "fatal"}
+    )
+    always_keep_fields: set[str] = field(default_factory=set)  # Fields to never drop
+
+
+# Default profiles for common MCP servers
+DEFAULT_MCP_PROFILES: list[MCPToolProfile] = [
+    # Slack - preserve errors and messages matching query
+    MCPToolProfile(
+        tool_name_pattern=r".*slack.*",
+        max_items=25,
+        preserve_error_keywords={"error", "failed", "exception", "bug", "issue", "broken"},
+    ),
+    # Database - preserve errors and anomalies
+    MCPToolProfile(
+        tool_name_pattern=r".*database.*|.*sql.*|.*query.*",
+        max_items=30,
+        preserve_error_keywords={"error", "null", "failed", "exception", "violation"},
+    ),
+    # GitHub - preserve errors and high-priority issues
+    MCPToolProfile(
+        tool_name_pattern=r".*github.*|.*git.*",
+        max_items=20,
+        preserve_error_keywords={"error", "bug", "critical", "urgent", "blocker"},
+    ),
+    # Logs - preserve ALL errors
+    MCPToolProfile(
+        tool_name_pattern=r".*log.*|.*trace.*",
+        max_items=40,  # Keep more for logs
+        preserve_error_keywords={"error", "fatal", "critical", "exception", "failed", "panic"},
+    ),
+    # File system - minimal compression (paths matter)
+    MCPToolProfile(
+        tool_name_pattern=r".*file.*|.*fs.*|.*directory.*",
+        max_items=50,
+        min_tokens_to_compress=1000,  # Higher threshold
+    ),
+    # Generic fallback
+    MCPToolProfile(
+        tool_name_pattern=r".*",
+        max_items=20,
+    ),
+]
+
+
+class HeadroomMCPCompressor:
+    """Core compression logic for MCP tool outputs.
+
+    This class handles the actual compression of MCP tool results.
+    It's used by both the proxy server and standalone functions.
+    """
+
+    def __init__(
+        self,
+        config: HeadroomConfig | None = None,
+        profiles: list[MCPToolProfile] | None = None,
+        token_counter: Callable[[str], int] | None = None,
+    ):
+        """Initialize MCP compressor.
+
+        Args:
+            config: Headroom configuration.
+            profiles: Tool-specific compression profiles.
+            token_counter: Function to count tokens. Uses tiktoken if None.
+        """
+        self.config = config or HeadroomConfig()
+        self.profiles = profiles or DEFAULT_MCP_PROFILES
+
+        # Initialize token counter
+        if token_counter:
+            self._count_tokens = token_counter
+        else:
+            provider = OpenAIProvider()
+            counter = provider.get_token_counter("gpt-4o")
+            self._count_tokens = counter.count_text
+
+    def get_profile(self, tool_name: str) -> MCPToolProfile:
+        """Get the compression profile for a tool."""
+        for profile in self.profiles:
+            if re.match(profile.tool_name_pattern, tool_name, re.IGNORECASE):
+                return profile
+        # Return last profile (generic fallback)
+        return self.profiles[-1]
+
+    def compress(
+        self,
+        content: str,
+        tool_name: str,
+        tool_args: dict[str, Any] | None = None,
+        user_query: str = "",
+    ) -> MCPCompressionResult:
+        """Compress MCP tool output.
+
+        Args:
+            content: Raw tool output (usually JSON string).
+            tool_name: Name of the MCP tool (e.g., "mcp__slack__search").
+            tool_args: Arguments passed to the tool (used for context).
+            user_query: User's original query (for relevance scoring).
+
+        Returns:
+            MCPCompressionResult with compressed content and metrics.
+        """
+        profile = self.get_profile(tool_name)
+        original_tokens = self._count_tokens(content)
+
+        # Build context for relevance scoring
+        context_parts = []
+        if user_query:
+            context_parts.append(user_query)
+        if tool_args:
+            context_parts.append(json.dumps(tool_args))
+        context = " ".join(context_parts)
+
+        # Check if compression is needed
+        if not profile.enabled or original_tokens < profile.min_tokens_to_compress:
+            return MCPCompressionResult(
+                original_content=content,
+                compressed_content=content,
+                original_tokens=original_tokens,
+                compressed_tokens=original_tokens,
+                tokens_saved=0,
+                compression_ratio=0.0,
+                items_before=None,
+                items_after=None,
+                errors_preserved=0,
+                was_compressed=False,
+                tool_name=tool_name,
+                context_used=context,
+            )
+
+        # Try to parse as JSON
+        try:
+            json.loads(content)
+        except json.JSONDecodeError:
+            # Not JSON, return as-is
+            return MCPCompressionResult(
+                original_content=content,
+                compressed_content=content,
+                original_tokens=original_tokens,
+                compressed_tokens=original_tokens,
+                tokens_saved=0,
+                compression_ratio=0.0,
+                items_before=None,
+                items_after=None,
+                errors_preserved=0,
+                was_compressed=False,
+                tool_name=tool_name,
+                context_used=context,
+            )
+
+        # Find arrays to compress
+        items_before = 0
+        items_after = 0
+        errors_preserved = 0
+
+        # Create SmartCrusher with profile settings
+        smart_config = SmartCrusherConfig(
+            enabled=True,
+            min_tokens_to_crush=profile.min_tokens_to_compress,
+            max_items_after_crush=profile.max_items,
+        )
+        crusher = SmartCrusher(config=smart_config)  # type: ignore[arg-type]
+
+        # Build messages for SmartCrusher (it expects conversation format)
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": context or f"Process {tool_name} results"},
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [
+                    {
+                        "id": "call_1",
+                        "function": {"name": tool_name, "arguments": json.dumps(tool_args or {})},
+                    }
+                ],
+            },
+            {"role": "tool", "content": content, "tool_call_id": "call_1"},
+        ]
+
+        # Create tokenizer wrapper
+        class TokenizerWrapper:
+            def __init__(self, count_fn: Any) -> None:
+                self._count = count_fn
+
+            def count_text(self, text: str) -> int:
+                result = self._count(text)
+                return int(result) if result is not None else 0
+
+            def count_messages(self, messages: list[dict[str, Any]]) -> int:
+                total = 0
+                for msg in messages:
+                    if msg.get("content"):
+                        total += self._count(str(msg["content"]))
+                return total
+
+        tokenizer = TokenizerWrapper(self._count_tokens)
+
+        # Apply SmartCrusher
+        result = crusher.apply(messages, tokenizer=tokenizer)  # type: ignore[arg-type]
+        compressed_content = result.messages[-1]["content"]
+
+        # Remove any Headroom markers for clean output
+        compressed_content = re.sub(r"\n<headroom:[^>]+>", "", compressed_content)
+
+        # Count items and errors
+        try:
+            original_data = json.loads(content)
+            compressed_data = json.loads(compressed_content)
+
+            # Find the array in original
+            for _key, value in original_data.items():
+                if isinstance(value, list):
+                    items_before = len(value)
+                    break
+
+            # Find the array in compressed
+            for _key, value in compressed_data.items():
+                if isinstance(value, list):
+                    items_after = len(value)
+                    # Count errors preserved
+                    for item in value:
+                        item_str = str(item).lower()
+                        if any(kw in item_str for kw in profile.preserve_error_keywords):
+                            errors_preserved += 1
+                    break
+        except (json.JSONDecodeError, AttributeError):
+            pass
+
+        compressed_tokens = self._count_tokens(compressed_content)
+        tokens_saved = original_tokens - compressed_tokens
+        compression_ratio = tokens_saved / original_tokens if original_tokens > 0 else 0.0
+
+        return MCPCompressionResult(
+            original_content=content,
+            compressed_content=compressed_content,
+            original_tokens=original_tokens,
+            compressed_tokens=compressed_tokens,
+            tokens_saved=tokens_saved,
+            compression_ratio=compression_ratio,
+            items_before=items_before,
+            items_after=items_after,
+            errors_preserved=errors_preserved,
+            was_compressed=True,
+            tool_name=tool_name,
+            context_used=context,
+        )
+
+
+def compress_tool_result(
+    content: str,
+    tool_name: str,
+    tool_args: dict[str, Any] | None = None,
+    user_query: str = "",
+    config: HeadroomConfig | None = None,
+) -> str:
+    """Compress an MCP tool result (standalone function).
+
+    This is the simplest way to use Headroom with MCP in your host application.
+
+    Args:
+        content: Raw tool output.
+        tool_name: Name of the MCP tool.
+        tool_args: Arguments passed to the tool.
+        user_query: User's query for relevance scoring.
+        config: Optional Headroom configuration.
+
+    Returns:
+        Compressed content string.
+
+    Example:
+        ```python
+        # In your MCP host application
+        result = await client.call_tool("search_logs", {"service": "api"})
+
+        compressed = compress_tool_result(
+            content=result,
+            tool_name="search_logs",
+            tool_args={"service": "api"},
+            user_query="find errors",
+        )
+
+        messages.append({"role": "tool", "content": compressed})
+        ```
+    """
+    compressor = HeadroomMCPCompressor(config=config)
+    result = compressor.compress(
+        content=content,
+        tool_name=tool_name,
+        tool_args=tool_args,
+        user_query=user_query,
+    )
+    return result.compressed_content
+
+
+def compress_tool_result_with_metrics(
+    content: str,
+    tool_name: str,
+    tool_args: dict[str, Any] | None = None,
+    user_query: str = "",
+    config: HeadroomConfig | None = None,
+) -> MCPCompressionResult:
+    """Compress an MCP tool result and return full metrics.
+
+    Same as compress_tool_result but returns detailed metrics.
+
+    Returns:
+        MCPCompressionResult with all compression metrics.
+    """
+    compressor = HeadroomMCPCompressor(config=config)
+    return compressor.compress(
+        content=content,
+        tool_name=tool_name,
+        tool_args=tool_args,
+        user_query=user_query,
+    )
+
+
+class HeadroomMCPClientWrapper:
+    """Wrapper for MCP clients that automatically compresses tool results.
+
+    This wraps an MCP client to transparently compress all tool outputs.
+
+    Example:
+        ```python
+        from mcp import Client
+        from headroom.integrations.mcp import HeadroomMCPClientWrapper
+
+        # Wrap your MCP client
+        base_client = Client(transport)
+        client = HeadroomMCPClientWrapper(base_client)
+
+        # Use normally - compression is automatic
+        result = await client.call_tool("search", {"query": "errors"})
+        ```
+    """
+
+    def __init__(
+        self,
+        client: Any,
+        config: HeadroomConfig | None = None,
+        user_query_extractor: Callable[[dict], str] | None = None,
+    ):
+        """Initialize wrapper.
+
+        Args:
+            client: The MCP client to wrap.
+            config: Headroom configuration.
+            user_query_extractor: Function to extract user query from context.
+        """
+        self._client = client
+        self._compressor = HeadroomMCPCompressor(config=config)
+        self._query_extractor = user_query_extractor or (lambda x: "")
+        self._metrics: list[MCPCompressionResult] = []
+
+    async def call_tool(
+        self,
+        name: str,
+        arguments: dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
+    ) -> str:
+        """Call an MCP tool and compress the result.
+
+        Args:
+            name: Tool name.
+            arguments: Tool arguments.
+            context: Optional context (for query extraction).
+
+        Returns:
+            Compressed tool result.
+        """
+        # Call the underlying client
+        result = await self._client.call_tool(name, arguments)
+
+        # Extract user query from context if available
+        user_query = ""
+        if context and self._query_extractor is not None:
+            user_query = self._query_extractor(context)
+
+        # Compress
+        compression_result = self._compressor.compress(
+            content=result,
+            tool_name=name,
+            tool_args=arguments,
+            user_query=user_query,
+        )
+
+        self._metrics.append(compression_result)
+        return compression_result.compressed_content
+
+    def get_metrics(self) -> list[MCPCompressionResult]:
+        """Get compression metrics for all tool calls."""
+        return self._metrics.copy()
+
+    def get_total_tokens_saved(self) -> int:
+        """Get total tokens saved across all tool calls."""
+        return sum(m.tokens_saved for m in self._metrics)
+
+    def __getattr__(self, name: str) -> Any:
+        """Forward all other attributes to the wrapped client."""
+        return getattr(self._client, name)
+
+
+# Type alias for MCP Server (will be properly typed when mcp package is used)
+MCPServer = Any
+
+
+def create_headroom_mcp_proxy(
+    upstream_servers: list[tuple[str, MCPServer]],
+    config: HeadroomConfig | None = None,
+) -> dict[str, Any]:
+    """Create configuration for a Headroom MCP proxy server.
+
+    This returns a configuration dict that can be used to set up
+    a proxy server that wraps upstream MCP servers.
+
+    Args:
+        upstream_servers: List of (name, server) tuples.
+        config: Headroom configuration.
+
+    Returns:
+        Configuration dict for proxy server.
+
+    Example:
+        ```python
+        # In your MCP server setup
+        proxy_config = create_headroom_mcp_proxy(
+            upstream_servers=[
+                ("slack", slack_server),
+                ("database", db_server),
+            ]
+        )
+
+        # Use proxy_config to initialize your proxy server
+        ```
+    """
+    return {
+        "upstream_servers": dict(upstream_servers),
+        "compressor": HeadroomMCPCompressor(config=config),
+        "config": config or HeadroomConfig(),
+    }

headroom/memory/__init__.py

@@ -0,0 +1,37 @@
+"""Headroom Memory - Simple, LLM-driven memory for AI applications.
+
+Two approaches available:
+
+1. Background extraction (original):
+    from headroom import with_memory
+    client = with_memory(OpenAI(), user_id="alice")
+
+2. Zero-latency inline extraction (recommended):
+    from headroom.memory import with_fast_memory
+    client = with_fast_memory(OpenAI(), user_id="alice")
+"""
+
+from headroom.memory.fast_store import FastMemoryStore, MemoryChunk
+from headroom.memory.fast_wrapper import with_fast_memory
+from headroom.memory.inline_extractor import (
+    InlineMemoryWrapper,
+    inject_memory_instruction,
+    parse_response_with_memory,
+)
+from headroom.memory.store import Memory, SQLiteMemoryStore
+from headroom.memory.wrapper import with_memory
+
+__all__ = [
+    # Original approach (background extraction)
+    "with_memory",
+    "Memory",
+    "SQLiteMemoryStore",
+    # Fast approach (inline extraction - recommended)
+    "with_fast_memory",
+    "FastMemoryStore",
+    "MemoryChunk",
+    # Low-level inline extraction
+    "InlineMemoryWrapper",
+    "inject_memory_instruction",
+    "parse_response_with_memory",
+]