6krrt/tests/test_poller_parsing.py

"""Tests for poller.py's pure catalog-parsing helpers.

The NeuralWatt catalog encodes serving class in the model id and access
gating only in prose, so both are parsed rather than read from a field.
Cases below are taken from the live catalog.
"""

import pytest

from poller import parse_access_level, parse_base_model_id, parse_serving_class


# --- serving class --------------------------------------------------------

@pytest.mark.parametrize(
    "model_id,expected",
    [
        # base rows take the schema defaults
        ("gemma-4-31b", ("standard", "default", "full")),
        ("kimi-k3", ("standard", "default", "full")),
        # one dimension at a time
        ("kimi-k3-fast", ("standard", "reduced", "full")),
        ("kimi-k3-flex", ("flex", "default", "full")),
        ("glm-5.2-short", ("standard", "default", "short")),
        # the dimensions are orthogonal and combine freely
        ("glm-5.2-short-fast", ("standard", "reduced", "short")),
        ("glm-5.2-short-flex", ("flex", "default", "short")),
        ("glm-5.2-short-fast-flex", ("flex", "reduced", "short")),
    ],
)
def test_parse_serving_class(model_id, expected):
    assert parse_serving_class(model_id) == expected


def test_flash_is_not_read_as_fast():
    # Given: a base model whose name ends in 'flash' — a near-miss for the
    # '-fast' suffix that a substring match would misclassify, wrongly
    # demoting the cheapest model in the catalog out of tier 1
    # Then: whole-segment matching leaves it a plain standard row
    assert parse_serving_class("deepseek-v4-flash") == ("standard", "default", "full")
    assert parse_serving_class("deepseek-v4-flash-flex") == ("flex", "default", "full")


def test_namespaced_and_mixed_case_ids():
    # Given: the HF-style duplicate row, which carries a namespace and caps
    assert parse_serving_class("deepseek-ai/DeepSeek-V4-Flash") == (
        "standard",
        "default",
        "full",
    )


def test_suffix_only_id_is_not_stripped_to_nothing():
    # Given: a degenerate id consisting solely of a suffix token, the loop
    # must leave at least one segment rather than consuming the whole id
    assert parse_serving_class("fast") == ("standard", "default", "full")


# --- access level ---------------------------------------------------------

@pytest.mark.parametrize(
    "display_name,description,expected",
    [
        ("GLM-5.2 (short)", "... Private preview (grant-gated).", "preview"),
        ("GLM-5.2", "Private GLM-5.2 test canary", "canary"),
        # gating can appear in the display name rather than the description
        ("DeepSeek V4 Flash 0731 (Canary)", "1M context window.", "canary"),
        ("Gemma 4 31B", "Google Gemma 4 31B — multimodal with tool calling.", "public"),
        (None, None, "public"),
    ],
)
def test_parse_access_level(display_name, description, expected):
    assert parse_access_level(display_name, description) == expected


def test_preview_takes_precedence_over_canary():
    # Given: prose mentioning both, the more restrictive label wins
    assert parse_access_level("x", "Private preview (grant-gated) canary") == "preview"


# --- model family ---------------------------------------------------------

@pytest.mark.parametrize(
    "model_id,expected",
    [
        ("gemma-4-31b", "gemma-4-31b"),
        ("kimi-k3", "kimi-k3"),
        ("kimi-k3-fast", "kimi-k3"),
        ("kimi-k3-flex", "kimi-k3"),
        # every serving dimension collapses, in any combination
        ("glm-5.2-short-fast-flex", "glm-5.2"),
        ("glm-5.2-short", "glm-5.2"),
        # the HF-style duplicate row folds into the same family
        ("deepseek-ai/DeepSeek-V4-Flash", "deepseek-v4-flash"),
        # ...without 'flash' being mistaken for the '-fast' suffix
        ("deepseek-v4-flash", "deepseek-v4-flash"),
        ("deepseek-v4-flash-flex", "deepseek-v4-flash"),
    ],
)
def test_parse_base_model_id(model_id, expected):
    assert parse_base_model_id(model_id) == expected


def test_family_never_collapses_to_empty():
    # A degenerate id made only of suffix tokens must keep a segment, or every
    # such row would share one meaningless family
    assert parse_base_model_id("flex") == "flex"


# --- per-model context overrides ------------------------------------------
#
# The knob validated and was never read, while config.yaml shipped a worked
# example for it. Anyone who followed that example got silence.

def _row(**overrides):
    from poller import ModelRow

    fields = dict(
        model_id="qwen3.6-35b", provider="neuralwatt", base_model_id="qwen3.6-35b",
        display_name=None, cost_per_1m_prompt=None, cost_per_1m_completion=None,
        cost_per_1m_prompt_cached=None, context_window=200_000,
        max_output_tokens=None, supports_tools=True, supports_json_mode=True,
        supports_vision=False, supports_reasoning=True,
        reasoning_default_enabled=True, latency_class="standard",
        reasoning_mode="default", context_variant="full", access_level="public",
        pricing_tbd=False, deprecated=False,
    )
    fields.update(overrides)
    return ModelRow(**fields)


def _cfg(overrides=None):
    import copy

    import yaml

    from config import RouterConfig

    with open("config.yaml") as fh:
        raw = copy.deepcopy(yaml.safe_load(fh))
    raw["context"]["per_model_overrides"] = overrides or {}
    return RouterConfig(**raw)


def test_the_global_safety_factor_applies_without_an_override():
    # 200000 * 0.75 - 4096
    assert _row().effective_context_window(_cfg()) == 145_904


def test_a_per_model_safety_factor_is_actually_applied():
    cfg = _cfg({"qwen3.6-35b": {"safety_factor": 0.85}})

    # 200000 * 0.85 - 4096
    assert _row().effective_context_window(cfg) == 165_904


def test_an_override_for_another_model_does_not_leak():
    cfg = _cfg({"kimi-k3": {"safety_factor": 0.85}})

    assert _row().effective_context_window(cfg) == 145_904


def test_a_zero_output_reserve_means_zero_not_unset():
    """Present-but-falsy: `or` would silently fall through to 4096."""
    cfg = _cfg({"qwen3.6-35b": {"output_reserve_tokens": 0}})

    assert _row().effective_context_window(cfg) == 150_000


def test_an_override_reserve_beats_the_models_advertised_ceiling():
    cfg = _cfg({"qwen3.6-35b": {"output_reserve_tokens": 8192}})

    assert _row(max_output_tokens=16384).effective_context_window(cfg) == 141_808


def test_each_key_falls_back_independently():
    cfg = _cfg({"qwen3.6-35b": {"safety_factor": 0.85}})

    # factor overridden, reserve still the global default
    assert _row().effective_context_window(cfg) == 165_904


def test_a_typo_inside_an_override_is_an_error():
    """The point of typing it: StrictModel has to reach inside the block."""
    with pytest.raises(Exception, match="safety_factr|Extra inputs"):
        _cfg({"qwen3.6-35b": {"safety_factr": 0.85}})


def test_an_out_of_range_override_is_an_error():
    with pytest.raises(Exception, match="safety_factor"):
        _cfg({"qwen3.6-35b": {"safety_factor": 1.5}})