headroom/integrations/strands/providers.py

"""Provider detection for Strands models.

Automatically detects the correct Headroom provider based on the Strands model type.
"""

from __future__ import annotations

import logging
from typing import Any

from headroom.providers import (
    AnthropicProvider,
    GoogleProvider,
    OpenAIProvider,
)
from headroom.providers.base import Provider

logger = logging.getLogger(__name__)

# Mapping from Strands model class names to Headroom providers
_STRANDS_MODEL_PROVIDERS: dict[str, type[Provider]] = {
    # Bedrock models (primarily Claude via Bedrock)
    "BedrockModel": AnthropicProvider,
    # Anthropic models (direct API)
    "AnthropicModel": AnthropicProvider,
    # OpenAI models
    "OpenAIModel": OpenAIProvider,
    # LiteLLM (uses OpenAI-compatible interface)
    "LiteLLMModel": OpenAIProvider,
    # Ollama (uses OpenAI-compatible interface)
    "OllamaModel": OpenAIProvider,
    # Google Gemini models
    "GeminiModel": GoogleProvider,
    # Writer models (uses OpenAI-compatible interface)
    "WriterModel": OpenAIProvider,
}


def get_headroom_provider(model: Any) -> Provider:
    """Get the appropriate Headroom provider for a Strands model.

    Detection strategy:
    1. Check model class name against known Strands model types
    2. Check for provider hints in model attributes
    3. Fall back to OpenAI provider (most compatible)

    Args:
        model: A Strands model instance (BedrockModel, AnthropicModel, etc.)

    Returns:
        Appropriate Headroom Provider instance.

    Example:
        from strands.models import BedrockModel
        from headroom.integrations.strands.providers import get_headroom_provider

        model = BedrockModel(model_id="anthropic.claude-3-5-sonnet-20241022-v2:0")
        provider = get_headroom_provider(model)  # Returns AnthropicProvider
    """
    # Strategy 1: Class name matching
    class_name = model.__class__.__name__
    if class_name in _STRANDS_MODEL_PROVIDERS:
        provider_class = _STRANDS_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 = model.__class__.__module__
    if "anthropic" in module_path.lower():
        logger.debug(f"Detected AnthropicProvider from module {module_path}")
        return AnthropicProvider()
    elif "bedrock" 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 "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 = _extract_model_id(model)
    if model_id:
        model_id_lower = model_id.lower()
        if "claude" in model_id_lower or "anthropic" 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()

    # Strategy 4: Default fallback
    logger.warning(
        f"Unknown Strands model class '{class_name}', defaulting to OpenAIProvider. "
        "Token counting may be inaccurate."
    )
    return OpenAIProvider()


def _extract_model_id(model: Any) -> str:
    """Extract model ID from a Strands model using various attribute names.

    Args:
        model: A Strands model instance

    Returns:
        Model ID string or empty string if not found
    """
    # Try common attribute names used by Strands models
    for attr in ["model_id", "model", "model_name", "id"]:
        value = getattr(model, attr, None)
        if value and isinstance(value, str):
            return str(value)

    # Try to get from config if available (config can be dict or object)
    config = getattr(model, "config", None)
    if config:
        for attr in ["model_id", "model", "model_name"]:
            # Handle dict-style config (Strands uses this)
            if isinstance(config, dict):
                value = config.get(attr)
            else:
                value = getattr(config, attr, None)
            if value and isinstance(value, str):
                return str(value)

    # Try get_config() method (Strands Model interface)
    if hasattr(model, "get_config"):
        try:
            config_dict = model.get_config()
            if isinstance(config_dict, dict):
                for attr in ["model_id", "model", "model_name"]:
                    value = config_dict.get(attr)
                    if value and isinstance(value, str):
                        return str(value)
        except Exception:
            pass

    return ""


def get_model_name_from_strands(model: Any) -> str:
    """Extract the model name/ID from a Strands model.

    Args:
        model: A Strands model instance

    Returns:
        Model name string (e.g., "anthropic.claude-3-5-sonnet-20241022-v2:0")
    """
    model_id = _extract_model_id(model)
    if model_id:
        return str(model_id)

    # Fallback with warning
    class_name = model.__class__.__name__
    logger.warning(
        f"Could not extract model name from {class_name} (no 'model_id', 'model', "
        f"'model_name', or 'id' attribute). Defaulting to 'gpt-4o'. "
        "Token counting may be inaccurate."
    )
    return "gpt-4o"