backend/models/registry.py

from typing import Dict, List, Optional, Set
from .models import Model, ModelProvider, ModelCapability, ModelPricing

DEFAULT_FREE_MODEL = "Kimi K2"
DEFAULT_PREMIUM_MODEL = "Claude Sonnet 4"

class ModelRegistry:
    def __init__(self):
        self._models: Dict[str, Model] = {}
        self._aliases: Dict[str, str] = {}
        self._initialize_models()
    
    def _initialize_models(self):
        self.register(Model(
            id="anthropic/claude-sonnet-4-20250514",
            name="Claude Sonnet 4",
            provider=ModelProvider.ANTHROPIC,
            aliases=["claude-sonnet-4", "anthropic/claude-sonnet-4", "Claude Sonnet 4", "claude-sonnet-4-20250514"],
            context_window=200_000,
            capabilities=[
                ModelCapability.CHAT,
                ModelCapability.FUNCTION_CALLING,
                ModelCapability.VISION,
                ModelCapability.THINKING,
            ],
            pricing=ModelPricing(
                input_cost_per_million_tokens=3.00,
                output_cost_per_million_tokens=15.00
            ),
            tier_availability=["paid"],
            priority=100,
            recommended=True,
            enabled=True
        ))
        
        self.register(Model(
            id="anthropic/claude-3-7-sonnet-latest",
            name="Claude 3.7 Sonnet",
            provider=ModelProvider.ANTHROPIC,
            aliases=["sonnet-3.7", "claude-3.7", "Claude 3.7 Sonnet", "claude-3-7-sonnet-latest"],
            context_window=200_000,
            capabilities=[
                ModelCapability.CHAT,
                ModelCapability.FUNCTION_CALLING,
                ModelCapability.VISION,
            ],
            pricing=ModelPricing(
                input_cost_per_million_tokens=3.00,
                output_cost_per_million_tokens=15.00
            ),
            tier_availability=["paid"],
            priority=93,
            enabled=True
        ))
        
        self.register(Model(
            id="anthropic/claude-3-5-sonnet-latest",
            name="Claude 3.5 Sonnet",
            provider=ModelProvider.ANTHROPIC,
            aliases=["sonnet-3.5", "claude-3.5", "Claude 3.5 Sonnet", "claude-3-5-sonnet-latest"],
            context_window=200_000,
            capabilities=[
                ModelCapability.CHAT,
                ModelCapability.FUNCTION_CALLING,
                ModelCapability.VISION,
            ],
            pricing=ModelPricing(
                input_cost_per_million_tokens=3.00,
                output_cost_per_million_tokens=15.00
            ),
            tier_availability=["paid"],
            priority=90,
            enabled=True
        ))
        
        self.register(Model(
            id="openai/gpt-5",
            name="GPT-5",
            provider=ModelProvider.OPENAI,
            aliases=["gpt-5", "GPT-5"],
            context_window=400_000,
            capabilities=[
                ModelCapability.CHAT,
                ModelCapability.FUNCTION_CALLING,
                ModelCapability.VISION,
                ModelCapability.STRUCTURED_OUTPUT,
            ],
            pricing=ModelPricing(
                input_cost_per_million_tokens=1.25,
                output_cost_per_million_tokens=10.00
            ),
            tier_availability=["paid"],
            priority=99,
            enabled=True
        ))
        
        self.register(Model(
            id="openai/gpt-5-mini",
            name="GPT-5 Mini",
            provider=ModelProvider.OPENAI,
            aliases=["gpt-5-mini", "GPT-5 Mini"],
            context_window=400_000,
            capabilities=[
                ModelCapability.CHAT,
                ModelCapability.FUNCTION_CALLING,
                ModelCapability.STRUCTURED_OUTPUT,
            ],
            pricing=ModelPricing(
                input_cost_per_million_tokens=0.25,
                output_cost_per_million_tokens=2.00
            ),
            tier_availability=["free", "paid"],
            priority=85,
            enabled=True
        ))
        
        self.register(Model(
            id="gemini/gemini-2.5-pro",
            name="Gemini 2.5 Pro",
            provider=ModelProvider.GOOGLE,
            aliases=["google/gemini-2.5-pro", "gemini-2.5-pro", "Gemini 2.5 Pro"],
            context_window=2_000_000,
            capabilities=[
                ModelCapability.CHAT,
                ModelCapability.FUNCTION_CALLING,
                ModelCapability.VISION,
                ModelCapability.STRUCTURED_OUTPUT,
            ],
            pricing=ModelPricing(
                input_cost_per_million_tokens=1.25,
                output_cost_per_million_tokens=10.00
            ),
            tier_availability=["paid"],
            priority=96,
            enabled=True
        ))
        
        self.register(Model(
            id="xai/grok-4",
            name="Grok 4",
            provider=ModelProvider.XAI,
            aliases=["grok-4", "x-ai/grok-4", "openrouter/x-ai/grok-4", "Grok 4"],
            context_window=128_000,
            capabilities=[
                ModelCapability.CHAT,
                ModelCapability.FUNCTION_CALLING,
            ],
            pricing=ModelPricing(
                input_cost_per_million_tokens=5.00,
                output_cost_per_million_tokens=15.00
            ),
            tier_availability=["paid"],
            priority=94,
            enabled=True
        ))
        
        self.register(Model(
            id="openrouter/moonshotai/kimi-k2",
            name="Kimi K2",
            provider=ModelProvider.MOONSHOTAI,
            aliases=["moonshotai/kimi-k2", "kimi-k2", "Kimi K2"],
            context_window=200_000,
            capabilities=[
                ModelCapability.CHAT,
                ModelCapability.FUNCTION_CALLING,
            ],
            pricing=ModelPricing(
                input_cost_per_million_tokens=1.00,
                output_cost_per_million_tokens=3.00
            ),
            tier_availability=["free", "paid"],
            priority=100,
            enabled=True
        ))

        """
        # DeepSeek Models
        self.register(Model(
            id="openrouter/deepseek/deepseek-chat",
            name="DeepSeek Chat",
            provider=ModelProvider.OPENROUTER,
            aliases=["deepseek", "deepseek-chat"],
            context_window=128_000,
            capabilities=[
                ModelCapability.CHAT, 
                ModelCapability.FUNCTION_CALLING
            ],
            pricing=ModelPricing(
                input_cost_per_million_tokens=0.38,
                output_cost_per_million_tokens=0.89
            ),
            tier_availability=["free", "paid"],
            priority=95,
            enabled=False  # Currently disabled
        ))
        
        # Qwen Models
        self.register(Model(
            id="openrouter/qwen/qwen3-235b-a22b",
            name="Qwen3 235B",
            provider=ModelProvider.OPENROUTER,
            aliases=["qwen3", "qwen-3"],
            context_window=128_000,
            capabilities=[
                ModelCapability.CHAT, 
                ModelCapability.FUNCTION_CALLING
            ],
            pricing=ModelPricing(
                input_cost_per_million_tokens=0.13,
                output_cost_per_million_tokens=0.60
            ),
            tier_availability=["free", "paid"],
            priority=90,
            enabled=False  # Currently disabled
        ))
        """
    
    def register(self, model: Model) -> None:
        self._models[model.id] = model
        for alias in model.aliases:
            self._aliases[alias] = model.id
    
    def get(self, model_id: str) -> Optional[Model]:
        if model_id in self._models:
            return self._models[model_id]
        
        if model_id in self._aliases:
            actual_id = self._aliases[model_id]
            return self._models.get(actual_id)
        
        return None
    
    def get_all(self, enabled_only: bool = True) -> List[Model]:
        models = list(self._models.values())
        if enabled_only:
            models = [m for m in models if m.enabled]
        return models
    
    def get_by_tier(self, tier: str, enabled_only: bool = True) -> List[Model]:
        models = self.get_all(enabled_only)
        return [m for m in models if tier in m.tier_availability]
    
    def get_by_provider(self, provider: ModelProvider, enabled_only: bool = True) -> List[Model]:
        models = self.get_all(enabled_only)
        return [m for m in models if m.provider == provider]
    
    def get_by_capability(self, capability: ModelCapability, enabled_only: bool = True) -> List[Model]:
        models = self.get_all(enabled_only)
        return [m for m in models if capability in m.capabilities]
    
    def resolve_model_id(self, model_id: str) -> Optional[str]:
        model = self.get(model_id)
        return model.id if model else None
    
    def get_aliases(self, model_id: str) -> List[str]:
        model = self.get(model_id)
        return model.aliases if model else []
    
    def enable_model(self, model_id: str) -> bool:
        model = self.get(model_id)
        if model:
            model.enabled = True
            return True
        return False
    
    def disable_model(self, model_id: str) -> bool:
        model = self.get(model_id)
        if model:
            model.enabled = False
            return True
        return False
    
    def get_context_window(self, model_id: str, default: int = 31_000) -> int:
        model = self.get(model_id)
        return model.context_window if model else default
    
    def get_pricing(self, model_id: str) -> Optional[ModelPricing]:
        model = self.get(model_id)
        return model.pricing if model else None
    
    def to_legacy_format(self) -> Dict:
        models_dict = {}
        aliases_dict = {}
        pricing_dict = {}
        context_windows_dict = {}
        
        for model in self.get_all(enabled_only=True):
            models_dict[model.id] = {
                "aliases": model.aliases,
                "pricing": {
                    "input_cost_per_million_tokens": model.pricing.input_cost_per_million_tokens,
                    "output_cost_per_million_tokens": model.pricing.output_cost_per_million_tokens,
                } if model.pricing else None,
                "context_window": model.context_window,
                "tier_availability": model.tier_availability,
            }
            
            for alias in model.aliases:
                aliases_dict[alias] = model.id
            
            if model.pricing:
                pricing_dict[model.id] = {
                    "input_cost_per_million_tokens": model.pricing.input_cost_per_million_tokens,
                    "output_cost_per_million_tokens": model.pricing.output_cost_per_million_tokens,
                }
            
            context_windows_dict[model.id] = model.context_window
        
        free_models = [m.id for m in self.get_by_tier("free")]
        paid_models = [m.id for m in self.get_by_tier("paid")]
        
        # Debug logging
        from utils.logger import logger
        logger.debug(f"Legacy format generation: {len(free_models)} free models, {len(paid_models)} paid models")
        logger.debug(f"Free models: {free_models}")
        logger.debug(f"Paid models: {paid_models}")
        
        return {
            "MODELS": models_dict,
            "MODEL_NAME_ALIASES": aliases_dict,
            "HARDCODED_MODEL_PRICES": pricing_dict,
            "MODEL_CONTEXT_WINDOWS": context_windows_dict,
            "FREE_TIER_MODELS": free_models,
            "PAID_TIER_MODELS": paid_models,
        }

registry = ModelRegistry()