opensquilla--opensquilla/tests/test_provider_model_catalog.py

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from opensquilla.provider.model_catalog import ModelCatalog, _corrections_budget_fallback


def test_deepseek_v4_direct_models_use_models_dev_limits() -> None:
    # The vendored models.dev snapshot supplies the real per-(provider, model)
    # budgets offline (PR #406 roadmap item 4); the packaged corrections
    # budget rows remain only the emergency floor beneath it.
    catalog = ModelCatalog()

    for model in ("deepseek-v4-flash", "deepseek-v4-pro"):
        assert catalog.resolve_context_window(model, "deepseek") == 1_000_000
        assert catalog.resolve_max_tokens(model, provider="deepseek") == 384_000
        caps = catalog.get_capabilities(model, provider_name="deepseek")
        assert caps.supports_reasoning is True
        assert caps.supports_tools is True
        assert caps.reasoning_format == "deepseek"


def test_provider_scoped_corrections_budget_outranks_snapshot_merge() -> None:
    """tokenrhythm has no models.dev table: without the provider-scoped
    corrections layer, the snapshot's cross-provider bare-id merge would
    serve the origin providers' budgets (202_752 zhipu glm-5, 262_144
    moonshot kimi) instead of the relay's own published limits. The rows
    mirror the platform listing (catalog_overrides.toml block comment);
    when the listing is reachable the boot-time live ingest supersedes
    them (see test_provider/test_live_catalog.py)."""
    catalog = ModelCatalog()

    assert catalog.resolve_context_window_with_source(
        "deepseek-v4-flash", provider="tokenrhythm"
    ) == (1_000_000, "catalog")
    assert catalog.resolve_max_tokens("deepseek-v4-flash", provider="tokenrhythm") == 384_000
    # Discriminating rows — the bare-id merge would report the origin
    # providers' windows here, not the relay's published ones.
    assert catalog.resolve_context_window("glm-5", provider="tokenrhythm") == 1_000_000
    assert catalog.resolve_context_window("kimi-k2.5", provider="tokenrhythm") == 256_000
    assert catalog.resolve_context_window("kimi-k2.7-code", provider="tokenrhythm") == 256_000
    assert catalog.resolve_max_tokens("kimi-k2.7-code", provider="tokenrhythm") == 128_000
    assert catalog.resolve_context_window("qwen3.7-max", provider="tokenrhythm") == 1_000_000
    assert catalog.resolve_max_tokens("qwen3.7-max", provider="tokenrhythm") == 256_000
    # The same bare id on direct DeepSeek keeps its own snapshot-table
    # budgets — the provider-scoped layer never leaks across providers.
    assert catalog.resolve_context_window("deepseek-v4-flash", "deepseek") == 1_000_000


def test_direct_profile_windows_resolve_from_models_dev_snapshot() -> None:
    catalog = ModelCatalog()

    expected_windows = {
        "gpt-5.4-nano": 400_000,
        "gpt-5.4-mini": 400_000,
        "gpt-5.5": 1_050_000,
        "glm-4.7-flashx": 200_000,
        # Real budget (202k) instead of the former 80k conservative placeholder.
        "glm-5": 202_752,
        "glm-5.1": 200_000,
        "z-ai/glm-5.2": 1_000_000,
        "moonshot-v1-8k": 8_192,
        "moonshot-v1-128k": 131_072,
        "kimi-k2.5": 262_144,
        "kimi-k2.6": 262_144,
    }

    for model_id, context_window in expected_windows.items():
        assert catalog.resolve_context_window(model_id) == context_window
        max_tokens = catalog.resolve_max_tokens(model_id)
        assert max_tokens > 0
        assert max_tokens <= context_window


def test_corrections_budget_fallback_is_provider_agnostic_by_basename() -> None:
    # The retired static table's budget slot now resolves from the packaged
    # corrections rows, keyed PROVIDER-AGNOSTICALLY by basename. The moonshot
    # window rows carry the exact values the static table did.
    assert _corrections_budget_fallback("moonshot-v1-8k") == (8_192, 8_192)
    assert _corrections_budget_fallback("moonshot-v1-32k") == (32_768, 32_768)
    assert _corrections_budget_fallback("moonshot-v1-128k") == (131_072, 131_072)
    # The vendor-qualified router-tier rows the snapshot only knows by their
    # slash id resolve for the bare basename too (grok's static tuple was
    # DEFAULT_MAX_TOKENS output / 1M window).
    assert _corrections_budget_fallback("grok-4.3") == (16_384, 1_000_000)
    assert _corrections_budget_fallback("step-3.5-flash") == (16_384, 256_000)
    # Glob capability-ladder rows are never consulted for budgets, and an
    # unknown basename yields None.
    assert _corrections_budget_fallback("model-nobody-knows") is None


def test_corrections_budget_qualified_and_unqualified_resolve_identically() -> None:
    catalog = ModelCatalog()
    for qualified, bare in (
        ("z-ai/glm-5", "glm-5"),
        ("deepseek/deepseek-v4-pro", "deepseek-v4-pro"),
        ("moonshot/moonshot-v1-8k", "moonshot-v1-8k"),
    ):
        assert catalog.resolve_context_window(qualified) == catalog.resolve_context_window(bare)
        assert catalog.resolve_max_tokens(qualified) == catalog.resolve_max_tokens(bare)


# Retirement parity net for the 25 keys of the deleted static fallback table.
# Expectations are LITERALS captured from the pre-change tree (both resolvers
# run against staging/provider-overhaul before the table was removed), for
# the two call shapes production uses: no provider (router-decision path)
# and the key's natural provider (turn-runner path). Any drift here means
# the retirement changed a resolution the static table used to decide.
#
# key → (natural provider,
#        (max_tokens, context_window) with provider="",
#        (max_tokens, context_window) with the natural provider)
_STATIC_RETIREMENT_PARITY: dict[str, tuple[str, tuple[int, int], tuple[int, int]]] = {
    "claude-opus-4.8": ("anthropic", (128_000, 1_000_000), (128_000, 1_000_000)),
    "claude-sonnet-4.6": ("anthropic", (128_000, 1_000_000), (128_000, 1_000_000)),
    "gemini-3.5-flash": ("gemini", (65_536, 1_048_576), (65_536, 1_048_576)),
    "gpt-5.4-nano": ("openai", (128_000, 400_000), (128_000, 400_000)),
    "gpt-5.4-mini": ("openai", (128_000, 400_000), (128_000, 400_000)),
    "gpt-5.5": ("openai", (128_000, 1_050_000), (128_000, 1_050_000)),
    "qwen3-coder-plus": ("dashscope", (65_536, 1_048_576), (65_536, 1_048_576)),
    "grok-4.3": ("xai", (16_384, 1_000_000), (16_384, 1_000_000)),
    "glm-4.5-air": ("zhipu", (98_304, 131_072), (98_304, 131_072)),
    "glm-4.6": ("zhipu", (131_072, 204_800), (131_072, 204_800)),
    "glm-4.7-flashx": ("zhipu", (131_072, 200_000), (131_072, 200_000)),
    "glm-5": ("zhipu", (16_384, 202_752), (131_072, 204_800)),
    "glm-5.1": ("zhipu", (128_000, 200_000), (131_072, 200_000)),
    "glm-5.2": ("zhipu", (128_000, 1_000_000), (131_072, 1_000_000)),
    "minimax-m2.5": ("minimax", (131_072, 204_800), (131_072, 204_800)),
    "minimax-m2.7": ("minimax", (131_072, 204_800), (131_072, 204_800)),
    "step-3.5-flash": ("stepfun", (16_384, 256_000), (16_384, 256_000)),
    "deepseek-v4-flash": ("deepseek", (384_000, 1_000_000), (384_000, 1_000_000)),
    "deepseek-v4-pro": ("deepseek", (384_000, 1_000_000), (384_000, 1_000_000)),
    "deepseek-v3.2": ("deepseek", (8_192, 128_000), (8_192, 128_000)),
    "moonshot-v1-8k": ("moonshot", (8_192, 8_192), (8_192, 8_192)),
    "moonshot-v1-32k": ("moonshot", (8_192, 32_768), (8_192, 32_768)),
    "moonshot-v1-128k": ("moonshot", (8_192, 131_072), (8_192, 131_072)),
    "kimi-k2.5": ("moonshot", (32_768, 262_144), (8_192, 262_144)),
    "kimi-k2.6": ("moonshot", (16_384, 262_144), (8_192, 262_144)),
}


def test_static_table_retirement_keeps_all_25_key_resolutions_identical() -> None:
    assert len(_STATIC_RETIREMENT_PARITY) == 25
    catalog = ModelCatalog()
    for model, (natural, bare_expected, natural_expected) in _STATIC_RETIREMENT_PARITY.items():
        observed_bare = (
            catalog.resolve_max_tokens(model),
            catalog.resolve_context_window(model),
        )
        assert observed_bare == bare_expected, (model, "")
        observed_natural = (
            catalog.resolve_max_tokens(model, provider=natural),
            catalog.resolve_context_window(model, natural),
        )
        assert observed_natural == natural_expected, (model, natural)


def test_populate_from_data_parses_openrouter_pricing() -> None:
    catalog = ModelCatalog()
    catalog._populate_from_data(
        [
            {
                "id": "vendor/priced-model",
                "context_length": 100_000,
                "pricing": {"prompt": "0.0000025", "completion": "0.00001"},
            },
            {"id": "vendor/free-model", "context_length": 8_192},
            {"id": "vendor/bad-pricing", "pricing": {"prompt": "n/a", "completion": None}},
        ]
    )

    priced = catalog.get("vendor/priced-model")
    assert priced is not None
    assert priced.input_cost_per_1k == pytest.approx(0.0025)
    assert priced.output_cost_per_1k == pytest.approx(0.01)
    # Missing pricing block → 0.0
    assert catalog.get("vendor/free-model").input_cost_per_1k == 0.0
    # Non-numeric / None → 0.0 (no crash)
    bad = catalog.get("vendor/bad-pricing")
    assert bad.input_cost_per_1k == 0.0
    assert bad.output_cost_per_1k == 0.0


def test_openrouter_near_context_completion_window_uses_safe_default() -> None:
    catalog = ModelCatalog()
    catalog._populate_from_data(
        [
            {
                "id": "provider/vision-model",
                "context_length": 262_144,
                "top_provider": {"max_completion_tokens": 262_142},
            }
        ]
    )

    assert catalog.resolve_context_window("provider/vision-model") == 262_144
    assert catalog.resolve_max_tokens("provider/vision-model") == 8192


def test_openrouter_safe_default_never_raises_smaller_provider_limit() -> None:
    catalog = ModelCatalog()
    catalog._populate_from_data(
        [
            {
                "id": "provider/smaller-output-model",
                "context_length": 12_000,
                "top_provider": {"max_completion_tokens": 4096},
            }
        ]
    )

    assert catalog.resolve_context_window("provider/smaller-output-model") == 12_000
    assert catalog.resolve_max_tokens("provider/smaller-output-model") == 4096


def _catalog_with_live_reasoning_model() -> ModelCatalog:
    catalog = ModelCatalog()
    catalog._populate_from_data(
        [
            {
                "id": "vendor/reasoning-model",
                "context_length": 200_000,
                "top_provider": {"max_completion_tokens": 128_000},
                "supported_parameters": ["reasoning", "tools", "tool_choice"],
                "architecture": {"input_modalities": ["text", "image"]},
            }
        ]
    )
    return catalog


def test_get_capabilities_honors_user_reasoning_override() -> None:
    catalog = _catalog_with_live_reasoning_model()
    catalog.set_user_overrides(
        {
            "vendor/reasoning-model": {
                "supports_reasoning": False,
                "reasoning_format": "none",
            }
        }
    )

    caps = catalog.get_capabilities("vendor/reasoning-model", provider_name="openrouter")

    assert caps.supports_reasoning is False
    assert caps.reasoning_format == "none"


def test_get_capabilities_honors_user_vision_override_on_live_reasoning_model() -> None:
    catalog = _catalog_with_live_reasoning_model()
    catalog.set_user_overrides({"vendor/reasoning-model": {"supports_vision": False}})

    caps = catalog.get_capabilities("vendor/reasoning-model", provider_name="openrouter")

    assert caps.supports_reasoning is True
    assert caps.supports_vision is False


@pytest.mark.parametrize(
    ("reasoning_format", "supports_reasoning"),
    [("deepseek", True), ("none", False)],
)
def test_get_capabilities_honors_user_reasoning_format_override_on_live_model(
    reasoning_format: str,
    supports_reasoning: bool,
) -> None:
    catalog = _catalog_with_live_reasoning_model()
    catalog.set_user_overrides(
        {"vendor/reasoning-model": {"reasoning_format": reasoning_format}}
    )

    caps = catalog.get_capabilities("vendor/reasoning-model", provider_name="openrouter")

    assert caps.supports_reasoning is supports_reasoning
    assert caps.reasoning_format == reasoning_format


def test_resolve_context_window_honors_user_override() -> None:
    catalog = ModelCatalog()
    catalog.set_user_overrides({"vllm/self-hosted-model": {"context_window": 131_072}})

    window, source = catalog.resolve_context_window_with_source("self-hosted-model", "vllm")

    assert window == 131_072
    assert source == "override"


def test_resolve_max_tokens_honors_user_override() -> None:
    catalog = ModelCatalog()
    catalog.set_user_overrides(
        {
            "vllm/self-hosted-model": {
                "context_window": 131_072,
                "max_output_tokens": 32_768,
            }
        }
    )

    effective, source = catalog.resolve_max_tokens_with_source(
        "self-hosted-model", provider="vllm"
    )

    assert effective == 32_768
    assert source == "override"


@pytest.mark.asyncio
async def test_fetch_openrouter_adds_app_attribution_headers() -> None:
    captured: dict[str, object] = {}
    mock_response = MagicMock()
    mock_response.raise_for_status = MagicMock()
    mock_response.json.return_value = {
        "data": [
            {
                "id": "openai/gpt-4o",
                "name": "GPT-4o",
                "context_length": 128_000,
                "top_provider": {"max_completion_tokens": 16_384},
            }
        ]
    }

    with patch("opensquilla.provider.model_catalog.httpx.AsyncClient") as mock_client_cls:
        mock_client = AsyncMock()
        mock_client.__aenter__ = AsyncMock(return_value=mock_client)
        mock_client.__aexit__ = AsyncMock(return_value=False)

        async def capture_get(url, *, headers):
            captured["url"] = url
            captured["headers"] = headers
            return mock_response

        mock_client.get = AsyncMock(side_effect=capture_get)
        mock_client_cls.return_value = mock_client

        catalog = ModelCatalog()
        await catalog.fetch_openrouter(api_key="test-key", base_url="https://openrouter.ai/api")

    assert captured["url"] == "https://openrouter.ai/api/v1/models"
    assert captured["headers"] == {
        "Authorization": "Bearer test-key",
        "HTTP-Referer": "https://opensquilla.ai",
        "X-Title": "OpenSquilla",
    }
    model = catalog.get("openai/gpt-4o")
    assert model is not None
    assert model.context_window == 128_000


def test_local_provider_context_window_uses_runtime_default_not_cloud() -> None:
    from opensquilla.provider.model_catalog import (
        _LOCAL_CONTEXT_WINDOW,
        DEFAULT_CONTEXT_WINDOW,
    )

    catalog = ModelCatalog()
    # Without provider context a bare ollama id falls to the 200k cloud default.
    assert catalog.resolve_context_window("qwen3:4b") == DEFAULT_CONTEXT_WINDOW
    # With the local provider it reports the runtime window instead.
    assert catalog.resolve_context_window("qwen3:4b", provider="ollama") == _LOCAL_CONTEXT_WINDOW
    assert _LOCAL_CONTEXT_WINDOW < DEFAULT_CONTEXT_WINDOW


def test_local_provider_max_tokens_clamped_to_local_window() -> None:
    from opensquilla.provider.model_catalog import _LOCAL_CONTEXT_WINDOW

    catalog = ModelCatalog()
    # max_tokens cannot exceed the (smaller) local context window.
    assert catalog.resolve_max_tokens("llama3.2:3b", provider="ollama") <= _LOCAL_CONTEXT_WINDOW


def test_cloud_provider_context_window_unchanged() -> None:
    from opensquilla.provider.model_catalog import DEFAULT_CONTEXT_WINDOW

    catalog = ModelCatalog()
    # An unknown cloud model id is unaffected by the provider argument.
    assert catalog.resolve_context_window("some-cloud-model", provider="openai") == (
        DEFAULT_CONTEXT_WINDOW
    )