pablan/backend/tests/test_llm_client.py

import pytest
from pydantic import BaseModel

from app.llm.client import chat_json, chat_stream, embed
from app.llm.errors import LLMError
from app.metrics import metrics
from tests.fake_openai import FakeOpenAI

PING = [{"role": "user", "content": "hi"}]


class Verdict(BaseModel):
    covered: list[str]
    done: bool


def _counter(name: str, **labels: str) -> float:
    for entry in metrics.snapshot()["counters"].get(name, []):
        if entry["labels"] == labels:
            return entry["value"]
    return 0.0


async def test_chat_stream_yields_deltas_and_records_metrics(
    fake_llm: FakeOpenAI,
) -> None:
    fake_llm.chat_responses.append({"chunks": ["Hel", "lo"]})
    out = [token async for token in chat_stream(PING)]
    assert out == ["Hel", "lo"]

    body = fake_llm.requests[-1]
    assert body["stream"] is True
    assert body["stream_options"] == {"include_usage": True}
    assert (
        _counter("llm_calls_total", role="chat", kind="chat_stream", status="ok") == 1
    )
    assert _counter("llm_tokens_total", role="chat", direction="completion") == 3


async def test_chat_stream_role_override(fake_llm: FakeOpenAI) -> None:
    [token async for token in chat_stream(PING, role="utility", max_tokens=1)]
    assert fake_llm.requests[-1]["max_tokens"] == 1
    assert (
        _counter("llm_calls_total", role="utility", kind="chat_stream", status="ok")
        == 1
    )


async def test_chat_stream_error_is_sanitized(fake_llm: FakeOpenAI) -> None:
    # Twice: the client retries once on connection/server errors.
    fake_llm.chat_responses.extend([{"status": 500}, {"status": 500}])
    with pytest.raises(LLMError) as excinfo:
        [token async for token in chat_stream(PING)]
    error = excinfo.value
    assert "chat_stream failed" in str(error)
    # The server error body must not leak into the exception message.
    assert "induced failure" not in str(error)
    assert error.role == "chat"
    assert error.kind == "chat_stream"
    assert error.status == "error"
    assert error.status_code == 500
    assert error.cause_type  # original exception class name only
    assert isinstance(error.duration_ms, int)
    assert (
        _counter("llm_calls_total", role="chat", kind="chat_stream", status="error")
        == 1
    )


async def test_chat_json_happy_path(fake_llm: FakeOpenAI) -> None:
    fake_llm.chat_responses.append({"content": '{"covered": ["a"], "done": true}'})
    result = await chat_json(PING, Verdict)
    assert result == Verdict(covered=["a"], done=True)

    body = fake_llm.requests[-1]
    assert body["response_format"]["type"] == "json_schema"
    assert body["response_format"]["json_schema"]["name"] == "Verdict"
    assert (
        _counter("llm_calls_total", role="utility", kind="chat_json", status="ok") == 1
    )


async def test_chat_json_retries_once_on_invalid_json(fake_llm: FakeOpenAI) -> None:
    fake_llm.chat_responses.extend(
        [
            {"content": "definitely not json"},
            {"content": '{"covered": [], "done": false}'},
        ]
    )
    result = await chat_json(PING, Verdict)
    assert result.done is False
    assert len(fake_llm.requests) == 2
    retry_messages = fake_llm.requests[-1]["messages"]
    assert retry_messages[-1]["role"] == "user"
    assert "JSON" in retry_messages[-1]["content"]
    assert retry_messages[-2] == {"role": "assistant", "content": "definitely not json"}
    assert (
        _counter("llm_calls_total", role="utility", kind="chat_json", status="invalid")
        == 1
    )


async def test_chat_json_gives_up_after_retry_without_leaking(
    fake_llm: FakeOpenAI,
) -> None:
    fake_llm.chat_responses.extend(
        [{"content": "SECRET-A 123"}, {"content": "SECRET-B 456"}]
    )
    with pytest.raises(LLMError) as excinfo:
        await chat_json(PING, Verdict)
    error = excinfo.value
    assert "Verdict" in str(error)
    # Structured debugging metadata is present ...
    assert error.role == "utility"
    assert error.kind == "chat_json"
    assert error.status == "invalid"
    assert error.cause_type == "ValidationError"
    assert error.attempt == 2
    assert error.status_code is None
    assert isinstance(error.duration_ms, int)
    # ... and neither message nor ANY metadata field carries content.
    everything = str(error) + repr(vars(error))
    assert "SECRET-A" not in everything
    assert "SECRET-B" not in everything


async def test_chat_json_http_error(fake_llm: FakeOpenAI) -> None:
    # Twice: the client retries once on connection/server errors.
    fake_llm.chat_responses.extend([{"status": 503}, {"status": 503}])
    with pytest.raises(LLMError) as excinfo:
        await chat_json(PING, Verdict)
    error = excinfo.value
    assert error.status == "error"
    assert error.status_code == 503
    assert error.attempt == 1
    assert error.cause_type
    assert "induced failure" not in str(error) + repr(vars(error))
    assert (
        _counter("llm_calls_total", role="utility", kind="chat_json", status="error")
        == 1
    )


async def test_embed_preserves_order(fake_llm: FakeOpenAI) -> None:
    vectors = await embed(["a", "b", "c"])
    assert len(vectors) == 3
    assert vectors[0][0] == 0.0
    assert vectors[2][0] == 2.0
    assert len(vectors[0]) == fake_llm.embedding_dim
    assert _counter("llm_calls_total", role="embedding", kind="embed", status="ok") == 1


@pytest.mark.parametrize(
    ("cause_type", "status_code", "expected"),
    [
        ("APIConnectionError", None, "llm_unreachable"),
        ("ConnectError", None, "llm_unreachable"),
        # A timeout is the busy case, not the down case: the endpoint took the
        # request and never came back.
        ("APITimeoutError", None, "llm_busy"),
        ("RateLimitError", 429, "llm_busy"),
        ("APIStatusError", 503, "llm_busy"),
        ("AuthenticationError", 401, "llm_misconfigured"),
        ("NotFoundError", 404, "llm_misconfigured"),
        ("BadRequestError", 400, "llm_failed"),
        ("ValueError", None, "llm_failed"),
    ],
)
def test_error_code_separates_down_from_busy_from_misconfigured(
    cause_type: str, status_code: int | None, expected: str
) -> None:
    """One classification for every surface: retry-in-a-moment, start the
    endpoint, and fix the config must not read the same to the user."""
    error = LLMError(
        "boom",
        role="chat",
        kind="chat_stream",
        status="error",
        cause_type=cause_type,
        status_code=status_code,
    )
    assert error.code == expected