core/model_config.py

"""Model configuration helpers and preset models."""

from typing import Any

# Preset provider models for quick selection
PROVIDER_PRESETS = {
    "GPT-5.2 Pro": {
        "model": "gpt-5.2-pro-2025-12-11",
        "tokenizer": "openai/gpt-5.2-pro-2025-12-11",
        "max_context": 272000,
        "api_key_env": "OPENAI_API_KEY",
    },
    "GPT-5.2": {
        "model": "gpt-5.2-2025-12-11",
        "tokenizer": "openai/gpt-5.2-2025-12-11",
        "max_context": 272000,
        "api_key_env": "OPENAI_API_KEY",
    },
    "GPT-5.1": {
        "model": "gpt-5.1-2025-11-13",
        "tokenizer": "openai/gpt-5.1-2025-11-13",
        "max_context": 272000,
        "api_key_env": "OPENAI_API_KEY",
    },
    "GPT-5": {
        "model": "gpt-5-2025-08-07",
        "tokenizer": "openai/gpt-5-2025-08-07",
        "max_context": 272000,
        "api_key_env": "OPENAI_API_KEY",
    },
    "GPT-5 Mini": {
        "model": "gpt-5-mini-2025-08-07",
        "tokenizer": "openai/gpt-5-mini-2025-08-07",
        "max_context": 272000,
        "api_key_env": "OPENAI_API_KEY",
    },
    "GPT-5 Nano": {
        "model": "gpt-5-nano-2025-08-07",
        "tokenizer": "openai/gpt-5-nano-2025-08-07",
        "max_context": 272000,
        "api_key_env": "OPENAI_API_KEY",
    },
    "Claude Opus 4.5": {
        "model": "claude-opus-4-5-202511019",
        "tokenizer": "anthropic/claude-opus-4-5-20251101",
        "max_context": 200000,
        "api_key_env": "ANTHROPIC_API_KEY",
    },
    "Claude Sonnet 4.5": {
        "model": "claude-sonnet-4-5-20250929",
        "tokenizer": "anthropic/claude-sonnet-4-5-20250929",
        "max_context": 200000,
        "api_key_env": "ANTHROPIC_API_KEY",
    },
    "Claude Haiku 4.5": {
        "model": "claude-haiku-4-5-20251001",
        "tokenizer": "anthropic/claude-haiku-4-5-20251001",
        "max_context": 200000,
        "api_key_env": "ANTHROPIC_API_KEY",
    },
    "Gemini 3.0 Pro": {
        "model": "gemini/gemini-3-pro-preview",
        "tokenizer": "gemini/gemini-3-pro-preview",
        "max_context": 2000000,
        "api_key_env": "GEMINI_API_KEY",
    },
    "Gemini 3.0 Flash": {
        "model": "gemini/gemini-3-flash-preview",
        "tokenizer": "gemini/gemini-3-flash-preview",
        "max_context": 1000000,
        "api_key_env": "GEMINI_API_KEY",
    },
    "Gemini 2.5 Pro": {
        "model": "gemini/gemini-2.5-pro",
        "tokenizer": "gemini/gemini-2.5-pro",
        "max_context": 2000000,
        "api_key_env": "GEMINI_API_KEY",
    },
    "Gemini 2.5 Flash": {
        "model": "gemini/gemini-2.5-flash",
        "tokenizer": "gemini/gemini-2.5-flash",
        "max_context": 1000000,
        "api_key_env": "GEMINI_API_KEY",
    },
}


def create_local_model_config(
    model: str,
    api_base: str | None = None,
    max_context: int = 131072,
    tokenizer: str | None = None,
) -> dict[str, Any]:
    """
    Create a local model configuration.
    
    Args:
        model: Model name (e.g., "ollama/llama2" or "gpt-3.5-turbo" for vLLM)
        api_base: API base URL (defaults based on model type)
        max_context: Maximum context window size
        tokenizer: Tokenizer name for token counting
        
    Returns:
        Model configuration dictionary
    """
    # Set default API base based on model type
    if api_base is None:
        if model.startswith("ollama/"):
            api_base = "http://localhost:11434"
        elif model.startswith("vllm/") or not model.startswith(("ollama/", "openrouter/")):
            # Assume OpenAI-compatible (vLLM)
            api_base = "http://localhost:8000/v1"
    
    # Infer tokenizer if not provided
    if tokenizer is None:
        if model.startswith("ollama/"):
            # Try to infer from model name
            model_name = model.replace("ollama/", "")
            tokenizer = f"hf/{model_name}"
        else:
            # For vLLM/OpenAI-compatible, try to infer
            tokenizer = model.replace("vllm/", "")
    
    return {
        "type": "local",
        "model": model,
        "api_base": api_base,
        "max_context": max_context,
        "tokenizer": tokenizer,
    }


def create_provider_model_config(
    model: str,
    api_key: str,
    max_context: int = 128000,
    tokenizer: str | None = None,
) -> dict[str, Any]:
    """
    Create a provider model configuration.
    
    Args:
        model: Model name in litellm format
        api_key: API key for the provider
        max_context: Maximum context window size
        tokenizer: Tokenizer name for token counting
        
    Returns:
        Model configuration dictionary
    """
    # Infer tokenizer if not provided
    if tokenizer is None:
        # Try to infer from model name
        if model.startswith("openai/") or not "/" in model:
            # OpenAI models
            model_name = model.replace("openai/", "")
            tokenizer = f"openai/{model_name}"
        elif model.startswith("anthropic/") or model.startswith("claude-"):
            # Anthropic models
            model_name = model.replace("anthropic/", "")
            tokenizer = f"anthropic/{model_name}"
        elif model.startswith("gemini/"):
            # Gemini models
            model_name = model.replace("gemini/", "")
            tokenizer = f"gemini/{model_name}"
        else:
            # Generic fallback
            tokenizer = "gpt2"
    
    return {
        "type": "provider",
        "model": model,
        "api_key": api_key,
        "max_context": max_context,
        "tokenizer": tokenizer,
    }


def get_provider_from_model(model: str) -> str:
    """
    Infer provider name from model identifier.
    
    Args:
        model: Model name in litellm format
        
    Returns:
        Provider name hint (e.g., "openai", "anthropic", "gemini")
    """
    model_lower = model.lower()
    if model_lower.startswith("openai/") or not "/" in model:
        return "openai"
    elif model_lower.startswith("anthropic/") or model_lower.startswith("claude-"):
        return "anthropic"
    elif model_lower.startswith("gemini/"):
        return "gemini"
    elif model_lower.startswith("openrouter/"):
        return "openrouter"
    elif model_lower.startswith("cohere/"):
        return "cohere"
    elif model_lower.startswith("mistral/"):
        return "mistral"
    else:
        return "other"


def get_api_key_env_name(provider: str) -> str:
    """
    Get the environment variable name for API key based on provider.
    
    Args:
        provider: Provider name
        
    Returns:
        Environment variable name for API key
    """
    provider_to_key = {
        "openai": "OPENAI_API_KEY",
        "anthropic": "ANTHROPIC_API_KEY",
        "gemini": "GEMINI_API_KEY",
        "openrouter": "OPENROUTER_API_KEY",
        "cohere": "COHERE_API_KEY",
        "mistral": "MISTRAL_API_KEY",
        "other": "API_KEY",
    }
    return provider_to_key.get(provider.lower(), "API_KEY")