chopratejas--headroom/headroom/integrations/strands/model.py

"""Strands SDK model wrapper for Headroom optimization.

This module provides HeadroomStrandsModel, which wraps any Strands model
to apply Headroom context optimization before API calls.
"""

from __future__ import annotations

import asyncio
import logging
import threading
from collections.abc import AsyncGenerator, AsyncIterable
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, TypeVar
from uuid import uuid4

# Strands imports - these are optional dependencies
try:
    from strands.models import Model
    from strands.types.content import Message, Messages, SystemContentBlock
    from strands.types.streaming import StreamEvent
    from strands.types.tools import ToolChoice, ToolSpec

    STRANDS_AVAILABLE = True
except ImportError:
    STRANDS_AVAILABLE = False
    Model = object  # type: ignore[misc,assignment]
    Message = dict  # type: ignore[misc,assignment]
    Messages = list  # type: ignore[misc,assignment]
    StreamEvent = dict  # type: ignore[misc,assignment]
    ToolChoice = dict  # type: ignore[misc,assignment]
    ToolSpec = dict  # type: ignore[misc,assignment]
    SystemContentBlock = dict  # type: ignore[misc,assignment]

T = TypeVar("T")

from headroom import HeadroomConfig  # noqa: E402
from headroom.providers import OpenAIProvider  # noqa: E402
from headroom.transforms import TransformPipeline  # noqa: E402

from .providers import get_headroom_provider, get_model_name_from_strands  # noqa: E402

logger = logging.getLogger(__name__)


def _check_strands_available() -> None:
    """Raise ImportError if Strands SDK is not installed."""
    if not STRANDS_AVAILABLE:
        raise ImportError(
            "Strands SDK is required for this integration. Install with: pip install strands-agents"
        )


def strands_available() -> bool:
    """Check if Strands SDK is installed."""
    return STRANDS_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


class HeadroomStrandsModel(Model):  # type: ignore[misc]
    """Strands model wrapper that applies Headroom optimizations.

    Wraps any Strands Model and automatically optimizes the context
    before each API call. Works with any Strands-compatible model provider.

    Example:
        from strands import Agent
        from strands.models.bedrock import BedrockModel
        from headroom.integrations.strands import HeadroomStrandsModel

        # Basic usage
        model = BedrockModel(model_id="us.anthropic.claude-sonnet-4-20250514-v1:0")
        optimized = HeadroomStrandsModel(wrapped_model=model)

        # Use with agent
        agent = Agent(model=optimized)
        response = agent("Hello!")

        # Access metrics
        print(f"Saved {optimized.total_tokens_saved} tokens")

        # With custom config
        from headroom import HeadroomConfig
        config = HeadroomConfig()
        optimized = HeadroomStrandsModel(wrapped_model=model, config=config)

    Attributes:
        wrapped_model: The underlying Strands model
        total_tokens_saved: Running total of tokens saved
        metrics_history: List of OptimizationMetrics from recent calls
    """

    def __init__(
        self,
        wrapped_model: Any,
        config: HeadroomConfig | None = None,
        auto_detect_provider: bool = True,
    ) -> None:
        """Initialize HeadroomStrandsModel.

        Args:
            wrapped_model: The Strands model to wrap (e.g., BedrockModel, OpenAIModel)
            config: Optional HeadroomConfig for optimization settings
            auto_detect_provider: Whether to auto-detect the Headroom provider
                based on the wrapped model type. Default True.
        """
        _check_strands_available()

        if wrapped_model is None:
            raise ValueError("wrapped_model cannot be None")

        self.wrapped_model = wrapped_model
        self.headroom_config = config or HeadroomConfig()
        self.auto_detect_provider = auto_detect_provider

        # Internal state
        self._metrics_history: list[OptimizationMetrics] = []
        self._total_tokens_saved: int = 0
        self._pipeline: TransformPipeline | None = None
        self._headroom_provider: Any = None
        self._lock = threading.Lock()

    @property
    def config(self) -> Any:
        """Forward config access to wrapped model (required by Strands Agent)."""
        return self.wrapped_model.config

    @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 Strands messages to OpenAI format for Headroom.

        Strands uses dict-based messages similar to OpenAI format:
        - {"role": "user", "content": "..."}
        - {"role": "assistant", "content": "...", "tool_calls": [...]}
        - {"role": "tool", "content": "...", "tool_call_id": "..."}

        Args:
            messages: List of Strands messages (typically dicts or Message objects)

        Returns:
            List of messages in OpenAI dict format
        """
        result = []
        for msg in messages:
            # Handle dict format (most common in Strands)
            if isinstance(msg, dict):
                entry: dict[str, Any] = {
                    "role": msg.get("role", "user"),
                }

                # Handle content
                content = msg.get("content")
                if content is None:
                    entry["content"] = ""
                elif isinstance(content, list):
                    # Content blocks - preserve structure
                    entry["content"] = content
                else:
                    entry["content"] = content

                # Handle tool calls
                if "tool_calls" in msg and msg["tool_calls"]:
                    entry["tool_calls"] = msg["tool_calls"]

                # Handle tool call ID for tool responses
                if "tool_call_id" in msg and msg["tool_call_id"]:
                    entry["tool_call_id"] = msg["tool_call_id"]

                # Handle name field (for tool messages)
                if "name" in msg and msg["name"]:
                    entry["name"] = msg["name"]

                result.append(entry)

            # Handle Strands Message objects (if they have role/content attrs)
            elif hasattr(msg, "role") and hasattr(msg, "content"):
                entry = {
                    "role": msg.role,
                }

                content = msg.content
                if content is None:
                    entry["content"] = ""
                elif isinstance(content, list):
                    entry["content"] = content
                else:
                    entry["content"] = content

                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
                if hasattr(msg, "name") and msg.name:
                    entry["name"] = msg.name

                result.append(entry)

            else:
                # Fallback: convert to string
                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]], original_messages: list[Any]
    ) -> list[dict[str, Any]]:
        """Convert OpenAI format messages back to Strands format.

        Since Strands uses dict-based messages similar to OpenAI,
        this is largely a passthrough, but ensures proper structure.

        Args:
            messages: The optimized messages in OpenAI dict format
            original_messages: The original Strands messages (for reference)

        Returns:
            List of messages in Strands dict format
        """
        result = []
        for msg in messages:
            entry: dict[str, Any] = {
                "role": msg.get("role", "user"),
            }

            # Handle content
            content = msg.get("content")
            if content is not None:
                entry["content"] = content

            # Preserve tool-related fields
            if "tool_calls" in msg and msg["tool_calls"]:
                entry["tool_calls"] = msg["tool_calls"]
            if "tool_call_id" in msg and msg["tool_call_id"]:
                entry["tool_call_id"] = msg["tool_call_id"]
            if "name" in msg and msg["name"]:
                entry["name"] = msg["name"]

            result.append(entry)

        return result

    def _optimize_messages(
        self, messages: list[Any]
    ) -> tuple[list[dict[str, Any]], OptimizationMetrics]:
        """Apply Headroom optimization to messages.

        Thread-safe with fallback on pipeline errors.

        Args:
            messages: List of Strands messages to optimize

        Returns:
            Tuple of (optimized_messages, metrics)
        """
        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_strands(self.wrapped_model),
            )
            return [], metrics

        # Get model name from wrapped model
        model = get_model_name_from_strands(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
            logger.warning(
                f"Headroom optimization failed, using original messages: {type(e).__name__}: {e}"
            )
            optimized = openai_messages
            # Estimate token count (rough approximation: ~4 chars/token)
            tokens_before = sum(len(str(m.get("content", ""))) // 4 for m in openai_messages)
            tokens_after = tokens_before
            transforms_applied = ["fallback:error"]

        # Create metrics
        tokens_saved = max(0, tokens_before - tokens_after)
        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 to Strands format
        optimized_messages = self._convert_messages_from_openai(optimized, messages)

        return optimized_messages, metrics

    async def stream(
        self,
        messages: Messages,
        tool_specs: list[ToolSpec] | None = None,
        system_prompt: str | None = None,
        *,
        tool_choice: ToolChoice | None = None,
        system_prompt_content: list[SystemContentBlock] | None = None,
        invocation_state: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> AsyncIterable[StreamEvent]:
        """Stream response with Headroom optimization.

        This is the main method required by Strands Model interface.
        Optimizes messages before delegating to the wrapped model's stream method.

        Args:
            messages: List of messages to send to the model
            tool_specs: Optional list of tool specifications
            system_prompt: Optional system prompt string
            tool_choice: Optional tool choice configuration
            system_prompt_content: Optional list of system content blocks
            invocation_state: Optional invocation state dictionary
            **kwargs: Additional arguments passed to the wrapped model

        Yields:
            Streaming events from the wrapped model
        """
        # 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 (stream): {metrics.tokens_before} -> "
            f"{metrics.tokens_after} tokens ({metrics.savings_percent:.1f}% saved)"
        )

        # Delegate to wrapped model's stream method with all parameters
        async for event in self.wrapped_model.stream(
            optimized_messages,
            tool_specs=tool_specs,
            system_prompt=system_prompt,
            tool_choice=tool_choice,
            system_prompt_content=system_prompt_content,
            invocation_state=invocation_state,
            **kwargs,
        ):
            yield event

    def get_config(self) -> Any:
        """Get the configuration of the wrapped model.

        Returns:
            The model configuration from the wrapped model.
        """
        return self.wrapped_model.get_config()

    def update_config(self, **model_config: Any) -> None:
        """Update the configuration of the wrapped model.

        Args:
            **model_config: Configuration options to update on the wrapped model.
        """
        self.wrapped_model.update_config(**model_config)

    async def structured_output(
        self,
        output_model: type[T],
        prompt: Messages,
        system_prompt: str | None = None,
        **kwargs: Any,
    ) -> AsyncGenerator[dict[str, T | Any], None]:
        """Generate structured output with Headroom optimization.

        Optimizes the prompt messages before delegating to the wrapped model's
        structured_output method.

        Args:
            output_model: The type/schema for the structured output
            prompt: List of prompt messages
            system_prompt: Optional system prompt
            **kwargs: Additional arguments passed to the wrapped model

        Yields:
            Structured output events from the wrapped model
        """
        # Run optimization in executor (CPU-bound)
        loop = asyncio.get_running_loop()
        optimized_prompt, metrics = await loop.run_in_executor(
            None, self._optimize_messages, prompt
        )

        logger.info(
            f"Headroom optimized (structured_output): {metrics.tokens_before} -> "
            f"{metrics.tokens_after} tokens ({metrics.savings_percent:.1f}% saved)"
        )

        # Delegate to wrapped model
        async for event in self.wrapped_model.structured_output(
            output_model, optimized_prompt, system_prompt=system_prompt, **kwargs
        ):
            yield event

    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,
                "total_tokens_before": 0,
                "total_tokens_after": 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

    # =========================================================================
    # Forward attribute access to wrapped model for compatibility
    # =========================================================================

    def __getattr__(self, name: str) -> Any:
        """Forward attribute access to wrapped model."""
        # Avoid infinite recursion for our own attributes
        if name in (
            "wrapped_model",
            "config",
            "auto_detect_provider",
            "_metrics_history",
            "_total_tokens_saved",
            "_pipeline",
            "_headroom_provider",
            "_lock",
            "pipeline",
            "total_tokens_saved",
            "metrics_history",
        ):
            raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")
        return getattr(self.wrapped_model, name)


def optimize_messages(
    messages: list[Any],
    config: HeadroomConfig | None = None,
    model: str = "gpt-4o",
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    """Standalone function to optimize Strands messages.

    Use this for manual optimization when you need fine-grained control.

    Args:
        messages: List of Strands messages (dicts)
        config: HeadroomConfig for optimization settings
        model: Model name for token estimation

    Returns:
        Tuple of (optimized_messages, metrics_dict)

    Example:
        from headroom.integrations.strands 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_strands_available()

    config = config or HeadroomConfig()
    provider = OpenAIProvider()
    pipeline = TransformPipeline(config=config, provider=provider)

    # Convert to OpenAI format (Strands uses similar format)
    openai_messages = []
    for msg in messages:
        if isinstance(msg, dict):
            entry: dict[str, Any] = {
                "role": msg.get("role", "user"),
                "content": msg.get("content", ""),
            }
            if "tool_calls" in msg and msg["tool_calls"]:
                entry["tool_calls"] = msg["tool_calls"]
            if "tool_call_id" in msg and msg["tool_call_id"]:
                entry["tool_call_id"] = msg["tool_call_id"]
            openai_messages.append(entry)
        elif hasattr(msg, "role") and hasattr(msg, "content"):
            entry = {"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)
        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