tests/test_model_router.py

import asyncio
from types import SimpleNamespace

from pnp_lab.config import Settings
from pnp_lab.model_router import ModelRouter


class FakeCompletions:
    def __init__(self, responses):
        self.responses = list(responses)
        self.calls = []

    async def create(self, **kwargs):
        self.calls.append(kwargs)
        item = self.responses.pop(0)
        if isinstance(item, Exception):
            raise item
        return item


class FakeClient:
    def __init__(self, responses):
        self.chat = SimpleNamespace(completions=FakeCompletions(responses))


def response(content, *, reasoning="", finish="stop", model="moonshotai/Kimi-K3", prompt_tokens=10, completion_tokens=20):
    msg = SimpleNamespace(content=content, reasoning_content=reasoning, model_extra={})
    choice = SimpleNamespace(message=msg, finish_reason=finish)
    usage = SimpleNamespace(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens)
    return SimpleNamespace(choices=[choice], usage=usage, model=model, id="req_test", _request_id="req_test")


def make_router(responses):
    s = Settings()
    s.hf_token = "hf_test_token"
    s.kimi_reasoning_effort = "high"
    router = object.__new__(ModelRouter)
    router.settings = s
    router._sem = asyncio.Semaphore(1)
    router._client = FakeClient(responses)
    router._prices = {"moonshotai/Kimi-K3": (3.0, 15.0), "zai-org/GLM-5.2": (0.75, 2.4)}
    router._models_cache = []
    return router


def test_empty_kimi_content_is_retried_with_lower_reasoning_effort():
    router = make_router([
        response("", reasoning="x" * 4000, finish="length"),
        response('{"ok": true}', reasoning="x" * 100, finish="stop"),
    ])
    diagnostics = []
    result = asyncio.run(router.chat(
        "moonshotai/Kimi-K3",
        "system",
        "user",
        max_tokens=1000,
        retries=2,
        diagnostic_cb=diagnostics.append,
    ))
    assert result.text == '{"ok": true}'
    assert [d["status"] for d in diagnostics] == ["empty_content", "success"]
    assert diagnostics[0]["reasoning_effort"] == "high"
    assert diagnostics[1]["reasoning_effort"] == "low"
    assert diagnostics[0]["reasoning_chars"] == 4000


def test_router_falls_back_to_second_model_after_empty_primary():
    router = make_router([
        response("", reasoning="x" * 20, finish="length"),
        response('{"ok": true}', reasoning="", finish="stop", model="zai-org/GLM-5.2"),
    ])
    diagnostics = []
    result = asyncio.run(router.chat(
        "moonshotai/Kimi-K3",
        "system",
        "user",
        max_tokens=1000,
        retries=1,
        fallback_models=["zai-org/GLM-5.2"],
        diagnostic_cb=diagnostics.append,
    ))
    assert result.model_used == "zai-org/GLM-5.2"
    assert [d["status"] for d in diagnostics] == ["empty_content", "success"]
    assert diagnostics[-1]["model"] == "zai-org/GLM-5.2"