src/model_config.py

#!/usr/bin/env python3
"""
Model Configuration for MCPMark
================================

This module provides configuration management for different LLM models,
automatically detecting the required API keys and base URLs based on the model name.
"""

import os
from typing import Dict, List

from src.logger import get_logger

# Initialize logger
logger = get_logger(__name__)


class ModelConfig:
    """
    Configuration container for a specific model.
    It loads the necessary API key and base URL from environment variables.
    """

    # Model configuration mapping
    MODEL_CONFIGS = {
        # OpenAI models
        "gpt-4o": {
            "provider": "openai",
            "api_key_var": "OPENAI_API_KEY",
            "litellm_input_model_name": "openai/gpt-4o",
        },
        "gpt-4.1": {
            "provider": "openai",
            "api_key_var": "OPENAI_API_KEY",
            "litellm_input_model_name": "openai/gpt-4.1",
        },
        "gpt-4.1-mini": {
            "provider": "openai",
            "api_key_var": "OPENAI_API_KEY",
            "litellm_input_model_name": "openai/gpt-4.1-mini",
        },
        "gpt-4.1-nano": {
            "provider": "openai",
            "api_key_var": "OPENAI_API_KEY",
            "litellm_input_model_name": "openai/gpt-4.1-nano",
        },
        "gpt-5.2": {
            "provider": "openai",
            "api_key_var": "OPENAI_API_KEY",
            "litellm_input_model_name": "openai/gpt-5.2",
        },
        "gpt-5": {
            "provider": "openai",
            "api_key_var": "OPENAI_API_KEY",
            "litellm_input_model_name": "openai/gpt-5",
        },
        "gpt-5-mini": {
            "provider": "openai",
            "api_key_var": "OPENAI_API_KEY",
            "litellm_input_model_name": "openai/gpt-5-mini",
        },
        "gpt-5-nano": {
            "provider": "openai",
            "api_key_var": "OPENAI_API_KEY",
            "litellm_input_model_name": "openai/gpt-5-nano",
        },
        "o3": {
            "provider": "openai",
            "api_key_var": "OPENAI_API_KEY",
            "litellm_input_model_name": "openai/o3",
        },
        "o4-mini": {
            "provider": "openai",
            "api_key_var": "OPENAI_API_KEY",
            "litellm_input_model_name": "openai/o4-mini",
        },
        "gpt-oss-120b": {
            "provider": "openai",
            "api_key_var": "OPENROUTER_API_KEY",
            "litellm_input_model_name": "openrouter/openai/gpt-oss-120b",
        },
        # DeepSeek models (dimcode OpenAI-compatible endpoint)
        "deepseek-v4-pro": {
            "provider": "openai",
            "api_key_var": "DEEPSEEK_API_KEY",
            "base_url_var": "DEEPSEEK_BASE_URL",
            "litellm_input_model_name": "openai/deepseek-v4-pro",
            # litellm's openai provider rejects a TOP-LEVEL reasoning_effort, so
            # both the thinking toggle and reasoning_effort ride in extra_body
            # (forwarded verbatim into the request JSON the dimcode endpoint reads).
            "extra_body": {"thinking": {"type": "enabled"}, "reasoning_effort": "high"},
        },
        # Anthropic models
        "claude-3.7-sonnet": {
            "provider": "anthropic",
            "api_key_var": "ANTHROPIC_API_KEY",
            "litellm_input_model_name": "anthropic/claude-3-7-sonnet-20250219",
        },
        "claude-sonnet-4": {
            "provider": "anthropic",
            "api_key_var": "ANTHROPIC_API_KEY",
            "litellm_input_model_name": "anthropic/claude-sonnet-4-20250514",
        },
        "claude-sonnet-4.5": {
            "provider": "anthropic",
            "api_key_var": "ANTHROPIC_API_KEY",
            "litellm_input_model_name": "anthropic/claude-sonnet-4-5-20250929",
        },
        "claude-opus-4": {
            "provider": "anthropic",
            "api_key_var": "ANTHROPIC_API_KEY",
            "litellm_input_model_name": "anthropic/claude-opus-4-20250514",
        },
        "claude-opus-4.1": {
            "provider": "anthropic",
            "api_key_var": "ANTHROPIC_API_KEY",
            "litellm_input_model_name": "anthropic/claude-opus-4-1-20250805",
        },
        "claude-opus-4.5": {
            "provider": "anthropic",
            "api_key_var": "ANTHROPIC_API_KEY",
            "litellm_input_model_name": "anthropic/claude-opus-4-5-20251101",
        },
        # Google models
        "gemini-2.5-pro": {
            "provider": "google",
            "api_key_var": "GEMINI_API_KEY",
            "litellm_input_model_name": "gemini/gemini-2.5-pro",
        },
        "gemini-2.5-flash": {
            "provider": "google",
            "api_key_var": "GEMINI_API_KEY",
            "litellm_input_model_name": "gemini/gemini-2.5-flash",
        },
        "gemini-3-pro": {
            "provider": "google",
            "api_key_var": "GEMINI_API_KEY",
            "litellm_input_model_name": "gemini/gemini-3-pro-preview",
        },
        # Moonshot models
        "kimi-k2-0711": {
            "provider": "moonshot",
            "api_key_var": "MOONSHOT_API_KEY",
            "litellm_input_model_name": "moonshot/kimi-k2-0711-preview",
        },
        "kimi-k2-0905": {
            "provider": "moonshot",
            "api_key_var": "MOONSHOT_API_KEY",
            "litellm_input_model_name": "moonshot/kimi-k2-0905-preview",
        },
        "kimi-k2-thinking": {
            "provider": "moonshot",
            "api_key_var": "OPENROUTER_API_KEY",
            "litellm_input_model_name": "openrouter/moonshotai/kimi-k2-thinking",
        },
        # Grok models
        "grok-4": {
            "provider": "xai",
            "api_key_var": "GROK_API_KEY",
            "litellm_input_model_name": "xai/grok-4-0709",
        },
        "grok-code-fast-1": {
            "provider": "xai",
            "api_key_var": "GROK_API_KEY",
            "litellm_input_model_name": "xai/grok-code-fast-1",
        },
        # Qwen models
        "qwen-3-coder-plus": {
            "provider": "qwen",
            "api_key_var": "DASHSCOPE_API_KEY",
            "litellm_input_model_name": "dashscope/qwen3-coder-plus",
        },
        "qwen-3-max": {
            "provider": "qwen",
            "api_key_var": "DASHSCOPE_API_KEY",
            "litellm_input_model_name": "dashscope/qwen3-max-preview",
        },
        # Zhipu
        "glm-4.5": {
            "provider": "zhipu",
            "api_key_var": "OPENROUTER_API_KEY",
            "litellm_input_model_name": "openrouter/z-ai/glm-4.5",
        }
    }

    def __init__(self, model_name: str):
        """
        Initializes the model configuration.

        Args:
            model_name: The name of the model (e.g., 'gpt-4o', 'deepseek-chat').

        Raises:
            ValueError: If the model is not supported or environment variables are missing.
        """
        self.short_model_name = model_name
        model_info = self._get_model_info(model_name)

        # Load API key, base URL and LiteLLM model name from environment variables
        if "base_url_var" in model_info:
            self.base_url = os.getenv(model_info["base_url_var"])
        else:
            self.base_url = None
        
        self.api_key = os.getenv(model_info["api_key_var"])
        if not self.api_key:
            raise ValueError(
                f"Missing required environment variable: {model_info['api_key_var']}"
            )

        self.litellm_input_model_name = model_info.get("litellm_input_model_name", model_name)

        # Optional extra request-body fields forwarded to litellm (e.g. DeepSeek
        # `thinking` toggle). None when the model has no extras.
        self.extra_body = model_info.get("extra_body")

    def _get_model_info(self, model_name: str) -> Dict[str, str]:
        """
        Retrieves the configuration details for a given model name.
        For unsupported models, defaults to using OPENAI_BASE_URL and OPENAI_API_KEY.
        """
        if model_name not in self.MODEL_CONFIGS:
            logger.warning(
                f"Model '{model_name}' not in supported list. Using default OpenAI configuration."
            )
            # Return default configuration for unsupported models
            return {
                "provider": "openai",
                "api_key_var": "OPENAI_API_KEY",
                "litellm_input_model_name": model_name,
            }
        return self.MODEL_CONFIGS[model_name]

    @classmethod
    def get_supported_models(cls) -> List[str]:
        """Returns a list of all supported model names."""
        return list(cls.MODEL_CONFIGS.keys())


def main():
    """Example usage of the ModelConfig class."""
    logger.info("Supported models: %s", ModelConfig.get_supported_models())

    try:
        # Example: Create a model config for DeepSeek
        model_config = ModelConfig("deepseek-chat")
        logger.info("✅ DeepSeek model config created successfully.")
        logger.info("Short model name: %s", model_config.short_model_name)
        logger.info("API key loaded: %s", bool(model_config.api_key))

    except ValueError as e:
        logger.error("⚠️  Configuration error: %s", e)


if __name__ == "__main__":
    main()