headroom-ai 0.2.13__py3-none-any.whl

headroom/integrations/agno/model.py

@@ -0,0 +1,625 @@
+"""Agno model wrapper for Headroom optimization.
+
+This module provides HeadroomAgnoModel, which wraps any Agno model
+to apply Headroom context optimization before API calls.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+import warnings
+from collections.abc import AsyncIterator, Iterator
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+from uuid import uuid4
+
+# Agno imports - these are optional dependencies
+try:
+    from agno.models.base import Model
+    from agno.models.message import Message
+    from agno.models.response import ModelResponse
+
+    AGNO_AVAILABLE = True
+except ImportError:
+    AGNO_AVAILABLE = False
+    Model = object  # type: ignore[misc,assignment]
+    Message = dict  # type: ignore[misc,assignment]
+    ModelResponse = dict  # type: ignore[misc,assignment]
+
+from headroom import HeadroomConfig, HeadroomMode
+from headroom.providers import OpenAIProvider
+from headroom.transforms import TransformPipeline
+
+from .providers import get_headroom_provider, get_model_name_from_agno
+
+logger = logging.getLogger(__name__)
+
+
+def _check_agno_available() -> None:
+    """Raise ImportError if Agno is not installed."""
+    if not AGNO_AVAILABLE:
+        raise ImportError("Agno is required for this integration. Install with: pip install agno")
+
+
+def agno_available() -> bool:
+    """Check if Agno is installed."""
+    return AGNO_AVAILABLE
+
+
+@dataclass
+class OptimizationMetrics:
+    """Metrics from a single optimization pass."""
+
+    request_id: str
+    timestamp: datetime
+    tokens_before: int
+    tokens_after: int
+    tokens_saved: int
+    savings_percent: float
+    transforms_applied: list[str]
+    model: str
+
+
+@dataclass
+class HeadroomAgnoModel(Model):  # type: ignore[misc]
+    """Agno model wrapper that applies Headroom optimizations.
+
+    Extends agno.models.base.Model to be fully compatible with Agno Agent.
+    Wraps any Agno Model and automatically optimizes the context
+    before each API call. Works with OpenAIChat, Claude, Gemini, and
+    other Agno model types.
+
+    Example:
+        from agno.agent import Agent
+        from agno.models.openai import OpenAIChat
+        from headroom.integrations.agno import HeadroomAgnoModel
+
+        # Basic usage
+        model = OpenAIChat(id="gpt-4o")
+        optimized = HeadroomAgnoModel(wrapped_model=model)
+
+        # Use with agent
+        agent = Agent(model=optimized)
+        response = agent.run("Hello!")
+
+        # Access metrics
+        print(f"Saved {optimized.total_tokens_saved} tokens")
+
+        # With custom config
+        from headroom import HeadroomConfig, HeadroomMode
+        config = HeadroomConfig(default_mode=HeadroomMode.OPTIMIZE)
+        optimized = HeadroomAgnoModel(wrapped_model=model, headroom_config=config)
+
+    Attributes:
+        wrapped_model: The underlying Agno model
+        total_tokens_saved: Running total of tokens saved
+        metrics_history: List of OptimizationMetrics from recent calls
+    """
+
+    # Required by Model base class - we'll derive from wrapped model
+    id: str = field(default="headroom-wrapper")
+    name: str | None = field(default=None)
+    provider: str | None = field(default=None)
+
+    # HeadroomAgnoModel specific fields
+    wrapped_model: Any = field(default=None)
+    headroom_config: HeadroomConfig | None = field(default=None)
+    headroom_mode: HeadroomMode | None = field(default=None)
+    auto_detect_provider: bool = field(default=True)
+
+    # Internal state (not part of dataclass comparison)
+    _metrics_history: list[OptimizationMetrics] = field(
+        default_factory=list, repr=False, compare=False
+    )
+    _total_tokens_saved: int = field(default=0, repr=False, compare=False)
+    _pipeline: TransformPipeline | None = field(default=None, repr=False, compare=False)
+    _headroom_provider: Any = field(default=None, repr=False, compare=False)
+    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False, compare=False)
+    _initialized: bool = field(default=False, repr=False, compare=False)
+
+    def __post_init__(self) -> None:
+        """Initialize HeadroomAgnoModel after dataclass construction."""
+        _check_agno_available()
+
+        if self.wrapped_model is None:
+            raise ValueError("wrapped_model cannot be None")
+
+        # Set id from wrapped model
+        if hasattr(self.wrapped_model, "id"):
+            self.id = f"headroom:{self.wrapped_model.id}"
+
+        # Set name and provider from wrapped model for compatibility
+        if self.name is None and hasattr(self.wrapped_model, "name"):
+            self.name = self.wrapped_model.name
+        if self.provider is None and hasattr(self.wrapped_model, "provider"):
+            self.provider = self.wrapped_model.provider
+
+        # Initialize config
+        if self.headroom_config is None:
+            self.headroom_config = HeadroomConfig()
+
+        # Handle deprecated mode parameter
+        if self.headroom_mode is not None:
+            warnings.warn(
+                "The 'headroom_mode' parameter is deprecated. Use HeadroomConfig(default_mode=...) instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        self._initialized = True
+
+        # Call parent __post_init__ if it exists
+        if hasattr(super(), "__post_init__"):
+            super().__post_init__()
+
+    # Forward attribute access to wrapped model for compatibility
+    def __getattr__(self, name: str) -> Any:
+        """Forward attribute access to wrapped model."""
+        # Avoid infinite recursion during initialization
+        if name.startswith("_") or not self.__dict__.get("_initialized", False):
+            raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")
+        if name in (
+            "wrapped_model",
+            "headroom_config",
+            "headroom_mode",
+            "auto_detect_provider",
+            "pipeline",
+            "total_tokens_saved",
+            "metrics_history",
+            "id",
+            "name",
+            "provider",
+        ):
+            raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")
+        return getattr(self.wrapped_model, name)
+
+    @property
+    def pipeline(self) -> TransformPipeline:
+        """Lazily initialize TransformPipeline (thread-safe)."""
+        if self._pipeline is None:
+            with self._lock:
+                # Double-check after acquiring lock
+                if self._pipeline is None:
+                    if self.auto_detect_provider:
+                        self._headroom_provider = get_headroom_provider(self.wrapped_model)
+                        logger.debug(
+                            f"Auto-detected provider: {self._headroom_provider.__class__.__name__}"
+                        )
+                    else:
+                        self._headroom_provider = OpenAIProvider()
+                    self._pipeline = TransformPipeline(
+                        config=self.headroom_config,
+                        provider=self._headroom_provider,
+                    )
+        return self._pipeline
+
+    @property
+    def total_tokens_saved(self) -> int:
+        """Total tokens saved across all calls."""
+        return self._total_tokens_saved
+
+    @property
+    def metrics_history(self) -> list[OptimizationMetrics]:
+        """History of optimization metrics."""
+        return self._metrics_history.copy()
+
+    def _convert_messages_to_openai(self, messages: list[Any]) -> list[dict[str, Any]]:
+        """Convert Agno messages to OpenAI format for Headroom."""
+        result = []
+        for msg in messages:
+            # Handle Agno Message objects
+            if hasattr(msg, "role") and hasattr(msg, "content"):
+                entry: dict[str, Any] = {
+                    "role": msg.role,
+                    "content": msg.content if msg.content is not None else "",
+                }
+                # Handle tool calls
+                if hasattr(msg, "tool_calls") and msg.tool_calls:
+                    entry["tool_calls"] = msg.tool_calls
+                # Handle tool call ID for tool responses
+                if hasattr(msg, "tool_call_id") and msg.tool_call_id:
+                    entry["tool_call_id"] = msg.tool_call_id
+                result.append(entry)
+            # Handle dict format
+            elif isinstance(msg, dict):
+                result.append(msg.copy())
+            else:
+                # Try to extract content
+                content = str(msg) if msg is not None else ""
+                result.append({"role": "user", "content": content})
+        return result
+
+    def _convert_messages_from_openai(self, messages: list[dict[str, Any]]) -> list[Any]:
+        """Convert OpenAI format messages back to Agno format.
+
+        Note: Agno typically accepts OpenAI-format dicts directly,
+        so we may not need full conversion.
+        """
+        # Agno models generally accept OpenAI-format messages
+        # Return as-is for compatibility
+        return messages
+
+    def _optimize_messages(self, messages: list[Any]) -> tuple[list[Any], OptimizationMetrics]:
+        """Apply Headroom optimization to messages.
+
+        Thread-safe with fallback on pipeline errors.
+        """
+        request_id = str(uuid4())
+
+        # Convert to OpenAI format
+        openai_messages = self._convert_messages_to_openai(messages)
+
+        # Handle empty messages gracefully
+        if not openai_messages:
+            metrics = OptimizationMetrics(
+                request_id=request_id,
+                timestamp=datetime.now(timezone.utc),
+                tokens_before=0,
+                tokens_after=0,
+                tokens_saved=0,
+                savings_percent=0,
+                transforms_applied=[],
+                model=get_model_name_from_agno(self.wrapped_model),
+            )
+            return openai_messages, metrics
+
+        # Get model name from wrapped model
+        model = get_model_name_from_agno(self.wrapped_model)
+
+        # Ensure pipeline is initialized
+        _ = self.pipeline
+
+        # Get model context limit
+        model_limit = (
+            self._headroom_provider.get_context_limit(model) if self._headroom_provider else 128000
+        )
+
+        try:
+            # Apply Headroom transforms via pipeline
+            result = self.pipeline.apply(
+                messages=openai_messages,
+                model=model,
+                model_limit=model_limit,
+            )
+            optimized = result.messages
+            tokens_before = result.tokens_before
+            tokens_after = result.tokens_after
+            transforms_applied = result.transforms_applied
+        except (
+            ValueError,
+            TypeError,
+            AttributeError,
+            RuntimeError,
+            KeyError,
+            IndexError,
+            ImportError,
+            OSError,
+        ) as e:
+            # Fallback to original messages on pipeline error
+            # Log at warning level (degraded behavior, not critical failure)
+            logger.warning(
+                f"Headroom optimization failed, using original messages: {type(e).__name__}: {e}"
+            )
+            optimized = openai_messages
+            # Estimate token count for unoptimized messages (rough approximation)
+            # Note: This uses ~4 chars/token which is approximate for English text
+            tokens_before = sum(len(str(m.get("content", ""))) // 4 for m in openai_messages)
+            tokens_after = tokens_before  # No optimization occurred
+            transforms_applied = ["fallback:error"]
+
+        # Create metrics
+        tokens_saved = max(0, tokens_before - tokens_after)  # Never negative
+        metrics = OptimizationMetrics(
+            request_id=request_id,
+            timestamp=datetime.now(timezone.utc),
+            tokens_before=tokens_before,
+            tokens_after=tokens_after,
+            tokens_saved=tokens_saved,
+            savings_percent=(tokens_saved / tokens_before * 100 if tokens_before > 0 else 0),
+            transforms_applied=transforms_applied,
+            model=model,
+        )
+
+        # Track metrics (thread-safe)
+        with self._lock:
+            self._metrics_history.append(metrics)
+            self._total_tokens_saved += metrics.tokens_saved
+
+            # Keep only last 100 metrics
+            if len(self._metrics_history) > 100:
+                self._metrics_history = self._metrics_history[-100:]
+
+        # Convert back (Agno accepts OpenAI format)
+        optimized_messages = self._convert_messages_from_openai(optimized)
+
+        return optimized_messages, metrics
+
+    def response(self, messages: list[Any], **kwargs: Any) -> Any:  # type: ignore[override]
+        """Generate response with Headroom optimization.
+
+        This is the core method that Agno agents call.
+        """
+        # Optimize messages
+        optimized_messages, metrics = self._optimize_messages(messages)
+
+        logger.info(
+            f"Headroom optimized: {metrics.tokens_before} -> {metrics.tokens_after} tokens "
+            f"({metrics.savings_percent:.1f}% saved)"
+        )
+
+        # Call wrapped model with optimized messages
+        return self.wrapped_model.response(optimized_messages, **kwargs)
+
+    def response_stream(self, messages: list[Any], **kwargs: Any) -> Iterator[Any]:  # type: ignore[override]
+        """Stream response with Headroom optimization."""
+        # Optimize messages
+        optimized_messages, metrics = self._optimize_messages(messages)
+
+        logger.info(
+            f"Headroom optimized (streaming): {metrics.tokens_before} -> "
+            f"{metrics.tokens_after} tokens"
+        )
+
+        # Stream from wrapped model
+        yield from self.wrapped_model.response_stream(optimized_messages, **kwargs)
+
+    async def aresponse(self, messages: list[Any], **kwargs: Any) -> Any:  # type: ignore[override]
+        """Async generate response with Headroom optimization."""
+        # Run optimization in executor (CPU-bound)
+        loop = asyncio.get_running_loop()
+        optimized_messages, metrics = await loop.run_in_executor(
+            None, self._optimize_messages, messages
+        )
+
+        logger.info(
+            f"Headroom optimized (async): {metrics.tokens_before} -> {metrics.tokens_after} tokens "
+            f"({metrics.savings_percent:.1f}% saved)"
+        )
+
+        # Call wrapped model's async method
+        if hasattr(self.wrapped_model, "aresponse"):
+            return await self.wrapped_model.aresponse(optimized_messages, **kwargs)
+        else:
+            # Fallback to sync in executor (non-blocking)
+            return await loop.run_in_executor(
+                None, lambda: self.wrapped_model.response(optimized_messages, **kwargs)
+            )
+
+    async def aresponse_stream(self, messages: list[Any], **kwargs: Any) -> AsyncIterator[Any]:  # type: ignore[override]
+        """Async stream response with Headroom optimization."""
+        # Run optimization in executor (CPU-bound)
+        loop = asyncio.get_running_loop()
+        optimized_messages, metrics = await loop.run_in_executor(
+            None, self._optimize_messages, messages
+        )
+
+        logger.info(
+            f"Headroom optimized (async streaming): {metrics.tokens_before} -> "
+            f"{metrics.tokens_after} tokens"
+        )
+
+        # Async stream from wrapped model
+        if hasattr(self.wrapped_model, "aresponse_stream"):
+            async for chunk in self.wrapped_model.aresponse_stream(optimized_messages, **kwargs):
+                yield chunk
+        else:
+            # Fallback: wrap sync streaming in async iterator (non-blocking)
+            # Run the entire sync iteration in executor to avoid blocking event loop
+            def _sync_stream() -> list[Any]:
+                return list(self.wrapped_model.response_stream(optimized_messages, **kwargs))
+
+            chunks = await loop.run_in_executor(None, _sync_stream)
+            for chunk in chunks:
+                yield chunk
+
+    def get_savings_summary(self) -> dict[str, Any]:
+        """Get summary of token savings."""
+        if not self._metrics_history:
+            return {
+                "total_requests": 0,
+                "total_tokens_saved": 0,
+                "average_savings_percent": 0,
+            }
+
+        return {
+            "total_requests": len(self._metrics_history),
+            "total_tokens_saved": self._total_tokens_saved,
+            "average_savings_percent": sum(m.savings_percent for m in self._metrics_history)
+            / len(self._metrics_history),
+            "total_tokens_before": sum(m.tokens_before for m in self._metrics_history),
+            "total_tokens_after": sum(m.tokens_after for m in self._metrics_history),
+        }
+
+    def reset(self) -> None:
+        """Reset all tracked metrics (thread-safe).
+
+        Clears the metrics history and resets the total tokens saved counter.
+        Useful for starting fresh measurements or between test runs.
+        """
+        with self._lock:
+            self._metrics_history = []
+            self._total_tokens_saved = 0
+
+    # =========================================================================
+    # Abstract method implementations required by agno.models.base.Model
+    # These delegate to the wrapped model after applying Headroom optimization
+    # =========================================================================
+
+    def invoke(self, messages: list[Any], **kwargs: Any) -> Any:
+        """Invoke the wrapped model with optimized messages.
+
+        This is required by agno.models.base.Model abstract interface.
+        """
+        # Optimize messages before invoking
+        optimized_messages, metrics = self._optimize_messages(messages)
+
+        logger.info(
+            f"Headroom optimized (invoke): {metrics.tokens_before} -> {metrics.tokens_after} tokens "
+            f"({metrics.savings_percent:.1f}% saved)"
+        )
+
+        # Delegate to wrapped model
+        return self.wrapped_model.invoke(optimized_messages, **kwargs)
+
+    async def ainvoke(self, messages: list[Any], **kwargs: Any) -> Any:
+        """Async invoke the wrapped model with optimized messages.
+
+        This is required by agno.models.base.Model abstract interface.
+        """
+        # Run optimization in executor (CPU-bound)
+        loop = asyncio.get_running_loop()
+        optimized_messages, metrics = await loop.run_in_executor(
+            None, self._optimize_messages, messages
+        )
+
+        logger.info(
+            f"Headroom optimized (ainvoke): {metrics.tokens_before} -> {metrics.tokens_after} tokens "
+            f"({metrics.savings_percent:.1f}% saved)"
+        )
+
+        # Delegate to wrapped model
+        if hasattr(self.wrapped_model, "ainvoke"):
+            return await self.wrapped_model.ainvoke(optimized_messages, **kwargs)
+        else:
+            # Fallback to sync in executor
+            return await loop.run_in_executor(
+                None, lambda: self.wrapped_model.invoke(optimized_messages, **kwargs)
+            )
+
+    def invoke_stream(self, messages: list[Any], **kwargs: Any) -> Iterator[Any]:
+        """Stream invoke the wrapped model with optimized messages.
+
+        This is required by agno.models.base.Model abstract interface.
+        """
+        # Optimize messages before streaming
+        optimized_messages, metrics = self._optimize_messages(messages)
+
+        logger.info(
+            f"Headroom optimized (invoke_stream): {metrics.tokens_before} -> {metrics.tokens_after} tokens "
+            f"({metrics.savings_percent:.1f}% saved)"
+        )
+
+        # Delegate to wrapped model
+        yield from self.wrapped_model.invoke_stream(optimized_messages, **kwargs)
+
+    async def ainvoke_stream(self, messages: list[Any], **kwargs: Any) -> AsyncIterator[Any]:
+        """Async stream invoke the wrapped model with optimized messages.
+
+        This is required by agno.models.base.Model abstract interface.
+        """
+        # Run optimization in executor (CPU-bound)
+        loop = asyncio.get_running_loop()
+        optimized_messages, metrics = await loop.run_in_executor(
+            None, self._optimize_messages, messages
+        )
+
+        logger.info(
+            f"Headroom optimized (ainvoke_stream): {metrics.tokens_before} -> {metrics.tokens_after} tokens "
+            f"({metrics.savings_percent:.1f}% saved)"
+        )
+
+        # Delegate to wrapped model
+        if hasattr(self.wrapped_model, "ainvoke_stream"):
+            async for chunk in self.wrapped_model.ainvoke_stream(optimized_messages, **kwargs):
+                yield chunk
+        else:
+            # Fallback: wrap sync streaming
+            def _sync_stream() -> list[Any]:
+                return list(self.wrapped_model.invoke_stream(optimized_messages, **kwargs))
+
+            chunks = await loop.run_in_executor(None, _sync_stream)
+            for chunk in chunks:
+                yield chunk
+
+    def _parse_provider_response(self, response: Any, **kwargs: Any) -> Any:
+        """Parse provider response - delegates to wrapped model.
+
+        This is required by agno.models.base.Model abstract interface.
+        """
+        return self.wrapped_model._parse_provider_response(response, **kwargs)
+
+    def _parse_provider_response_delta(self, response: Any) -> Any:
+        """Parse streaming response delta - delegates to wrapped model.
+
+        This is required by agno.models.base.Model abstract interface.
+        """
+        return self.wrapped_model._parse_provider_response_delta(response)
+
+
+def optimize_messages(
+    messages: list[Any],
+    config: HeadroomConfig | None = None,
+    mode: HeadroomMode = HeadroomMode.OPTIMIZE,
+    model: str = "gpt-4o",
+) -> tuple[list[dict[str, Any]], dict[str, Any]]:
+    """Standalone function to optimize Agno messages.
+
+    Use this for manual optimization when you need fine-grained control.
+
+    Args:
+        messages: List of Agno Message objects or dicts
+        config: HeadroomConfig for optimization settings
+        mode: HeadroomMode (AUDIT, OPTIMIZE, or SIMULATE)
+        model: Model name for token estimation
+
+    Returns:
+        Tuple of (optimized_messages, metrics_dict)
+
+    Example:
+        from headroom.integrations.agno import optimize_messages
+
+        messages = [
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "What is 2+2?"},
+        ]
+
+        optimized, metrics = optimize_messages(messages)
+        print(f"Saved {metrics['tokens_saved']} tokens")
+    """
+    _check_agno_available()
+
+    config = config or HeadroomConfig()
+    provider = OpenAIProvider()
+    pipeline = TransformPipeline(config=config, provider=provider)
+
+    # Convert to OpenAI format
+    openai_messages = []
+    for msg in messages:
+        if hasattr(msg, "role") and hasattr(msg, "content"):
+            entry: dict[str, Any] = {"role": msg.role, "content": msg.content or ""}
+            if hasattr(msg, "tool_calls") and msg.tool_calls:
+                entry["tool_calls"] = msg.tool_calls
+            if hasattr(msg, "tool_call_id") and msg.tool_call_id:
+                entry["tool_call_id"] = msg.tool_call_id
+            openai_messages.append(entry)
+        elif isinstance(msg, dict):
+            openai_messages.append(msg.copy())
+        else:
+            openai_messages.append({"role": "user", "content": str(msg)})
+
+    # Get model context limit
+    model_limit = provider.get_context_limit(model)
+
+    # Apply transforms
+    result = pipeline.apply(
+        messages=openai_messages,
+        model=model,
+        model_limit=model_limit,
+    )
+
+    metrics = {
+        "tokens_before": result.tokens_before,
+        "tokens_after": result.tokens_after,
+        "tokens_saved": result.tokens_before - result.tokens_after,
+        "savings_percent": (
+            (result.tokens_before - result.tokens_after) / result.tokens_before * 100
+            if result.tokens_before > 0
+            else 0
+        ),
+        "transforms_applied": result.transforms_applied,
+    }
+
+    return result.messages, metrics