Newer
Older
navi-1 / tests / unit / llm / test_model_stamping.py
"""The resolved model and target reach the UI.

The route engine picks the entry that actually answered; the UI needs to know
both which model and which host — "qwen3.6:35b" on the LAN box means something
different from the same name in the cloud, and a backup entry is worth
surfacing as such.
"""

import asyncio

from navi.exceptions import LLMModelNotFoundError
from navi.llm.base import LLMChunk, LLMResponse
from navi.llm.config import LLMConfig, Route, RouteEntry, Target
from navi.llm.fallback import FallbackOllamaBackend


def _build(monkeypatch, serve, chain):
    """A backend whose per-target clients are ``serve``-driven fakes."""
    targets = {}
    entries = []
    for model, target_name in chain:
        target = targets.setdefault(
            target_name, Target(name=target_name, host=f"http://{target_name}.test")
        )
        entries.append(RouteEntry(model=model, target=target))
    config = LLMConfig(
        targets=targets,
        routes={"default": Route("default", tuple(entries))},
    )
    backend = FallbackOllamaBackend(config)
    backend._client = lambda entry: serve(entry)  # type: ignore[assignment]
    return backend


class FakeClient:
    def __init__(self, served, fail_models=(), chunks=None):
        self.served = served
        self.fail_models = set(fail_models)
        self.chunks = chunks

    async def complete(self, *args, model=None, **kwargs):
        self.served.append(model)
        if model in self.fail_models:
            raise LLMModelNotFoundError("missing")
        return LLMResponse(content="ok", tool_calls=None, finish_reason="stop")

    async def stream_complete(self, *args, model=None, **kwargs):
        self.served.append(model)
        if model in self.fail_models:
            raise LLMModelNotFoundError("missing")
        for chunk in self.chunks or []:
            yield chunk


def test_fallback_stamps_resolved_model_on_response(monkeypatch):
    served: list[str] = []
    client = FakeClient(served)
    backend = _build(monkeypatch, lambda entry: client, [("first-choice", "home")])

    result = asyncio.run(backend.complete([], model=["first-choice", "second-choice"]))

    assert served == ["first-choice"]
    assert result.model == "first-choice"


def test_a_missing_model_moves_down_the_chain_and_stamps_the_survivor(monkeypatch):
    served: list[str] = []
    client = FakeClient(served, fail_models={"down-model"})
    backend = _build(
        monkeypatch, lambda entry: client, [("down-model", "cloud"), ("alive-model", "home")]
    )

    result = asyncio.run(backend.complete([], route="default"))

    assert served == ["down-model", "alive-model"]
    assert result.model == "alive-model"
    assert result.target == "home"


def test_fallback_stamps_resolved_model_on_first_chunk(monkeypatch):
    chunks = [
        LLMChunk(delta="hello"),
        LLMChunk(delta=" world"),
        LLMChunk(delta=None, finish_reason="stop", prompt_tokens=5, completion_tokens=2),
    ]
    served: list[str] = []
    client = FakeClient(served, chunks=chunks)
    backend = _build(monkeypatch, lambda entry: client, [("the-model", "home")])

    out = []

    async def collect():
        async for chunk in backend.stream_complete([], route="default"):
            out.append(chunk)

    asyncio.run(collect())

    assert served == ["the-model"]
    assert len(out) == 3
    assert out[0].model == "the-model"
    assert out[0].target == "home"
    assert out[1].model is None
    assert out[2].model is None


def test_ollama_backend_stamps_first_chunk_only(monkeypatch):
    """OllamaBackend stamps the resolved model on the first chunk, not later ones."""
    from navi.llm.ollama import OllamaBackend

    class FakeChunk:
        def __init__(self, content, done=False):
            self.message = type(
                "M", (), {"content": content, "thinking": None, "tool_calls": None}
            )()
            self.done = done
            self.prompt_eval_count = 5 if done else None
            self.eval_count = 2 if done else None

    class FakeClient:
        async def chat(self, **kwargs):
            async def _gen():
                yield FakeChunk("hi")
                yield FakeChunk(" there", done=True)

            return _gen()

    backend = OllamaBackend(model="default-model", host="http://ollama.test", api_key="")
    monkeypatch.setattr(backend, "_client", FakeClient())

    out = []

    async def collect():
        async for chunk in backend.stream_complete([], model=["configured-model"]):
            out.append(chunk)

    asyncio.run(collect())

    assert out[0].model == "configured-model"
    assert out[1].model is None