431 lines
15 KiB
Python
431 lines
15 KiB
Python
"""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,
|
|
StreamPiece,
|
|
)
|
|
|
|
|
|
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, reasoning: str | None = None
|
|
):
|
|
"""One fake ChatCompletionChunk (``choices[].delta`` shape).
|
|
|
|
``reasoning_content`` is present on the delta only when *reasoning*
|
|
is not None — mirroring the real wire, where the field exists only
|
|
when the model sends it.
|
|
"""
|
|
if empty:
|
|
return SimpleNamespace(choices=[])
|
|
delta: SimpleNamespace = SimpleNamespace(content=content)
|
|
if reasoning is not None:
|
|
delta.reasoning_content = reasoning
|
|
return SimpleNamespace(choices=[SimpleNamespace(delta=delta)])
|
|
|
|
|
|
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,
|
|
**settings_kwargs: Any,
|
|
) -> tuple[LLMClient, _FakeCompletions]:
|
|
completions = _FakeCompletions(chunks, fail)
|
|
fake_openai = SimpleNamespace(chat=SimpleNamespace(completions=completions))
|
|
llm = LLMClient(_settings(**settings_kwargs))
|
|
llm._client = fake_openai # pyright: ignore[reportAttributeAccessIssue]
|
|
return llm, completions
|
|
|
|
|
|
async def _collect(llm: LLMClient, messages: list[dict[str, str]]) -> list[StreamPiece]:
|
|
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"}]))
|
|
# Content-only chunks yield content pieces in wire order.
|
|
assert [(p.kind, p.text) for p in pieces] == [
|
|
("content", "Hey "),
|
|
("content", "you've "),
|
|
("content", "got this! 🧠"),
|
|
]
|
|
assert all(isinstance(p, StreamPiece) for p in pieces)
|
|
|
|
|
|
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
|
|
# Phase 11: the old hard 700-token cap is gone — answers may run up to
|
|
# BOR_MAX_OUTPUT_TOKENS (default 32 768) so they are not cut off.
|
|
assert completions.kwargs["max_tokens"] == 32_768
|
|
assert completions.kwargs["messages"] == messages
|
|
|
|
|
|
def test_chat_stream_max_tokens_comes_from_settings() -> None:
|
|
"""The output cap is operator-configurable, not a client constant."""
|
|
llm, completions = _make_stream_client(
|
|
[_chunk("x")], max_output_tokens=1234 # pyright: ignore[reportArgumentType]
|
|
)
|
|
asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))
|
|
assert completions.kwargs is not None
|
|
assert completions.kwargs["max_tokens"] == 1234
|
|
|
|
|
|
def test_chat_stream_skips_empty_deltas_and_choiceless_chunks() -> None:
|
|
llm, _ = _make_stream_client([_chunk("a"), _chunk(empty=True), _chunk(None), _chunk("b")])
|
|
pieces = asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))
|
|
assert [(p.kind, p.text) for p in pieces] == [("content", "a"), ("content", "b")]
|
|
|
|
|
|
def test_chat_stream_maps_reasoning_content_to_thinking_pieces() -> None:
|
|
"""The verified aipi wire field (``delta.reasoning_content``) maps to
|
|
``thinking`` pieces; content chunks are untouched by the presence of
|
|
reasoning elsewhere in the stream."""
|
|
llm, _ = _make_stream_client(
|
|
[
|
|
_chunk("", reasoning="Step 1: parse the question."),
|
|
_chunk("", reasoning="Step 2: cite the doc."),
|
|
_chunk("Talos."),
|
|
]
|
|
)
|
|
pieces = asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))
|
|
assert [(p.kind, p.text) for p in pieces] == [
|
|
("thinking", "Step 1: parse the question."),
|
|
("thinking", "Step 2: cite the doc."),
|
|
("content", "Talos."),
|
|
]
|
|
|
|
|
|
def test_chat_stream_falls_back_to_reasoning_field() -> None:
|
|
"""Future-proofing: a bare ``delta.reasoning`` field (no
|
|
``reasoning_content``) is picked up by the fallback getattr."""
|
|
chunk = SimpleNamespace(
|
|
choices=[
|
|
SimpleNamespace(delta=SimpleNamespace(content="ans", reasoning="why not"))
|
|
]
|
|
)
|
|
llm, _ = _make_stream_client([chunk])
|
|
pieces = asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))
|
|
assert [(p.kind, p.text) for p in pieces] == [
|
|
("thinking", "why not"),
|
|
("content", "ans"),
|
|
]
|
|
|
|
|
|
def test_chat_stream_thinking_yields_before_content_in_chunk() -> None:
|
|
"""One chunk carrying both fields yields the thinking piece first."""
|
|
llm, _ = _make_stream_client([_chunk("answer", reasoning="hmm")])
|
|
pieces = asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))
|
|
assert [(p.kind, p.text) for p in pieces] == [
|
|
("thinking", "hmm"),
|
|
("content", "answer"),
|
|
]
|
|
|
|
|
|
def test_chat_stream_interleaved_thinking_and_content_order_preserved() -> None:
|
|
"""The piece sequence must match the chunk sequence exactly — a late
|
|
or interleaved thinking chunk is emitted at its wire position."""
|
|
llm, _ = _make_stream_client(
|
|
[
|
|
_chunk("", reasoning="t1"),
|
|
_chunk("c1"),
|
|
_chunk("", reasoning="t2"),
|
|
_chunk("c2"),
|
|
]
|
|
)
|
|
pieces = asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))
|
|
assert [(p.kind, p.text) for p in pieces] == [
|
|
("thinking", "t1"),
|
|
("content", "c1"),
|
|
("thinking", "t2"),
|
|
("content", "c2"),
|
|
]
|
|
|
|
|
|
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"}]))
|