src/model_registry.py

from dataclasses import dataclass
from typing import Dict


@dataclass(frozen=True)
class ModelConfig:
    alias: str
    deployment_name: str
    is_openai_compatible: bool = True
    backend: str = "openai"  # "openai", "google", "huggingface", or "vllm"
    is_thinking_model: bool = False  # Uses thinking tokens, needs higher max_completion_tokens
    vllm_base_url: str = "http://localhost:8000/v1"  # Only used when backend="vllm"


MODELS: Dict[str, ModelConfig] = {
    "gemini-3-flash": ModelConfig(
        alias="gemini-3-flash",
        deployment_name="gemini-3-flash-preview",
        is_openai_compatible=True,
        backend="google",
        is_thinking_model=True,
    ),
    "gpt-5.1": ModelConfig(
        alias="gpt-5.1",
        deployment_name="gpt-5.1",
        is_openai_compatible=True,
        backend="openai",
    ),
    "gemma-3-27b-it": ModelConfig(
        alias="gemma-3-27b-it",
        deployment_name="google/gemma-3-27b-it",
        is_openai_compatible=False,
        backend="vllm",
    ),
}


DEFAULT_MODEL_ALIAS = "gemini-3-flash"


def get_model_config(alias: str) -> ModelConfig:
    if alias not in MODELS:
        available = ", ".join(sorted(MODELS.keys()))
        raise ValueError(
            f"Unknown model alias '{alias}'. Available models: {available}"
        )
    return MODELS[alias]