brain-of-reese/tests/unit/test_llm_client.py

"""Unit tests: LLMClient embeddings (batching, order, loud dim failure).

The fakes stand in at the httpx-transport layer — that is where LLMClient
actually talks to the endpoint (see ``LLMClient._embed_batch`` in
``app/rag/llm.py`` for why the openai SDK's own ``embeddings.create`` is
bypassed: it injects ``encoding_format``, which aipi's litellm proxy
rejects).
"""
from __future__ import annotations

import asyncio
import json
from types import SimpleNamespace
from typing import Any

import pytest

from app.config import Settings
from app.rag.llm import EmbeddingDimensionError, EmbeddingError, LLMClient, LLMError


def _settings(**kwargs: Any) -> Settings:
    kwargs.setdefault("_env_file", None)
    return Settings(**kwargs)  # pyright: ignore[reportCallIssue]  (kwarg exists at runtime)


class _Row:
    def __init__(self, index: int, embedding: list[float]) -> None:
        self.index = index
        self.embedding = embedding


class _Response:
    def __init__(self, rows: list[_Row]) -> None:
        self.data = rows


class _FakeEmbeddingsService:
    """Simulates the /embeddings endpoint; records calls; can fail."""

    def __init__(
        self,
        dim: int = 768,
        fail: Exception | None = None,
        drop_index: int = -1,
        http_error: int | None = None,
        too_large_min: int | None = None,
    ) -> None:
        self.dim = dim
        self.fail = fail
        self.drop_index = drop_index
        self.http_error = http_error
        self.too_large_min = too_large_min
        self.calls: list[list[str]] = []

    async def create(self, *, model: str, input: list[str]) -> _Response:
        self.calls.append(list(input))
        if self.fail is not None:
            raise self.fail
        return _Response(
            [
                _Row(i, [0.5] * self.dim)
                for i in range(len(input))
                if i != self.drop_index
            ]
        )


class _FakeHttpResponse:
    def __init__(
        self, status_code: int, payload: dict[str, Any] | None = None, text: str = ""
    ) -> None:
        self.status_code = status_code
        self._payload = payload
        self.text = text or (json.dumps(payload) if payload is not None else "boom")

    def json(self) -> Any:
        if self._payload is None:
            raise ValueError("no json body")
        return self._payload


#: The endpoint's real error phrasing (litellm) — the client keys off it.
_TOO_LARGE_TEXT = 'input (9999 tokens) is too large to process. increase the physical batch size'


class _FakeHttp:
    """Stands in for the httpx transport the openai client owns."""

    def __init__(self, service: _FakeEmbeddingsService) -> None:
        self.service = service
        self.bodies: list[dict[str, Any]] = []

    async def post(
        self, url: str, *, json: dict[str, Any], headers: dict[str, str] | None = None
    ) -> _FakeHttpResponse:
        self.bodies.append(json)
        assert "Authorization" in (headers or {})
        if self.service.http_error is not None:
            return _FakeHttpResponse(self.service.http_error)
        if (
            self.service.too_large_min is not None
            and len(json["input"]) >= self.service.too_large_min
        ):
            return _FakeHttpResponse(500, None, _TOO_LARGE_TEXT)
        rows = await self.service.create(model=json["model"], input=json["input"])
        payload = {"data": [{"index": r.index, "embedding": r.embedding} for r in rows.data]}
        return _FakeHttpResponse(200, payload)


class _FakeClient:
    """Stands in for the openai AsyncOpenAI object (only its transport is used)."""

    def __init__(self, http: _FakeHttp) -> None:
        self._client = http


def _make_client(service: _FakeEmbeddingsService, **kwargs: Any) -> tuple[LLMClient, _FakeHttp]:
    kwargs.setdefault("embed_batch_size", 2)
    llm = LLMClient(_settings(**kwargs))
    http = _FakeHttp(service)
    llm._client = _FakeClient(http)  # pyright: ignore[reportAttributeAccessIssue]
    return llm, http


def test_embed_batches_by_batch_size_and_keeps_order() -> None:
    service = _FakeEmbeddingsService()
    llm, http = _make_client(service)
    texts = [f"t{i}" for i in range(5)]
    vecs = asyncio.run(llm.embed(texts))

    assert [len(c) for c in service.calls] == [2, 2, 1]
    assert [t for call in service.calls for t in call] == texts
    assert len(vecs) == 5
    assert all(len(v) == 768 for v in vecs)
    assert llm.embed_batches == 3
    # aipi (litellm) rejects the SDK's injected "encoding_format" — the
    # payload must stay a minimal {model, input} body.
    assert all(set(b) == {"model", "input"} for b in http.bodies)


def test_embed_empty_returns_empty_without_calling_endpoint() -> None:
    service = _FakeEmbeddingsService()
    llm, http = _make_client(service)
    assert asyncio.run(llm.embed([])) == []
    assert http.bodies == []
    assert service.calls == []
    assert llm.embed_batches == 0


def test_embed_one_returns_single_vector() -> None:
    llm, _ = _make_client(_FakeEmbeddingsService())
    vec = asyncio.run(llm.embed_one("hello"))
    assert len(vec) == 768


def test_dim_mismatch_fails_loudly_with_actionable_message() -> None:
    llm, _ = _make_client(_FakeEmbeddingsService(dim=512))
    with pytest.raises(EmbeddingDimensionError) as exc:
        asyncio.run(llm.embed(["hello"]))
    msg = str(exc.value)
    assert "512" in msg and "768" in msg
    assert "BOR_EMBEDDING_DIM" in msg
    assert "llm_probe" in msg


def test_endpoint_error_is_wrapped() -> None:
    llm, _ = _make_client(_FakeEmbeddingsService(fail=RuntimeError("connection refused")))
    with pytest.raises(EmbeddingError, match="connection refused"):
        asyncio.run(llm.embed(["hello"]))
    assert llm.embed_batches == 0


def test_http_error_surfaces_status() -> None:
    llm, _ = _make_client(_FakeEmbeddingsService(http_error=502))
    with pytest.raises(EmbeddingError, match="HTTP 502"):
        asyncio.run(llm.embed(["hello"]))
    assert llm.embed_batches == 0


def test_missing_vector_row_is_rejected() -> None:
    llm, _ = _make_client(_FakeEmbeddingsService(drop_index=1))
    with pytest.raises(EmbeddingError, match="returned 1 vectors for 2 inputs"):
        asyncio.run(llm.embed(["a", "b"]))


def test_batch_size_one_forces_one_call_per_text() -> None:
    service = _FakeEmbeddingsService()
    llm, _ = _make_client(service, embed_batch_size=1)
    asyncio.run(llm.embed(["a", "b", "c"]))
    assert [len(c) for c in service.calls] == [1, 1, 1]


def test_token_budget_limits_texts_per_request() -> None:
    """~2000-char chunks must not stack up past aipi's ~1024-token cap."""
    service = _FakeEmbeddingsService()
    llm, _ = _make_client(service, embed_batch_size=16)  # high count cap
    texts = ["x" * 2000 for _ in range(4)]
    vecs = asyncio.run(llm.embed(texts))
    # 2000 + 2000 chars > 3600-char (≈900-token) budget ⇒ one chunk per request
    assert [len(c) for c in service.calls] == [1, 1, 1, 1]
    assert len(vecs) == 4


def test_small_chunks_pack_up_to_count_cap() -> None:
    service = _FakeEmbeddingsService()
    llm, _ = _make_client(service, embed_batch_size=4)  # count cap binds
    texts = ["short text" for _ in range(9)]
    vecs = asyncio.run(llm.embed(texts))
    assert [len(c) for c in service.calls] == [4, 4, 1]
    assert len(vecs) == 9


def test_too_large_response_halves_batch_until_it_fits() -> None:
    """The tokenizer estimate can be wrong for dense content — the client
    must halve an over-large request and preserve order."""
    service = _FakeEmbeddingsService(too_large_min=3)
    llm, http = _make_client(service, embed_batch_size=16)  # all 8 fit one request
    texts = [f"t{i}" for i in range(8)]
    vecs = asyncio.run(llm.embed(texts))
    # http.bodies sees every request (including the rejected ones); the
    # left-half recursion completes before the right half starts.
    assert [len(b["input"]) for b in http.bodies] == [8, 4, 2, 2, 4, 2, 2]
    assert len(vecs) == 8
    assert all(len(v) == 768 for v in vecs)


def test_single_oversized_text_fails_actionably() -> None:
    service = _FakeEmbeddingsService(too_large_min=1)
    llm, _ = _make_client(service)
    with pytest.raises(EmbeddingError, match="token cap"):
        asyncio.run(llm.embed(["x" * 3000]))
    assert llm.embed_batches == 0


# ---------- chat streaming (phase 03) ----------


def _chunk(content: str | None = "text", empty: bool = False):
    """One fake ChatCompletionChunk (``choices[].delta.content`` shape)."""
    if empty:
        return SimpleNamespace(choices=[])
    return SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content=content))])


class _FakeChatStream:
    def __init__(self, chunks: list) -> None:
        self._chunks = list(chunks)

    def __aiter__(self):
        self._i = 0
        return self

    async def __anext__(self):
        if self._i >= len(self._chunks):
            raise StopAsyncIteration
        chunk = self._chunks[self._i]
        self._i += 1
        return chunk


class _FakeCompletions:
    def __init__(self, chunks: list | None = None, fail: Exception | None = None) -> None:
        self.chunks = chunks or []
        self.fail = fail
        self.kwargs: dict | None = None

    async def create(self, **kwargs) -> _FakeChatStream:
        self.kwargs = kwargs
        if self.fail is not None:
            raise self.fail
        return _FakeChatStream(self.chunks)


def _make_stream_client(
    chunks: list | None = None, fail: Exception | None = None
) -> tuple[LLMClient, _FakeCompletions]:
    completions = _FakeCompletions(chunks, fail)
    fake_openai = SimpleNamespace(chat=SimpleNamespace(completions=completions))
    llm = LLMClient(_settings())
    llm._client = fake_openai  # pyright: ignore[reportAttributeAccessIssue]
    return llm, completions


async def _collect(llm: LLMClient, messages: list[dict[str, str]]) -> list[str]:
    return [p async for p in llm.chat_stream(messages)]


def test_chat_stream_yields_deltas_in_order() -> None:
    llm, completions = _make_stream_client(
        [_chunk("Hey "), _chunk("you've "), _chunk("got this! 🧠")]
    )
    pieces = asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))
    assert pieces == ["Hey ", "you've ", "got this! 🧠"]


def test_chat_stream_uses_locked_generation_params() -> None:
    llm, completions = _make_stream_client([_chunk("x")])
    messages = [{"role": "system", "content": "s"}, {"role": "user", "content": "u"}]
    asyncio.run(_collect(llm, messages))
    assert completions.kwargs is not None
    assert completions.kwargs["model"] == "turbo"
    assert completions.kwargs["stream"] is True
    assert completions.kwargs["temperature"] == 0.4
    assert completions.kwargs["max_tokens"] == 700
    assert completions.kwargs["messages"] == messages


def test_chat_stream_skips_empty_deltas_and_choiceless_chunks() -> None:
    llm, _ = _make_stream_client([_chunk("a"), _chunk(empty=True), _chunk(None), _chunk("b")])
    assert asyncio.run(_collect(llm, [{"role": "user", "content": "q"}])) == ["a", "b"]


def test_chat_stream_wraps_failures_as_llm_error() -> None:
    llm, _ = _make_stream_client(fail=RuntimeError("connection reset by peer"))

    async def drain() -> None:
        async for _ in llm.chat_stream([{"role": "user", "content": "q"}]):
            pass

    with pytest.raises(LLMError, match="connection reset by peer"):
        asyncio.run(drain())


def test_chat_stream_llm_error_passes_through_unwrapped() -> None:
    llm, _ = _make_stream_client(fail=LLMError("already wrapped"))
    with pytest.raises(LLMError, match="already wrapped"):
        asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))