my-pal-mcp-server/providers/openrouter_registry.py

"""OpenRouter model registry for managing model configurations and aliases."""

import json
import logging
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

from .base import ModelCapabilities, ProviderType, RangeTemperatureConstraint


@dataclass
class OpenRouterModelConfig:
    """Configuration for an OpenRouter model."""

    model_name: str
    aliases: list[str] = field(default_factory=list)
    context_window: int = 32768  # Total context window size in tokens
    supports_extended_thinking: bool = False
    supports_system_prompts: bool = True
    supports_streaming: bool = True
    supports_function_calling: bool = False
    supports_json_mode: bool = False
    description: str = ""

    def to_capabilities(self) -> ModelCapabilities:
        """Convert to ModelCapabilities object."""
        return ModelCapabilities(
            provider=ProviderType.OPENROUTER,
            model_name=self.model_name,
            friendly_name="OpenRouter",
            max_tokens=self.context_window,  # ModelCapabilities still uses max_tokens
            supports_extended_thinking=self.supports_extended_thinking,
            supports_system_prompts=self.supports_system_prompts,
            supports_streaming=self.supports_streaming,
            supports_function_calling=self.supports_function_calling,
            temperature_constraint=RangeTemperatureConstraint(0.0, 2.0, 1.0),
        )


class OpenRouterModelRegistry:
    """Registry for managing OpenRouter model configurations and aliases."""

    def __init__(self, config_path: Optional[str] = None):
        """Initialize the registry.

        Args:
            config_path: Path to config file. If None, uses default locations.
        """
        self.alias_map: dict[str, str] = {}  # alias -> model_name
        self.model_map: dict[str, OpenRouterModelConfig] = {}  # model_name -> config

        # Determine config path
        if config_path:
            self.config_path = Path(config_path)
        else:
            # Check environment variable first
            env_path = os.getenv("OPENROUTER_MODELS_PATH")
            if env_path:
                self.config_path = Path(env_path)
            else:
                # Default to conf/openrouter_models.json
                self.config_path = Path(__file__).parent.parent / "conf" / "openrouter_models.json"

        # Load configuration
        self.reload()

    def reload(self) -> None:
        """Reload configuration from disk."""
        try:
            configs = self._read_config()
            self._build_maps(configs)
            logging.info(f"Loaded {len(self.model_map)} OpenRouter models with {len(self.alias_map)} aliases")
        except ValueError as e:
            # Re-raise ValueError only for duplicate aliases (critical config errors)
            logging.error(f"Failed to load OpenRouter model configuration: {e}")
            # Initialize with empty maps on failure
            self.alias_map = {}
            self.model_map = {}
            if "Duplicate alias" in str(e):
                raise
        except Exception as e:
            logging.error(f"Failed to load OpenRouter model configuration: {e}")
            # Initialize with empty maps on failure
            self.alias_map = {}
            self.model_map = {}

    def _read_config(self) -> list[OpenRouterModelConfig]:
        """Read configuration from file.

        Returns:
            List of model configurations
        """
        if not self.config_path.exists():
            logging.warning(f"OpenRouter model config not found at {self.config_path}")
            return []

        try:
            with open(self.config_path) as f:
                data = json.load(f)

            # Parse models
            configs = []
            for model_data in data.get("models", []):
                # Handle backwards compatibility - rename max_tokens to context_window
                if "max_tokens" in model_data and "context_window" not in model_data:
                    model_data["context_window"] = model_data.pop("max_tokens")

                config = OpenRouterModelConfig(**model_data)
                configs.append(config)

            return configs
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON in {self.config_path}: {e}")
        except Exception as e:
            raise ValueError(f"Error reading config from {self.config_path}: {e}")

    def _build_maps(self, configs: list[OpenRouterModelConfig]) -> None:
        """Build alias and model maps from configurations.

        Args:
            configs: List of model configurations
        """
        alias_map = {}
        model_map = {}

        for config in configs:
            # Add to model map
            model_map[config.model_name] = config

            # Add aliases
            for alias in config.aliases:
                alias_lower = alias.lower()
                if alias_lower in alias_map:
                    existing_model = alias_map[alias_lower]
                    raise ValueError(
                        f"Duplicate alias '{alias}' found for models " f"'{existing_model}' and '{config.model_name}'"
                    )
                alias_map[alias_lower] = config.model_name

        # Atomic update
        self.alias_map = alias_map
        self.model_map = model_map

    def resolve(self, name_or_alias: str) -> Optional[OpenRouterModelConfig]:
        """Resolve a model name or alias to configuration.

        Args:
            name_or_alias: Model name or alias to resolve

        Returns:
            Model configuration if found, None otherwise
        """
        # Try alias first (case-insensitive)
        alias_lower = name_or_alias.lower()
        if alias_lower in self.alias_map:
            model_name = self.alias_map[alias_lower]
            return self.model_map.get(model_name)

        # Try as direct model name
        return self.model_map.get(name_or_alias)

    def get_capabilities(self, name_or_alias: str) -> Optional[ModelCapabilities]:
        """Get model capabilities for a name or alias.

        Args:
            name_or_alias: Model name or alias

        Returns:
            ModelCapabilities if found, None otherwise
        """
        config = self.resolve(name_or_alias)
        if config:
            return config.to_capabilities()
        return None

    def list_models(self) -> list[str]:
        """List all available model names."""
        return list(self.model_map.keys())

    def list_aliases(self) -> list[str]:
        """List all available aliases."""
        return list(self.alias_map.keys())