Newer
Older
navi-1 / tests / unit / llm / test_ollama.py
"""Unit tests for Ollama backend helpers."""

import asyncio

import httpx
import ollama as ollama_client
import pytest

from navi.exceptions import (
    LLMAuthError,
    LLMBackendError,
    LLMConnectionError,
    LLMModelNotFoundError,
    LLMQuotaExceededError,
    LLMServerUnavailableError,
)
from navi.llm.config import LLMConfig, Route, RouteEntry, Target
from navi.llm.fallback import FallbackOllamaBackend
from navi.llm.ollama import _base_options, _classify_error, _resolve_think


def _response_error(message: str, status: int) -> ollama_client.ResponseError:
    return ollama_client.ResponseError(message, status_code=status)


def test_classify_read_timeout_as_connection_error():
    err = _classify_error(httpx.ReadTimeout("timed out"))

    assert isinstance(err, LLMConnectionError)


def test_classify_empty_timeout_message_as_connection_error():
    err = _classify_error(httpx.ReadTimeout(""))

    assert isinstance(err, LLMConnectionError)
    assert str(err) == "ReadTimeout"


# ── response-error classification ────────────────────────────────────────────
# Each class decides what the route engine does next, so these are the whole
# reason quota exhaustion now fails over instead of killing the turn.

def test_a_429_is_a_quota_error():
    assert isinstance(_classify_error(_response_error("rate limit", 429)), LLMQuotaExceededError)


def test_a_quota_message_without_a_status_still_counts():
    assert isinstance(
        _classify_error(_response_error("quota exhausted", 400)), LLMQuotaExceededError
    )


def test_401_and_403_are_auth_errors():
    assert isinstance(_classify_error(_response_error("unauthorized", 401)), LLMAuthError)
    assert isinstance(_classify_error(_response_error("forbidden", 403)), LLMAuthError)


def test_a_404_is_a_missing_model():
    assert isinstance(
        _classify_error(_response_error("model not found", 404)), LLMModelNotFoundError
    )


def test_5xx_is_a_server_error():
    for status in (500, 502, 503, 504):
        assert isinstance(
            _classify_error(_response_error("boom", status)), LLMServerUnavailableError
        )


def test_a_plain_400_stays_generic():
    """A malformed request is not worth failing over for — every entry would say no."""
    err = _classify_error(_response_error("bad request", 400))
    assert type(err) is LLMBackendError


# ── options ─────────────────────────────────────────────────────────────────

def test_think_resolves_from_global_setting(monkeypatch):
    import navi.config as config_mod
    from navi.config import Settings

    monkeypatch.setattr(config_mod, "settings", Settings(ollama_think=True))

    assert _resolve_think(None) is True
    assert _resolve_think(False) is False


def test_base_options_do_not_include_think(monkeypatch):
    import navi.llm.ollama as ollama_mod
    from navi.config import Settings

    monkeypatch.setattr(ollama_mod, "settings", Settings(ollama_num_ctx=1234))

    opts = _base_options(0.2, top_k=20, top_p=0.8)

    assert opts == {
        "temperature": 0.2,
        "num_ctx": 1234,
        "top_k": 20,
        "top_p": 0.8,
    }


def test_extra_options_override_the_defaults(monkeypatch):
    import navi.llm.ollama as ollama_mod
    from navi.config import Settings

    monkeypatch.setattr(ollama_mod, "settings", Settings(ollama_num_ctx=1234))

    opts = _base_options(0.2, extra={"presence_penalty": 0.0, "temperature": 0.4})

    assert opts["presence_penalty"] == 0.0
    assert opts["temperature"] == 0.4


# ── client construction ─────────────────────────────────────────────────────

def _backend_for(host: str, api_key: str = "") -> FallbackOllamaBackend:
    target = Target(name="t", host=host, api_key=api_key)
    entry = RouteEntry(model="m", target=target)
    config = LLMConfig(
        targets={"t": target}, routes={"default": Route("default", (entry,))}
    )
    return FallbackOllamaBackend(config)


def test_fallback_client_uses_expanded_timeout(monkeypatch):
    import navi.llm.fallback as fallback_mod
    from navi.config import Settings

    captured = {}

    class FakeOllamaBackend:
        def __init__(self, **kwargs):
            captured.update(kwargs)

    monkeypatch.setattr(fallback_mod, "OllamaBackend", FakeOllamaBackend)
    monkeypatch.setattr(
        fallback_mod,
        "settings",
        Settings(
            ollama_request_timeout=30,
            llm_complete_timeout=120,
            llm_stream_first_chunk_timeout=180,
        ),
    )

    backend = _backend_for("http://ollama.test", api_key="secret")
    backend._client(backend.config.routes["default"].entries[0])

    assert captured["host"] == "http://ollama.test"
    assert captured["api_key"] == "secret"
    assert captured["timeout"] == 180


def test_clients_are_cached_per_target(monkeypatch):
    import navi.llm.fallback as fallback_mod

    built = []

    class FakeOllamaBackend:
        def __init__(self, **kwargs):
            built.append(kwargs)

    monkeypatch.setattr(fallback_mod, "OllamaBackend", FakeOllamaBackend)
    backend = _backend_for("http://ollama.test")
    entry = backend.config.routes["default"].entries[0]

    assert backend._client(entry) is backend._client(entry)
    assert len(built) == 1