headroom-ai 0.2.13__py3-none-any.whl

headroom/integrations/agno/providers.py

@@ -0,0 +1,154 @@
+"""Provider detection for Agno models.
+
+Automatically detects the correct Headroom provider based on the Agno model type.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from headroom.providers import (
+    AnthropicProvider,
+    CohereProvider,
+    GoogleProvider,
+    OpenAIProvider,
+)
+from headroom.providers.base import Provider
+
+logger = logging.getLogger(__name__)
+
+# Mapping from Agno model class names to Headroom providers
+_AGNO_MODEL_PROVIDERS: dict[str, type[Provider]] = {
+    # OpenAI models
+    "OpenAIChat": OpenAIProvider,
+    "OpenAILike": OpenAIProvider,
+    # Anthropic models (direct and cloud variants)
+    "Claude": AnthropicProvider,
+    "Anthropic": AnthropicProvider,
+    "AwsBedrock": AnthropicProvider,  # Bedrock Claude models
+    "BedrockClaude": AnthropicProvider,
+    # Google models
+    "Gemini": GoogleProvider,
+    "GoogleGenerativeAI": GoogleProvider,
+    "VertexAI": GoogleProvider,
+    # LiteLLM (uses OpenAI-compatible interface, provider detected from model ID)
+    "LiteLLM": OpenAIProvider,
+    "LiteLLMChat": OpenAIProvider,
+    # Others default to OpenAI-compatible tokenization
+    "Groq": OpenAIProvider,
+    "Mistral": OpenAIProvider,
+    "MistralChat": OpenAIProvider,
+    "Together": OpenAIProvider,
+    "TogetherChat": OpenAIProvider,
+    "Fireworks": OpenAIProvider,
+    "FireworksChat": OpenAIProvider,
+    "Ollama": OpenAIProvider,
+    "OllamaChat": OpenAIProvider,
+    "DeepSeek": OpenAIProvider,
+    "DeepSeekChat": OpenAIProvider,
+    "xAI": OpenAIProvider,
+    "XAI": OpenAIProvider,
+    "Grok": OpenAIProvider,
+    "Cohere": CohereProvider,
+    "CohereChat": CohereProvider,
+    "Perplexity": OpenAIProvider,
+    "Anyscale": OpenAIProvider,
+    "OpenRouter": OpenAIProvider,
+    "Replicate": OpenAIProvider,
+    "HuggingFace": OpenAIProvider,
+    "HuggingFaceChat": OpenAIProvider,
+}
+
+
+def get_headroom_provider(agno_model: Any) -> Provider:
+    """Get the appropriate Headroom provider for an Agno model.
+
+    Detection strategy:
+    1. Check model class name against known Agno model types
+    2. Check for provider hints in model attributes
+    3. Fall back to OpenAI provider (most compatible)
+
+    Args:
+        agno_model: An Agno model instance (OpenAIChat, Claude, etc.)
+
+    Returns:
+        Appropriate Headroom Provider instance.
+
+    Example:
+        from agno.models.openai import OpenAIChat
+        from headroom.integrations.agno.providers import get_headroom_provider
+
+        model = OpenAIChat(id="gpt-4o")
+        provider = get_headroom_provider(model)  # Returns OpenAIProvider
+    """
+    # Strategy 1: Class name matching
+    class_name = agno_model.__class__.__name__
+    if class_name in _AGNO_MODEL_PROVIDERS:
+        provider_class = _AGNO_MODEL_PROVIDERS[class_name]
+        logger.debug(f"Detected provider {provider_class.__name__} from class {class_name}")
+        return provider_class()
+
+    # Strategy 2: Check module path
+    module_path = agno_model.__class__.__module__
+    if "anthropic" in module_path.lower():
+        logger.debug(f"Detected AnthropicProvider from module {module_path}")
+        return AnthropicProvider()
+    elif "google" in module_path.lower() or "gemini" in module_path.lower():
+        logger.debug(f"Detected GoogleProvider from module {module_path}")
+        return GoogleProvider()
+    elif "cohere" in module_path.lower():
+        logger.debug(f"Detected CohereProvider from module {module_path}")
+        return CohereProvider()
+    elif "openai" in module_path.lower() or "litellm" in module_path.lower():
+        logger.debug(f"Detected OpenAIProvider from module {module_path}")
+        return OpenAIProvider()
+
+    # Strategy 3: Check model ID/name for hints
+    model_id = getattr(agno_model, "id", "") or getattr(agno_model, "model", "")
+    if isinstance(model_id, str) and model_id:
+        model_id_lower = model_id.lower()
+        if "claude" in model_id_lower:
+            logger.debug(f"Detected AnthropicProvider from model ID {model_id}")
+            return AnthropicProvider()
+        elif "gemini" in model_id_lower:
+            logger.debug(f"Detected GoogleProvider from model ID {model_id}")
+            return GoogleProvider()
+        elif "gpt" in model_id_lower or "o1" in model_id_lower or "o3" in model_id_lower:
+            logger.debug(f"Detected OpenAIProvider from model ID {model_id}")
+            return OpenAIProvider()
+        elif "command" in model_id_lower or "cohere" in model_id_lower:
+            logger.debug(f"Detected CohereProvider from model ID {model_id}")
+            return CohereProvider()
+
+    # Strategy 4: Default fallback
+    logger.warning(
+        f"Unknown Agno model class '{class_name}', defaulting to OpenAIProvider. "
+        "Token counting may be inaccurate."
+    )
+    return OpenAIProvider()
+
+
+def get_model_name_from_agno(agno_model: Any) -> str:
+    """Extract the model name/ID from an Agno model.
+
+    Args:
+        agno_model: An Agno model instance
+
+    Returns:
+        Model name string (e.g., "gpt-4o", "claude-3-5-sonnet-20241022")
+    """
+    # Try common attribute names
+    for attr in ["id", "model", "model_name", "model_id"]:
+        value = getattr(agno_model, attr, None)
+        if value and isinstance(value, str):
+            return str(value)
+
+    # Fallback with warning
+    class_name = agno_model.__class__.__name__
+    logger.warning(
+        f"Could not extract model name from {class_name} (no 'id', 'model', "
+        f"'model_name', or 'model_id' attribute). Defaulting to 'gpt-4o'. "
+        "Token counting may be inaccurate."
+    )
+    return "gpt-4o"

headroom/integrations/langchain/__init__.py

@@ -0,0 +1,106 @@
+"""LangChain integration for Headroom.
+
+This package provides seamless integration with LangChain, including:
+- HeadroomChatModel: Drop-in wrapper for any LangChain chat model
+- HeadroomChatMessageHistory: Automatic conversation compression
+- HeadroomDocumentCompressor: Relevance-based document filtering
+- HeadroomToolWrapper: Tool output compression for agents
+- StreamingMetricsTracker: Token counting during streaming
+- HeadroomLangSmithCallbackHandler: LangSmith trace enrichment
+
+Example:
+    from langchain_openai import ChatOpenAI
+    from headroom.integrations.langchain import HeadroomChatModel
+
+    # Wrap any LangChain model
+    llm = HeadroomChatModel(ChatOpenAI(model="gpt-4o"))
+
+    # Use like normal - optimization happens automatically
+    response = llm.invoke("Hello!")
+
+Install: pip install headroom[langchain]
+"""
+
+# Core chat model wrapper
+# Agent tool wrapping
+from .agents import (
+    HeadroomToolWrapper,
+    ToolCompressionMetrics,
+    ToolMetricsCollector,
+    get_tool_metrics,
+    reset_tool_metrics,
+    wrap_tools_with_headroom,
+)
+from .chat_model import (
+    HeadroomCallbackHandler,
+    HeadroomChatModel,
+    HeadroomRunnable,
+    OptimizationMetrics,
+    langchain_available,
+    optimize_messages,
+)
+
+# LangSmith integration
+from .langsmith import (
+    HeadroomLangSmithCallbackHandler,
+    is_langsmith_available,
+    is_langsmith_tracing_enabled,
+)
+
+# Memory integration
+from .memory import HeadroomChatMessageHistory
+
+# Provider auto-detection
+from .providers import (
+    detect_provider,
+    get_headroom_provider,
+    get_model_name_from_langchain,
+)
+
+# Retriever integration
+from .retriever import CompressionMetrics, HeadroomDocumentCompressor
+
+# Streaming metrics
+from .streaming import (
+    StreamingMetrics,
+    StreamingMetricsCallback,
+    StreamingMetricsTracker,
+    track_async_streaming_response,
+    track_streaming_response,
+)
+
+__all__ = [
+    # Core
+    "HeadroomChatModel",
+    "HeadroomCallbackHandler",
+    "HeadroomRunnable",
+    "OptimizationMetrics",
+    "optimize_messages",
+    "langchain_available",
+    # Provider Detection
+    "detect_provider",
+    "get_headroom_provider",
+    "get_model_name_from_langchain",
+    # Memory
+    "HeadroomChatMessageHistory",
+    # Retrievers
+    "HeadroomDocumentCompressor",
+    "CompressionMetrics",
+    # Agents
+    "HeadroomToolWrapper",
+    "ToolCompressionMetrics",
+    "ToolMetricsCollector",
+    "wrap_tools_with_headroom",
+    "get_tool_metrics",
+    "reset_tool_metrics",
+    # LangSmith
+    "HeadroomLangSmithCallbackHandler",
+    "is_langsmith_available",
+    "is_langsmith_tracing_enabled",
+    # Streaming
+    "StreamingMetricsTracker",
+    "StreamingMetricsCallback",
+    "StreamingMetrics",
+    "track_streaming_response",
+    "track_async_streaming_response",
+]

headroom/integrations/langchain/agents.py

@@ -0,0 +1,326 @@
+"""Agent tool integration for LangChain with output compression.
+
+This module provides HeadroomToolWrapper and wrap_tools_with_headroom
+for wrapping LangChain tools to automatically compress their outputs
+and track per-tool compression metrics.
+
+Example:
+    from langchain.agents import create_openai_tools_agent
+    from langchain.tools import Tool
+    from headroom.integrations import wrap_tools_with_headroom
+
+    # Define tools
+    tools = [
+        Tool(name="search", func=search_func, description="Search"),
+        Tool(name="database", func=db_func, description="Query DB"),
+    ]
+
+    # Wrap with Headroom compression
+    wrapped_tools = wrap_tools_with_headroom(tools)
+
+    # Use in agent - outputs are automatically compressed
+    agent = create_openai_tools_agent(llm, wrapped_tools, prompt)
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+# LangChain imports - these are optional dependencies
+try:
+    from langchain_core.tools import BaseTool, StructuredTool, Tool
+
+    LANGCHAIN_AVAILABLE = True
+except ImportError:
+    LANGCHAIN_AVAILABLE = False
+    BaseTool = object  # type: ignore[misc,assignment]
+    StructuredTool = object  # type: ignore[misc,assignment]
+    Tool = object  # type: ignore[misc,assignment]
+
+from headroom.integrations.mcp import compress_tool_result
+
+logger = logging.getLogger(__name__)
+
+
+def _check_langchain_available() -> None:
+    """Raise ImportError if LangChain is not installed."""
+    if not LANGCHAIN_AVAILABLE:
+        raise ImportError(
+            "LangChain is required for this integration. "
+            "Install with: pip install headroom[langchain] "
+            "or: pip install langchain-core"
+        )
+
+
+@dataclass
+class ToolCompressionMetrics:
+    """Metrics from a single tool compression."""
+
+    tool_name: str
+    timestamp: datetime
+    chars_before: int
+    chars_after: int
+    chars_saved: int
+    compression_ratio: float
+    was_compressed: bool
+
+
+@dataclass
+class ToolMetricsCollector:
+    """Collects compression metrics across all tool invocations."""
+
+    metrics: list[ToolCompressionMetrics] = field(default_factory=list)
+
+    def add(self, metric: ToolCompressionMetrics) -> None:
+        """Add a metric entry."""
+        self.metrics.append(metric)
+        # Keep only last 1000
+        if len(self.metrics) > 1000:
+            self.metrics = self.metrics[-1000:]
+
+    def get_summary(self) -> dict[str, Any]:
+        """Get summary statistics."""
+        if not self.metrics:
+            return {
+                "total_invocations": 0,
+                "total_compressions": 0,
+                "total_chars_saved": 0,
+            }
+
+        compressed = [m for m in self.metrics if m.was_compressed]
+        return {
+            "total_invocations": len(self.metrics),
+            "total_compressions": len(compressed),
+            "total_chars_saved": sum(m.chars_saved for m in self.metrics),
+            "average_compression_ratio": (
+                sum(m.compression_ratio for m in compressed) / len(compressed) if compressed else 0
+            ),
+            "by_tool": self._get_by_tool_stats(),
+        }
+
+    def _get_by_tool_stats(self) -> dict[str, dict[str, Any]]:
+        """Get per-tool statistics."""
+        by_tool: dict[str, list[ToolCompressionMetrics]] = {}
+        for m in self.metrics:
+            if m.tool_name not in by_tool:
+                by_tool[m.tool_name] = []
+            by_tool[m.tool_name].append(m)
+
+        result = {}
+        for name, tool_metrics in by_tool.items():
+            compressed = [m for m in tool_metrics if m.was_compressed]
+            result[name] = {
+                "invocations": len(tool_metrics),
+                "compressions": len(compressed),
+                "chars_saved": sum(m.chars_saved for m in tool_metrics),
+            }
+        return result
+
+
+# Global metrics collector
+_global_metrics = ToolMetricsCollector()
+
+
+def get_tool_metrics() -> ToolMetricsCollector:
+    """Get the global tool metrics collector."""
+    return _global_metrics
+
+
+def reset_tool_metrics() -> None:
+    """Reset global tool metrics."""
+    global _global_metrics
+    _global_metrics = ToolMetricsCollector()
+
+
+class HeadroomToolWrapper:
+    """Wraps a LangChain tool to compress its output.
+
+    Applies SmartCrusher compression to tool outputs, particularly
+    useful for tools that return large JSON arrays (search results,
+    database queries, etc.).
+
+    Example:
+        from langchain.tools import Tool
+        from headroom.integrations import HeadroomToolWrapper
+
+        def search(query: str) -> str:
+            # Returns large JSON with 1000 results
+            return json.dumps({"results": [...1000 items...]})
+
+        search_tool = Tool(name="search", func=search, description="Search")
+        wrapped = HeadroomToolWrapper(search_tool)
+
+        # Use wrapped tool - output automatically compressed
+        result = wrapped("python tutorials")
+
+    Attributes:
+        tool: The wrapped LangChain tool
+        min_chars_to_compress: Minimum output size to trigger compression
+        metrics_collector: Collector for compression metrics
+    """
+
+    def __init__(
+        self,
+        tool: BaseTool,
+        min_chars_to_compress: int = 1000,
+        metrics_collector: ToolMetricsCollector | None = None,
+    ):
+        """Initialize HeadroomToolWrapper.
+
+        Args:
+            tool: The LangChain BaseTool to wrap.
+            min_chars_to_compress: Minimum character count for output
+                before compression is applied. Default 1000.
+            metrics_collector: Collector for metrics. Uses global
+                collector if not specified.
+        """
+        _check_langchain_available()
+
+        self.tool = tool
+        self.min_chars_to_compress = min_chars_to_compress
+        self._metrics = metrics_collector or _global_metrics
+
+        # Copy tool metadata
+        self.name = tool.name
+        self.description = tool.description
+
+    def __call__(self, *args: Any, **kwargs: Any) -> str:
+        """Invoke the tool and compress output.
+
+        Args:
+            *args: Arguments to pass to the tool.
+            **kwargs: Keyword arguments to pass to the tool.
+
+        Returns:
+            Compressed tool output as string.
+        """
+        # Invoke underlying tool
+        result = self.tool.invoke(*args, **kwargs)
+
+        # Convert to string if needed
+        if not isinstance(result, str):
+            result = str(result)
+
+        # Check if compression is needed
+        if len(result) < self.min_chars_to_compress:
+            self._record_metrics(result, result, was_compressed=False)
+            return result
+
+        # Try to compress
+        compressed = self._compress_output(result)
+        self._record_metrics(result, compressed, was_compressed=True)
+
+        return compressed
+
+    def invoke(self, *args: Any, **kwargs: Any) -> str:
+        """Invoke the tool (alias for __call__)."""
+        return self(*args, **kwargs)
+
+    def _compress_output(self, output: str) -> str:
+        """Apply compression to tool output.
+
+        Args:
+            output: Tool output string.
+
+        Returns:
+            Compressed output.
+        """
+        try:
+            return compress_tool_result(
+                content=output,
+                tool_name=self.name,
+            )
+        except Exception as e:
+            logger.debug(f"Tool compression failed: {e}")
+            return output
+
+    def _record_metrics(self, original: str, compressed: str, was_compressed: bool) -> None:
+        """Record compression metrics.
+
+        Args:
+            original: Original output.
+            compressed: Compressed output.
+            was_compressed: Whether compression was applied.
+        """
+        chars_before = len(original)
+        chars_after = len(compressed)
+        chars_saved = chars_before - chars_after
+
+        metric = ToolCompressionMetrics(
+            tool_name=self.name,
+            timestamp=datetime.now(),
+            chars_before=chars_before,
+            chars_after=chars_after,
+            chars_saved=max(0, chars_saved),
+            compression_ratio=chars_after / chars_before if chars_before > 0 else 1.0,
+            was_compressed=was_compressed and chars_saved > 0,
+        )
+
+        self._metrics.add(metric)
+
+        if was_compressed and chars_saved > 0:
+            logger.info(
+                f"HeadroomToolWrapper[{self.name}]: {chars_before} -> {chars_after} chars "
+                f"({chars_saved} saved, {metric.compression_ratio:.1%} of original)"
+            )
+
+    def as_langchain_tool(self) -> StructuredTool:
+        """Convert wrapper back to a LangChain tool.
+
+        Useful when you need to pass the wrapped tool to APIs
+        that expect a LangChain tool type.
+
+        Returns:
+            StructuredTool that wraps this wrapper.
+        """
+        return StructuredTool.from_function(
+            func=self.__call__,
+            name=self.name,
+            description=self.description,
+        )
+
+
+def wrap_tools_with_headroom(
+    tools: list[BaseTool],
+    min_chars_to_compress: int = 1000,
+    metrics_collector: ToolMetricsCollector | None = None,
+) -> list[StructuredTool]:
+    """Wrap multiple LangChain tools with Headroom compression.
+
+    Convenience function to wrap all tools in a list at once.
+
+    Args:
+        tools: List of LangChain tools to wrap.
+        min_chars_to_compress: Minimum output size for compression.
+        metrics_collector: Shared metrics collector for all tools.
+
+    Returns:
+        List of wrapped tools as StructuredTools.
+
+    Example:
+        from langchain.tools import Tool
+        from headroom.integrations import wrap_tools_with_headroom
+
+        tools = [search_tool, database_tool, api_tool]
+        wrapped = wrap_tools_with_headroom(tools)
+
+        # Use wrapped tools in agent
+        agent = create_openai_tools_agent(llm, wrapped, prompt)
+    """
+    _check_langchain_available()
+
+    collector = metrics_collector or _global_metrics
+
+    wrapped = []
+    for tool in tools:
+        wrapper = HeadroomToolWrapper(
+            tool=tool,
+            min_chars_to_compress=min_chars_to_compress,
+            metrics_collector=collector,
+        )
+        wrapped.append(wrapper.as_langchain_tool())
+
+    return wrapped