insights/services/llm-gateway/tests/test_model_registry.py

import os
from unittest.mock import MagicMock, patch

import pytest

from llm_gateway.cloudflare import CLOUDFLARE_ALLOWED_MODELS
from llm_gateway.rate_limiting.cost_refresh import COST_ALIASES
from llm_gateway.rate_limiting.model_cost_service import ModelCost, ModelCostService
from llm_gateway.services.model_registry import (
    ModelInfo,
    ModelRegistryService,
    _model_matches_allowlist,
    get_available_models,
    is_model_available,
)

PROVIDER_ENV_VARS = [
    "OPENAI_API_KEY",
    "ANTHROPIC_API_KEY",
    "OPENROUTER_API_KEY",
    "FIREWORKS_API_KEY",
]

MOCK_COST_DATA: dict[str, ModelCost] = {
    "gpt-4o": {
        "litellm_provider": "openai",
        "max_input_tokens": 128000,
        "supports_vision": True,
        "mode": "chat",
    },
    "gpt-4o-mini": {
        "litellm_provider": "openai",
        "max_input_tokens": 128000,
        "supports_vision": True,
        "mode": "chat",
    },
    "gpt-5.2": {
        "litellm_provider": "openai",
        "max_input_tokens": 200000,
        "supports_vision": True,
        "mode": "chat",
    },
    "gpt-5-mini": {
        "litellm_provider": "openai",
        "max_input_tokens": 200000,
        "supports_vision": True,
        "mode": "chat",
    },
    "gpt-5.3-codex": {
        "litellm_provider": "openai",
        "max_input_tokens": 200000,
        "supports_vision": False,
        "mode": "chat",
    },
    "o1": {
        "litellm_provider": "openai",
        "max_input_tokens": 200000,
        "supports_vision": True,
        "mode": "chat",
    },
    "claude-3-5-sonnet-20241022": {
        "litellm_provider": "anthropic",
        "max_input_tokens": 200000,
        "supports_vision": True,
        "mode": "chat",
    },
    "claude-haiku-4-5-20251001": {
        "litellm_provider": "anthropic",
        "max_input_tokens": 200000,
        "supports_vision": True,
        "mode": "chat",
    },
    "claude-opus-4-5": {
        "litellm_provider": "anthropic",
        "max_input_tokens": 200000,
        "supports_vision": True,
        "mode": "chat",
    },
    "claude-opus-4-6": {
        "litellm_provider": "anthropic",
        "max_input_tokens": 200000,
        "supports_vision": True,
        "mode": "chat",
    },
    "claude-sonnet-4-5": {
        "litellm_provider": "anthropic",
        "max_input_tokens": 200000,
        "supports_vision": True,
        "mode": "chat",
    },
    "claude-sonnet-4-6": {
        "litellm_provider": "anthropic",
        "max_input_tokens": 200000,
        "supports_vision": True,
        "mode": "chat",
    },
    "claude-haiku-4-5": {
        "litellm_provider": "anthropic",
        "max_input_tokens": 200000,
        "supports_vision": True,
        "mode": "chat",
    },
    "openrouter/anthropic/claude-3.5-sonnet": {
        "litellm_provider": "openrouter",
        "max_input_tokens": 200000,
        "supports_vision": True,
        "mode": "chat",
    },
    "fireworks_ai/accounts/fireworks/models/llama-v3p1-70b-instruct": {
        "litellm_provider": "fireworks_ai",
        "max_input_tokens": 131072,
        "supports_vision": False,
        "mode": "chat",
    },
    "text-embedding-ada-002": {
        "litellm_provider": "openai",
        "max_input_tokens": 8191,
        "supports_vision": False,
        "mode": "embedding",
    },
}


def mock_get_costs(self: ModelCostService, model: str) -> ModelCost | None:
    return MOCK_COST_DATA.get(model)


def mock_get_all_models(self: ModelCostService) -> dict[str, ModelCost]:
    return MOCK_COST_DATA


def create_mock_settings(
    openai: bool = True,
    anthropic: bool = True,
    openrouter: bool = False,
    fireworks: bool = False,
    cloudflare: bool = False,
    modal: bool = False,
    baseten: bool = False,
) -> MagicMock:
    settings = MagicMock()
    settings.openai_api_key = "sk-test" if openai else None
    settings.anthropic_api_key = "sk-ant-test" if anthropic else None
    settings.openrouter_api_key = "or-test" if openrouter else None
    settings.fireworks_api_key = "fw-test" if fireworks else None
    # CF needs both a key and an account id; default off so a MagicMock's truthy
    # auto-attributes don't silently enable Cloudflare model advertising.
    settings.cloudflare_api_key = "cf-test" if cloudflare else None
    settings.cloudflare_account_id = "acct-test" if cloudflare else None
    # Modal needs all three; same MagicMock-truthiness hazard as CF above.
    settings.modal_api_base = "https://modal.test/v1" if modal else None
    settings.modal_kimi_api_base = "https://kimi.modal.test/v1" if modal else None
    settings.modal_key = "wk-test" if modal else None
    settings.modal_secret = "ws-test" if modal else None
    settings.baseten_api_base = "https://baseten.test/v1" if baseten else None
    settings.baseten_api_key = "baseten-test" if baseten else None
    return settings


@pytest.fixture(autouse=True)
def reset_services():
    ModelRegistryService.reset_instance()
    ModelCostService.reset_instance()
    yield
    ModelRegistryService.reset_instance()
    ModelCostService.reset_instance()


@pytest.fixture(autouse=True)
def mock_cost_service():
    with (
        patch.object(ModelCostService, "get_costs", mock_get_costs),
        patch.object(ModelCostService, "get_all_models", mock_get_all_models),
    ):
        yield


@pytest.fixture(autouse=True)
def clear_env_api_keys():
    """Prevent real API keys in CI from leaking into unit tests.

    _get_configured_providers() checks both settings and env vars, so we need
    to clear the env vars to ensure unit tests only reflect mock settings.
    """
    with patch.dict(os.environ, {}, clear=False):
        for var in PROVIDER_ENV_VARS:
            os.environ.pop(var, None)
        yield


@pytest.fixture(autouse=True)
def mock_settings():
    with patch(
        "llm_gateway.services.model_registry.get_settings",
        return_value=create_mock_settings(),
    ):
        yield


class TestModelRegistryService:
    def test_get_instance_returns_singleton(self):
        instance1 = ModelRegistryService.get_instance()
        instance2 = ModelRegistryService.get_instance()
        assert instance1 is instance2

    def test_reset_instance_clears_singleton(self):
        instance1 = ModelRegistryService.get_instance()
        ModelRegistryService.reset_instance()
        instance2 = ModelRegistryService.get_instance()
        assert instance1 is not instance2


class TestGetModel:
    @pytest.mark.parametrize(
        "model_id,expected_provider",
        [
            ("gpt-4o", "openai"),
            ("claude-sonnet-4-5", "anthropic"),
            ("openrouter/anthropic/claude-3.5-sonnet", "openrouter"),
            ("fireworks_ai/accounts/fireworks/models/llama-v3p1-70b-instruct", "fireworks_ai"),
        ],
    )
    def test_returns_model_info_for_known_model(self, model_id: str, expected_provider: str):
        service = ModelRegistryService.get_instance()
        model = service.get_model(model_id)
        assert model is not None
        assert model.id == model_id
        assert model.provider == expected_provider

    def test_returns_none_for_unknown_model(self):
        service = ModelRegistryService.get_instance()
        model = service.get_model("unknown-model")
        assert model is None


class TestGetAvailableModels:
    def test_returns_all_chat_models_from_configured_providers(self):
        models = get_available_models("llm_gateway")
        model_ids = {m.id for m in models}
        assert "gpt-4o" in model_ids
        assert "claude-sonnet-4-5" in model_ids
        assert "claude-3-5-sonnet-20241022" in model_ids
        # New providers not configured by default
        assert "openrouter/anthropic/claude-3.5-sonnet" not in model_ids
        assert "fireworks_ai/accounts/fireworks/models/llama-v3p1-70b-instruct" not in model_ids

    def test_excludes_embedding_models(self):
        models = get_available_models("llm_gateway")
        model_ids = {m.id for m in models}
        assert "text-embedding-ada-002" not in model_ids

    def test_returns_model_info_objects(self):
        models = get_available_models("llm_gateway")
        assert all(isinstance(m, ModelInfo) for m in models)

    @pytest.mark.parametrize("alias", COST_ALIASES)
    def test_excludes_internal_cost_aliases(self, alias: str) -> None:
        with patch.dict(
            MOCK_COST_DATA,
            {alias: {"litellm_provider": "openai", "max_input_tokens": 128000, "mode": "chat"}},
        ):
            model_ids = {model.id for model in get_available_models("llm_gateway")}

        assert alias not in model_ids


class TestProviderFiltering:
    def test_only_returns_models_from_configured_providers(self):
        with patch(
            "llm_gateway.services.model_registry.get_settings",
            return_value=create_mock_settings(openai=True, anthropic=False),
        ):
            models = get_available_models("llm_gateway")
            providers = {m.provider for m in models}
            assert providers == {"openai"}
            model_ids = {m.id for m in models}
            assert "gpt-4o" in model_ids
            assert "claude-sonnet-4-5" not in model_ids

    def test_returns_empty_when_no_providers_configured(self):
        with patch(
            "llm_gateway.services.model_registry.get_settings",
            return_value=create_mock_settings(openai=False, anthropic=False),
        ):
            models = get_available_models("llm_gateway")
            assert len(models) == 0

    def test_returns_multiple_providers_when_configured(self):
        with patch(
            "llm_gateway.services.model_registry.get_settings",
            return_value=create_mock_settings(openai=True, anthropic=True),
        ):
            models = get_available_models("llm_gateway")
            providers = {m.provider for m in models}
            assert providers == {"openai", "anthropic"}

    @pytest.mark.parametrize(
        "provider_kwargs,expected_provider,expected_model_id",
        [
            (
                {"openai": False, "anthropic": False, "openrouter": True},
                "openrouter",
                "openrouter/anthropic/claude-3.5-sonnet",
            ),
            (
                {"openai": False, "anthropic": False, "fireworks": True},
                "fireworks_ai",
                "fireworks_ai/accounts/fireworks/models/llama-v3p1-70b-instruct",
            ),
        ],
    )
    def test_returns_single_new_provider_models_when_configured(
        self, provider_kwargs, expected_provider, expected_model_id
    ):
        with patch(
            "llm_gateway.services.model_registry.get_settings",
            return_value=create_mock_settings(**provider_kwargs),
        ):
            models = get_available_models("llm_gateway")
            providers = {m.provider for m in models}
            assert providers == {expected_provider}
            model_ids = {m.id for m in models}
            assert expected_model_id in model_ids

    def test_returns_all_four_providers_when_configured(self):
        with patch(
            "llm_gateway.services.model_registry.get_settings",
            return_value=create_mock_settings(openai=True, anthropic=True, openrouter=True, fireworks=True),
        ):
            models = get_available_models("llm_gateway")
            providers = {m.provider for m in models}
            assert providers == {"openai", "anthropic", "openrouter", "fireworks_ai"}


class TestCloudflareModelAdvertising:
    """CF Workers AI models aren't in litellm's cost map, so /v1/models advertises them explicitly
    when CF creds are set — gated on both key + account id, filtered by the product allowlist."""

    def _cf_ids(self, product: str) -> set[str]:
        return {m.id for m in get_available_models(product) if m.id.startswith("@cf/")}

    def test_cf_models_listed_when_configured(self):
        with patch(
            "llm_gateway.services.model_registry.get_settings",
            return_value=create_mock_settings(cloudflare=True),
        ):
            models = get_available_models("llm_gateway")
            cf = [m for m in models if m.id.startswith("@cf/")]
            assert {m.id for m in cf} == set(CLOUDFLARE_ALLOWED_MODELS)
            assert all(m.provider == "cloudflare" for m in cf)

    def test_cf_models_absent_when_not_configured(self):
        with patch(
            "llm_gateway.services.model_registry.get_settings",
            return_value=create_mock_settings(cloudflare=False),
        ):
            assert self._cf_ids("llm_gateway") == set()

    def test_cf_requires_both_key_and_account_id(self):
        settings = create_mock_settings(cloudflare=True)
        settings.cloudflare_account_id = None
        with patch("llm_gateway.services.model_registry.get_settings", return_value=settings):
            assert self._cf_ids("llm_gateway") == set()

    def test_cf_models_filtered_by_product_allowlist(self):
        # insights_code's allowlist includes @cf/zai-org/glm-5.2 but not the other CF model(s).
        with patch(
            "llm_gateway.services.model_registry.get_settings",
            return_value=create_mock_settings(cloudflare=True),
        ):
            assert self._cf_ids("insights_code") == {"@cf/zai-org/glm-5.2"}

    def test_modal_only_advertises_modal_served_models(self):
        # If CF creds are ever pulled after the Modal migration, GLM must stay advertised (it has a
        # Modal backend) while CF-only models like kimi drop off the listing.
        with patch(
            "llm_gateway.services.model_registry.get_settings",
            return_value=create_mock_settings(cloudflare=False, modal=True),
        ):
            assert self._cf_ids("llm_gateway") == {"@cf/zai-org/glm-5.2"}


class TestBasetenModelAdvertising:
    def test_deepseek_is_available_to_code_and_review(self):
        with patch(
            "llm_gateway.services.model_registry.get_settings",
            return_value=create_mock_settings(baseten=True),
        ):
            assert is_model_available("deepseek-ai/deepseek-v4-flash-0731", "review") is True
            assert is_model_available("deepseek-ai/deepseek-v4-flash-0731", "insights_code") is True
            assert is_model_available("deepseek-ai/deepseek-v4-flash-0731", "llm_gateway") is False
            assert is_model_available("deepseek-ai/deepseek-v4-flash-0731", "slack_app") is False
            for product in ("review", "insights_code"):
                model = next(
                    model
                    for model in ModelRegistryService.get_instance().get_available_models(product)
                    if model.id == "deepseek-ai/deepseek-v4-flash-0731"
                )
                assert model.context_window == 1_048_000


class TestModelMatchesAllowlist:
    @pytest.mark.parametrize(
        "model_id,expected",
        [
            ("gpt-4o", True),
            ("GPT-4O", True),
            ("Gpt-4o", True),
            ("claude-sonnet-4-5", True),
            ("CLAUDE-SONNET-4-5", True),
            ("unknown-model", False),
            ("gpt-4o-extra", False),
        ],
    )
    def test_case_insensitive_exact_matching(self, model_id: str, expected: bool):
        allowlist = frozenset({"gpt-4o", "claude-sonnet-4-5"})
        assert _model_matches_allowlist(model_id, allowlist) == expected


class TestIsModelAvailable:
    @pytest.mark.parametrize(
        "model_id,product,expected",
        [
            ("gpt-4o", "llm_gateway", True),
            ("gpt-4o", "insights_code", False),
            ("o1", "llm_gateway", True),
            ("o1", "insights_code", False),
            ("claude-opus-4-6", "insights_code", True),
            ("gpt-5.3-codex", "insights_code", True),
            ("gpt-5.2", "insights_code", True),
            ("gpt-5-mini", "insights_code", True),
            ("claude-opus-4-5", "insights_code", True),
            ("claude-sonnet-4-5", "insights_code", True),
            ("claude-haiku-4-5", "insights_code", True),
            # New providers not configured by default
            ("openrouter/anthropic/claude-3.5-sonnet", "llm_gateway", False),
            ("fireworks_ai/accounts/fireworks/models/llama-v3p1-70b-instruct", "llm_gateway", False),
            ("unknown-model", "llm_gateway", False),
            # Legacy aliases still work
            ("gpt-4o", "twig", False),
            ("claude-opus-4-6", "twig", True),
            ("gpt-4o", "array", False),
            ("claude-opus-4-5", "array", True),
        ],
    )
    def test_model_availability(self, model_id: str, product: str, expected: bool):
        assert is_model_available(model_id, product) == expected

    @pytest.mark.parametrize(
        "model_id,product,cloudflare,modal,expected",
        [
            # CF creds present + model priced/allowed -> available.
            ("@cf/zai-org/glm-5.2", "llm_gateway", True, False, True),
            # Same model, but no backend configured -> the runtime gate refuses it.
            ("@cf/zai-org/glm-5.2", "llm_gateway", False, False, False),
            # No CF creds, but Modal serves GLM -> still available.
            ("@cf/zai-org/glm-5.2", "llm_gateway", False, True, True),
            # Modal alone can't serve models without a Modal-served equivalent.
            ("@cf/moonshotai/kimi-k2.6", "llm_gateway", False, True, False),
            # CF configured but the product allowlist excludes this CF model -> unavailable.
            ("@cf/moonshotai/kimi-k2.6", "insights_code", True, False, False),
            # CF configured and the product allowlist includes it -> available.
            ("@cf/zai-org/glm-5.2", "insights_code", True, False, True),
        ],
    )
    def test_cf_model_availability_gated_on_creds(
        self, model_id: str, product: str, cloudflare: bool, modal: bool, expected: bool
    ):
        with patch(
            "llm_gateway.services.model_registry.get_settings",
            return_value=create_mock_settings(cloudflare=cloudflare, modal=modal),
        ):
            assert is_model_available(model_id, product) is expected

    def test_model_not_available_when_provider_not_configured(self):
        with patch.dict(os.environ, {}, clear=False):
            for var in PROVIDER_ENV_VARS:
                os.environ.pop(var, None)
            with patch(
                "llm_gateway.services.model_registry.get_settings",
                return_value=create_mock_settings(openai=False, anthropic=True),
            ):
                assert is_model_available("gpt-4o", "llm_gateway") is False
                assert is_model_available("claude-sonnet-4-5", "llm_gateway") is True