1281 lines
46 KiB
Python
1281 lines
46 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 collections.abc import AsyncGenerator
|
|
from types import SimpleNamespace
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
from app.config import Settings
|
|
from app.rag.llm import (
|
|
EmbeddingDimensionError,
|
|
EmbeddingError,
|
|
LLMClient,
|
|
LLMError,
|
|
RetryPiece,
|
|
StreamPiece,
|
|
ToolCallPiece,
|
|
chat_stream_retried,
|
|
)
|
|
from app.rag.scaffolding import ScaffoldingFilter
|
|
|
|
|
|
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 _tool_call(
|
|
index: int,
|
|
id: str | None = None,
|
|
name: str | None = None,
|
|
arguments: str | None = None,
|
|
):
|
|
"""One fake ``delta.tool_calls[]`` partial (openai SDK shape, phase 37).
|
|
|
|
``function`` is None when neither *name* nor *arguments* is given —
|
|
mirroring the real wire, where id-only fragments carry no function.
|
|
"""
|
|
fn = None
|
|
if name is not None or arguments is not None:
|
|
fn = SimpleNamespace(name=name, arguments=arguments)
|
|
return SimpleNamespace(index=index, id=id, function=fn)
|
|
|
|
|
|
def _chunk(
|
|
content: str | None = "text",
|
|
empty: bool = False,
|
|
reasoning: str | None = None,
|
|
tool_calls: list | None = None,
|
|
finish_reason: str | None = None,
|
|
):
|
|
"""One fake ChatCompletionChunk (``choices[].delta`` shape).
|
|
|
|
``reasoning_content``, ``tool_calls`` and ``finish_reason`` are
|
|
present only when provided — mirroring the real wire, where the
|
|
fields exist only when the model sends them.
|
|
"""
|
|
if empty:
|
|
return SimpleNamespace(choices=[])
|
|
delta: SimpleNamespace = SimpleNamespace(content=content)
|
|
if reasoning is not None:
|
|
delta.reasoning_content = reasoning
|
|
if tool_calls is not None:
|
|
delta.tool_calls = tool_calls
|
|
choice = SimpleNamespace(delta=delta)
|
|
if finish_reason is not None:
|
|
choice.finish_reason = finish_reason
|
|
return SimpleNamespace(choices=[choice])
|
|
|
|
|
|
class _FakeChatStream:
|
|
def __init__(self, chunks: list) -> None:
|
|
self._chunks = list(chunks)
|
|
self.close_calls = 0
|
|
|
|
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
|
|
|
|
async def close(self) -> None:
|
|
"""The openai SDK ``AsyncStream.close()`` (phase 48): ``chat_stream``
|
|
awaits it on every exit after a successful ``create()``."""
|
|
self.close_calls += 1
|
|
|
|
|
|
class _FakeCompletion:
|
|
"""One fake non-streaming ChatCompletion (``choices[].message`` shape).
|
|
|
|
``content=None`` mirrors the real wire where the field can be absent or
|
|
empty (reasoning-only replies, provider quirks).
|
|
"""
|
|
|
|
def __init__(self, content: str | None, empty_choices: bool = False) -> None:
|
|
if empty_choices:
|
|
self.choices = []
|
|
else:
|
|
self.choices = [SimpleNamespace(message=SimpleNamespace(content=content))]
|
|
|
|
|
|
class _FakeCompletions:
|
|
def __init__(
|
|
self,
|
|
chunks: list | None = None,
|
|
fail: Exception | None = None,
|
|
completion: _FakeCompletion | None = None,
|
|
) -> None:
|
|
self.chunks = chunks or []
|
|
self.fail = fail
|
|
self.completion = completion
|
|
self.kwargs: dict | None = None
|
|
self.chat_kwargs: dict | None = None
|
|
#: Every SDK-shaped stream handed out — teardown tests assert the
|
|
#: phase-48 ``close()`` on them (phase 71 task 02: with/without
|
|
#: a filter, the teardown path is the same object).
|
|
self.streams: list[_FakeChatStream] = []
|
|
|
|
async def create(self, **kwargs) -> _FakeChatStream | _FakeCompletion:
|
|
self.kwargs = kwargs
|
|
if self.fail is not None:
|
|
raise self.fail
|
|
if kwargs.get("stream"):
|
|
stream = _FakeChatStream(self.chunks)
|
|
self.streams.append(stream)
|
|
return stream
|
|
self.chat_kwargs = kwargs
|
|
assert self.completion is not None
|
|
return self.completion
|
|
|
|
|
|
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]:
|
|
"""Collect pieces from a tools-less stream (phase 37 task 02, test (a):
|
|
without tools, no ToolCallPiece can appear)."""
|
|
pieces = [p async for p in llm.chat_stream(messages)]
|
|
assert all(isinstance(p, StreamPiece) for p in pieces)
|
|
return cast("list[StreamPiece]", pieces)
|
|
|
|
|
|
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
|
|
# Phase 37: no tools passed ⇒ no `tools` key at all (byte-identical
|
|
# request to pre-phase-37).
|
|
assert "tools" not in completions.kwargs
|
|
|
|
|
|
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"}]))
|
|
|
|
|
|
# ---------- tool-call streaming (phase 37, task 02) ----------
|
|
|
|
#: The agent's tool list (phase 70: the harness-aligned surface) — the
|
|
#: exact wire shape AGENT_TOOLS passes through (the names are whatever
|
|
#: the caller's tools list names).
|
|
_AGENT_TOOLS: list[dict[str, Any]] = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "ls",
|
|
"description": "List the indexed documents.",
|
|
"parameters": {"type": "object", "properties": {}, "required": []},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "read",
|
|
"description": "Add one indexed document's full text to the context.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"path": {"type": "string"}},
|
|
"required": ["path"],
|
|
},
|
|
},
|
|
},
|
|
]
|
|
|
|
|
|
def _collect_with_tools(
|
|
llm: LLMClient, messages: list[dict[str, str]], tools: list[dict[str, Any]]
|
|
) -> list[StreamPiece | ToolCallPiece]:
|
|
async def run() -> list[StreamPiece | ToolCallPiece]:
|
|
return [p async for p in llm.chat_stream(messages, tools=tools)]
|
|
|
|
return asyncio.run(run())
|
|
|
|
|
|
def test_chat_stream_passes_tools_when_given() -> None:
|
|
"""(e) A non-None tools list is forwarded verbatim to create()."""
|
|
llm, completions = _make_stream_client([_chunk("ok")], llm_chat_model="turbo")
|
|
_collect_with_tools(llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS)
|
|
assert completions.kwargs is not None
|
|
assert completions.kwargs["tools"] == _AGENT_TOOLS
|
|
|
|
|
|
def test_chat_stream_accumulates_tool_call_across_chunk_partials() -> None:
|
|
"""(b) name on the first partial, arguments in fragments — merged into
|
|
one ToolCallPiece with the concatenated JSON, at finish_reason."""
|
|
llm, _ = _make_stream_client(
|
|
[
|
|
_chunk(
|
|
None,
|
|
tool_calls=[
|
|
_tool_call(
|
|
0,
|
|
id="call_abc",
|
|
name="read",
|
|
arguments='{"path": "Homelab/ku',
|
|
)
|
|
],
|
|
),
|
|
_chunk(None, tool_calls=[_tool_call(0, arguments='bernetes.md"}')]),
|
|
_chunk(None, finish_reason="tool_calls"),
|
|
]
|
|
)
|
|
pieces = _collect_with_tools(
|
|
llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS
|
|
)
|
|
assert pieces == [
|
|
ToolCallPiece(
|
|
id="call_abc",
|
|
name="read",
|
|
arguments={"path": "Homelab/kubernetes.md"},
|
|
)
|
|
]
|
|
|
|
|
|
def test_chat_stream_two_tool_calls_yielded_in_index_order() -> None:
|
|
"""(c) Indices 0 and 1, interleaved partials (index 1 seen first) —
|
|
both calls, in index order, each merged from its own fragments."""
|
|
llm, _ = _make_stream_client(
|
|
[
|
|
_chunk(
|
|
None,
|
|
tool_calls=[
|
|
_tool_call(1, id="call_b", name="read", arguments='{"pa')
|
|
],
|
|
),
|
|
_chunk(
|
|
None,
|
|
tool_calls=[
|
|
_tool_call(0, id="call_a", name="ls"),
|
|
_tool_call(1, arguments='th": "Homelab/a.md"}')
|
|
],
|
|
),
|
|
_chunk(None, finish_reason="tool_calls"),
|
|
]
|
|
)
|
|
pieces = _collect_with_tools(
|
|
llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS
|
|
)
|
|
assert pieces == [
|
|
ToolCallPiece(id="call_a", name="ls", arguments={}),
|
|
ToolCallPiece(
|
|
id="call_b",
|
|
name="read",
|
|
arguments={"path": "Homelab/a.md"},
|
|
),
|
|
]
|
|
|
|
|
|
def test_chat_stream_tool_calls_yielded_at_stream_end_without_finish_reason() -> None:
|
|
"""The spec's other emission point: stream ends without a
|
|
finish_reason="tool_calls" chunk — pieces still materialize."""
|
|
llm, _ = _make_stream_client(
|
|
[
|
|
_chunk(
|
|
None,
|
|
tool_calls=[_tool_call(0, id="call_z", name="ls")],
|
|
)
|
|
]
|
|
)
|
|
pieces = _collect_with_tools(
|
|
llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS
|
|
)
|
|
assert pieces == [ToolCallPiece(id="call_z", name="ls", arguments={})]
|
|
|
|
|
|
def test_chat_stream_synthesizes_call_id_when_absent() -> None:
|
|
"""Wire never carried the call id ⇒ synthesized "call_<index>"."""
|
|
llm, _ = _make_stream_client(
|
|
[
|
|
_chunk(None, tool_calls=[_tool_call(2, name="read", arguments="{}")]),
|
|
_chunk(None, finish_reason="tool_calls"),
|
|
]
|
|
)
|
|
pieces = _collect_with_tools(
|
|
llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS
|
|
)
|
|
assert pieces == [
|
|
ToolCallPiece(
|
|
id="call_2",
|
|
name="read",
|
|
arguments={},
|
|
)
|
|
]
|
|
|
|
|
|
def test_chat_stream_null_arguments_become_empty_dict() -> None:
|
|
"""JSON "null" (and, by the same branch, absent arguments) ⇒ {}."""
|
|
llm, _ = _make_stream_client(
|
|
[
|
|
_chunk(
|
|
None,
|
|
tool_calls=[
|
|
_tool_call(0, id="call_n", name="ls", arguments="null")
|
|
],
|
|
),
|
|
_chunk(None, finish_reason="tool_calls"),
|
|
]
|
|
)
|
|
pieces = _collect_with_tools(
|
|
llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS
|
|
)
|
|
assert pieces == [ToolCallPiece(id="call_n", name="ls", arguments={})]
|
|
|
|
|
|
def test_chat_stream_malformed_tool_arguments_raise_llm_error() -> None:
|
|
"""(d) A silently dropped tool call would corrupt the loop — malformed
|
|
arguments JSON must fail loudly."""
|
|
llm, _ = _make_stream_client(
|
|
[
|
|
_chunk(
|
|
None,
|
|
tool_calls=[
|
|
_tool_call(
|
|
0,
|
|
id="call_x",
|
|
name="read",
|
|
arguments='{"path": "Homelab",',
|
|
)
|
|
],
|
|
),
|
|
_chunk(None, finish_reason="tool_calls"),
|
|
]
|
|
)
|
|
|
|
async def drain() -> None:
|
|
async for _ in llm.chat_stream(
|
|
[{"role": "user", "content": "q"}], tools=_AGENT_TOOLS
|
|
):
|
|
pass
|
|
|
|
with pytest.raises(LLMError, match="malformed tool-call arguments"):
|
|
asyncio.run(drain())
|
|
|
|
|
|
def test_chat_stream_non_object_tool_arguments_raise_llm_error() -> None:
|
|
"""The OpenAI contract says arguments is a JSON *object* — a bare array
|
|
is malformed too."""
|
|
llm, _ = _make_stream_client(
|
|
[
|
|
_chunk(
|
|
None,
|
|
tool_calls=[
|
|
_tool_call(0, id="call_y", name="grep", arguments='[1, 2]')
|
|
],
|
|
),
|
|
_chunk(None, finish_reason="tool_calls"),
|
|
]
|
|
)
|
|
|
|
async def drain() -> None:
|
|
async for _ in llm.chat_stream(
|
|
[{"role": "user", "content": "q"}], tools=_AGENT_TOOLS
|
|
):
|
|
pass
|
|
|
|
with pytest.raises(LLMError, match="non-object tool-call arguments"):
|
|
asyncio.run(drain())
|
|
|
|
|
|
# ---------- scaffolding filter integration (phase 71, task 02) ----------
|
|
|
|
#: The incident's raw span (2026-09-03) — the filter's reason to exist.
|
|
_INCIDENT_SPAN = "<|tool_call_start|>[read(path='/homelab/backup-notes.md')]<|tool_call_end|>"
|
|
|
|
|
|
def _collect_filtered(
|
|
llm: LLMClient,
|
|
messages: list[dict[str, str]],
|
|
scaffolding: ScaffoldingFilter | None,
|
|
) -> list[StreamPiece]:
|
|
"""Collect pieces from a tools-less filtered stream (phase 71): without
|
|
tools, no ToolCallPiece can appear (the phase-37 contract)."""
|
|
async def run() -> list[StreamPiece]:
|
|
pieces = [p async for p in llm.chat_stream(messages, scaffolding=scaffolding)]
|
|
assert all(isinstance(p, StreamPiece) for p in pieces)
|
|
return cast("list[StreamPiece]", pieces)
|
|
|
|
return asyncio.run(run())
|
|
|
|
|
|
def test_chat_stream_with_filter_strips_span_mid_stream() -> None:
|
|
"""A span mid-stream never reaches the pieces: the clean remainder
|
|
flows, ``stripped_chars`` is exact, and no piece carries scaffolding."""
|
|
f = ScaffoldingFilter()
|
|
llm, _ = _make_stream_client(
|
|
[_chunk("Hello "), _chunk(_INCIDENT_SPAN + " world"), _chunk("!")]
|
|
)
|
|
pieces = _collect_filtered(llm, _RETRY_MSGS, f)
|
|
assert [(p.kind, p.text) for p in pieces] == [
|
|
("content", "Hello "),
|
|
("content", " world"),
|
|
("content", "!"),
|
|
]
|
|
assert f.stripped_chars == len(_INCIDENT_SPAN)
|
|
assert all("<|" not in p.text for p in pieces if p.kind == "content")
|
|
|
|
|
|
def test_chat_stream_with_filter_span_split_across_chunks() -> None:
|
|
"""A span split across two chunks never emits a partial marker —
|
|
nothing leaks until the span completes, then the clean text on both
|
|
sides flows and the whole span counts as stripped."""
|
|
span = _INCIDENT_SPAN
|
|
cut = len("<|tool_call_start|>") # split right after the start token
|
|
f = ScaffoldingFilter()
|
|
llm, _ = _make_stream_client([_chunk("A " + span[:cut]), _chunk(span[cut:] + " B")])
|
|
pieces = _collect_filtered(llm, _RETRY_MSGS, f)
|
|
assert [(p.kind, p.text) for p in pieces] == [
|
|
("content", "A "),
|
|
("content", " B"),
|
|
]
|
|
assert f.stripped_chars == len(span)
|
|
|
|
|
|
def test_chat_stream_with_filter_yields_nothing_for_pure_scaffolding() -> None:
|
|
"""A content stream of pure scaffolding yields ZERO content pieces —
|
|
an empty clean result yields nothing (no empty ``delta`` frames)."""
|
|
f = ScaffoldingFilter()
|
|
llm, _ = _make_stream_client(
|
|
[_chunk(_INCIDENT_SPAN), _chunk("<|tool_calls|>")]
|
|
)
|
|
pieces = _collect_filtered(llm, _RETRY_MSGS, f)
|
|
assert pieces == []
|
|
assert f.stripped_chars == len(_INCIDENT_SPAN) + len("<|tool_calls|>")
|
|
|
|
|
|
def test_chat_stream_with_filter_leaves_thinking_raw() -> None:
|
|
"""Locked (phase 71): the scratchpad stays raw — a span inside
|
|
``reasoning_content`` is yielded verbatim and never counts as
|
|
stripped."""
|
|
f = ScaffoldingFilter()
|
|
llm, _ = _make_stream_client(
|
|
[_chunk("", reasoning=_INCIDENT_SPAN), _chunk("ok")]
|
|
)
|
|
pieces = _collect_filtered(llm, _RETRY_MSGS, f)
|
|
assert [(p.kind, p.text) for p in pieces] == [
|
|
("thinking", _INCIDENT_SPAN),
|
|
("content", "ok"),
|
|
]
|
|
assert f.stripped_chars == 0
|
|
|
|
|
|
def test_chat_stream_without_filter_is_the_raw_path() -> None:
|
|
"""``scaffolding=None`` (the default, and the explicit opt-out): a span
|
|
in content is yielded verbatim — byte-identical to the pre-phase-71
|
|
raw path for callers that do not pass a filter."""
|
|
llm, _ = _make_stream_client([_chunk(_INCIDENT_SPAN)])
|
|
explicit_none = _collect_filtered(llm, _RETRY_MSGS, None)
|
|
default = asyncio.run(_collect(llm, _RETRY_MSGS))
|
|
assert explicit_none == default == [StreamPiece("content", _INCIDENT_SPAN)]
|
|
|
|
|
|
def test_chat_stream_flushed_tail_precedes_tool_call_pieces() -> None:
|
|
"""Content-before-tools wire convention (phase 71): a stream that ends
|
|
with a held filter tail (a partial marker at EOF — flushed as-is) +
|
|
tool_calls deltas yields the flushed tail content piece BEFORE the
|
|
materialized ``ToolCallPiece``."""
|
|
f = ScaffoldingFilter()
|
|
llm, _ = _make_stream_client(
|
|
[
|
|
_chunk("done <|tool_call_st"), # held: a live prefix of the start token
|
|
_chunk(
|
|
None,
|
|
tool_calls=[_tool_call(0, id="call_t", name="ls")],
|
|
finish_reason="tool_calls",
|
|
),
|
|
]
|
|
)
|
|
|
|
async def run() -> list[StreamPiece | ToolCallPiece]:
|
|
return [
|
|
p
|
|
async for p in llm.chat_stream(_RETRY_MSGS, tools=_AGENT_TOOLS, scaffolding=f)
|
|
]
|
|
|
|
pieces = asyncio.run(run())
|
|
assert pieces == [
|
|
StreamPiece("content", "done "),
|
|
StreamPiece("content", "<|tool_call_st"),
|
|
ToolCallPiece(id="call_t", name="ls", arguments={}),
|
|
]
|
|
|
|
|
|
def test_chat_stream_abandon_with_filter_closes_stream() -> None:
|
|
"""Phase-48 teardown is independent of the phase-71 filter: a consumer
|
|
abandon mid-filter still closes the endpoint stream exactly once."""
|
|
f = ScaffoldingFilter()
|
|
llm, completions = _make_stream_client([_chunk(f"piece {i} ") for i in range(1, 6)])
|
|
|
|
async def run() -> None:
|
|
gen = llm.chat_stream(_RETRY_MSGS, scaffolding=f)
|
|
first = await gen.__anext__()
|
|
assert isinstance(first, StreamPiece)
|
|
assert first.text == "piece 1 "
|
|
await gen.aclose() # the consumer stops after the first piece
|
|
|
|
asyncio.run(run())
|
|
assert completions.streams[0].close_calls == 1
|
|
|
|
|
|
# ---------- one-shot chat: LLMClient.chat (phase 30, task 01) ----------
|
|
|
|
|
|
def _make_chat_client(
|
|
completion: _FakeCompletion | None = None,
|
|
fail: Exception | None = None,
|
|
**settings_kwargs: Any,
|
|
) -> tuple[LLMClient, _FakeCompletions]:
|
|
completions = _FakeCompletions(fail=fail, completion=completion)
|
|
fake_openai = SimpleNamespace(chat=SimpleNamespace(completions=completions))
|
|
llm = LLMClient(_settings(**settings_kwargs))
|
|
llm._client = fake_openai # pyright: ignore[reportAttributeAccessIssue]
|
|
return llm, completions
|
|
|
|
|
|
def test_chat_returns_trimmed_content_with_locked_params() -> None:
|
|
"""Default model is ``lite`` (BOR_LLM_SUMMARY_MODEL), non-streaming,
|
|
low temperature, fixed 2048-token budget — summaries are short."""
|
|
llm, completions = _make_chat_client(_FakeCompletion(" Summary text.\n"))
|
|
messages = [{"role": "system", "content": "s"}, {"role": "user", "content": "u"}]
|
|
out = asyncio.run(llm.chat(messages))
|
|
assert out == "Summary text."
|
|
assert completions.chat_kwargs is not None
|
|
assert completions.chat_kwargs["model"] == "lite"
|
|
assert completions.chat_kwargs["stream"] is False
|
|
assert completions.chat_kwargs["temperature"] == 0.2
|
|
assert completions.chat_kwargs["max_tokens"] == 2048
|
|
assert completions.chat_kwargs["messages"] == messages
|
|
|
|
|
|
def test_chat_default_model_comes_from_llm_summary_model_setting() -> None:
|
|
llm, completions = _make_chat_client(
|
|
_FakeCompletion("x"), llm_summary_model="tiny"
|
|
)
|
|
asyncio.run(llm.chat([{"role": "user", "content": "q"}]))
|
|
assert completions.chat_kwargs is not None
|
|
assert completions.chat_kwargs["model"] == "tiny"
|
|
|
|
|
|
def test_chat_explicit_model_overrides_the_default() -> None:
|
|
llm, completions = _make_chat_client(
|
|
_FakeCompletion("x"), llm_summary_model="tiny"
|
|
)
|
|
asyncio.run(llm.chat([{"role": "user", "content": "q"}], model="special"))
|
|
assert completions.chat_kwargs is not None
|
|
assert completions.chat_kwargs["model"] == "special"
|
|
|
|
|
|
def test_chat_transport_failure_wrapped_as_llm_error_with_base_url() -> None:
|
|
"""HTTP/transport failures (incl. >=400 surfaced by the SDK) are wrapped
|
|
with the base URL in the message — same style as chat_stream."""
|
|
llm, _ = _make_chat_client(fail=RuntimeError("HTTP 502 Bad Gateway"))
|
|
with pytest.raises(LLMError, match="HTTP 502") as exc:
|
|
asyncio.run(llm.chat([{"role": "user", "content": "q"}]))
|
|
assert "aipi.reeseapps.com" in str(exc.value)
|
|
|
|
|
|
def test_chat_llm_error_passes_through_unwrapped() -> None:
|
|
llm, _ = _make_chat_client(fail=LLMError("already wrapped"))
|
|
with pytest.raises(LLMError, match="already wrapped"):
|
|
asyncio.run(llm.chat([{"role": "user", "content": "q"}]))
|
|
|
|
|
|
def test_chat_empty_choices_raises_llm_error() -> None:
|
|
llm, _ = _make_chat_client(_FakeCompletion(None, empty_choices=True))
|
|
with pytest.raises(LLMError, match="no choices"):
|
|
asyncio.run(llm.chat([{"role": "user", "content": "q"}]))
|
|
|
|
|
|
def test_chat_missing_content_raises_llm_error() -> None:
|
|
"""A silent empty summary must never be stored — None content fails."""
|
|
llm, _ = _make_chat_client(_FakeCompletion(None))
|
|
with pytest.raises(LLMError, match="empty content"):
|
|
asyncio.run(llm.chat([{"role": "user", "content": "q"}]))
|
|
|
|
|
|
def test_chat_whitespace_only_content_raises_llm_error() -> None:
|
|
llm, _ = _make_chat_client(_FakeCompletion(" \n\t "))
|
|
with pytest.raises(LLMError, match="empty content"):
|
|
asyncio.run(llm.chat([{"role": "user", "content": "q"}]))
|
|
|
|
|
|
# ---------- chat_stream_retried (phase 67, task 01) ----------
|
|
|
|
_RETRY_MSGS: list[dict[str, str]] = [{"role": "user", "content": "q"}]
|
|
|
|
|
|
class _ScriptedClient(LLMClient):
|
|
"""An LLMClient whose ``chat_stream`` is scripted per attempt — no
|
|
endpoint. ``attempts`` scripts the Nth call: ``(pieces, error)`` — the
|
|
stream yields *pieces*, then raises *error* if not None (an error after
|
|
zero pieces = "the endpoint died before the first token"; after some
|
|
pieces = a mid-stream drop). Records every request's messages/tools
|
|
and every attempt's stream teardown (the phase-48 close analog)."""
|
|
|
|
def __init__(
|
|
self, attempts: list[tuple[list[StreamPiece | ToolCallPiece], Exception | None]]
|
|
) -> None:
|
|
super().__init__(_settings())
|
|
self.attempts = list(attempts)
|
|
self.request_args: list[
|
|
tuple[list[dict[str, str]], list[dict[str, Any]] | None]
|
|
] = []
|
|
#: Indices of attempts whose stream teardown has run.
|
|
self.closed: list[int] = []
|
|
#: The caller-owned filter each attempt's ``chat_stream`` received
|
|
#: (phase 71 task 02 — the retry primitive forwards it).
|
|
self.scaffoldings: list[ScaffoldingFilter | None] = []
|
|
|
|
def chat_stream(
|
|
self,
|
|
messages: list[dict[str, str]],
|
|
tools: list[dict[str, Any]] | None = None,
|
|
scaffolding: ScaffoldingFilter | None = None,
|
|
) -> AsyncGenerator[StreamPiece | ToolCallPiece, None]:
|
|
index = len(self.request_args)
|
|
pieces, error = (
|
|
self.attempts[index]
|
|
if index < len(self.attempts)
|
|
else ([], LLMError("script exhausted"))
|
|
)
|
|
self.request_args.append(
|
|
(list(messages), list(tools) if tools is not None else None)
|
|
)
|
|
self.scaffoldings.append(scaffolding)
|
|
return self._attempt(index, pieces, error, scaffolding)
|
|
|
|
async def _attempt(
|
|
self,
|
|
index: int,
|
|
pieces: list[StreamPiece | ToolCallPiece],
|
|
error: Exception | None,
|
|
scaffolding: ScaffoldingFilter | None,
|
|
) -> AsyncGenerator[StreamPiece | ToolCallPiece, None]:
|
|
try:
|
|
for piece in pieces:
|
|
if (
|
|
scaffolding is not None
|
|
and isinstance(piece, StreamPiece)
|
|
and piece.kind == "content"
|
|
):
|
|
# Emulate the real contract (phase 71): content feeds
|
|
# the filter, an empty clean result yields nothing.
|
|
cleaned = scaffolding.feed(piece.text)
|
|
if cleaned:
|
|
yield StreamPiece("content", cleaned)
|
|
else:
|
|
yield piece
|
|
if error is not None:
|
|
raise error
|
|
finally:
|
|
self.closed.append(index)
|
|
|
|
|
|
def _record_sleeps(monkeypatch: pytest.MonkeyPatch) -> list[float]:
|
|
"""Monkeypatch ``asyncio.sleep`` (what ``chat_stream_retried`` awaits
|
|
for the flat pre-retry delay) and record every awaited delay."""
|
|
sleeps: list[float] = []
|
|
|
|
async def fake_sleep(seconds: float) -> None:
|
|
sleeps.append(seconds)
|
|
|
|
monkeypatch.setattr(asyncio, "sleep", fake_sleep)
|
|
return sleeps
|
|
|
|
|
|
def _collect_retried(
|
|
client: _ScriptedClient,
|
|
messages: list[dict[str, str]],
|
|
*,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
retries: int,
|
|
delay: float,
|
|
scaffolding: ScaffoldingFilter | None = None,
|
|
) -> list[StreamPiece | ToolCallPiece | RetryPiece]:
|
|
async def run() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
|
|
return [
|
|
p
|
|
async for p in chat_stream_retried(
|
|
client,
|
|
messages,
|
|
tools=tools,
|
|
retries=retries,
|
|
delay=delay,
|
|
scaffolding=scaffolding,
|
|
)
|
|
]
|
|
|
|
return asyncio.run(run())
|
|
|
|
|
|
def test_retried_retries_a_dead_attempt_before_the_first_piece(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Failure on attempt 1, success on attempt 2 → [RetryPiece(2, N)]
|
|
(the attempt about to be tried, 1-based) then the answer pieces; the
|
|
request is restarted byte-identical and the flat delay is awaited
|
|
exactly once."""
|
|
answer: list[StreamPiece | ToolCallPiece] = [
|
|
StreamPiece("content", "A "),
|
|
StreamPiece("content", "B"),
|
|
]
|
|
client = _ScriptedClient([([], LLMError("connection refused")), (answer, None)])
|
|
sleeps = _record_sleeps(monkeypatch)
|
|
pieces = _collect_retried(client, _RETRY_MSGS, retries=3, delay=2.5)
|
|
assert pieces == [RetryPiece(2, 4), *answer]
|
|
assert len(client.request_args) == 2
|
|
# The restart is byte-identical: same messages, same (absent) tools.
|
|
assert client.request_args[0] == client.request_args[1]
|
|
assert client.request_args[0][1] is None
|
|
assert sleeps == [2.5]
|
|
|
|
|
|
def test_retried_exhaustion_yields_all_retries_then_raises(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""retries=2 → 3 attempts; each pre-first-piece failure yields a
|
|
RetryPiece naming the attempt about to be tried (attempts 2 and 3 of
|
|
3), the final failure raises the terminal LLMError, and no sleep
|
|
follows the last attempt."""
|
|
client = _ScriptedClient(
|
|
[([], LLMError("down 1")), ([], LLMError("down 2")), ([], LLMError("down 3"))]
|
|
)
|
|
sleeps = _record_sleeps(monkeypatch)
|
|
|
|
async def run() -> list[RetryPiece]:
|
|
out: list[RetryPiece] = []
|
|
with pytest.raises(LLMError, match="down 3"):
|
|
async for p in chat_stream_retried(
|
|
client, _RETRY_MSGS, retries=2, delay=0.5
|
|
):
|
|
assert isinstance(p, RetryPiece)
|
|
out.append(p)
|
|
return out
|
|
|
|
out = asyncio.run(run())
|
|
assert out == [RetryPiece(2, 3), RetryPiece(3, 3)]
|
|
assert len(client.request_args) == 3
|
|
assert sleeps == [0.5, 0.5]
|
|
|
|
|
|
def test_retried_failure_after_first_piece_is_terminal(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Locked A2: a piece has already flowed → the LLMError is re-raised
|
|
immediately — no RetryPiece, no sleep, no second call (a partial
|
|
answer is never redone)."""
|
|
client = _ScriptedClient(
|
|
[([StreamPiece("content", "partial ")], LLMError("mid-stream drop"))],
|
|
)
|
|
sleeps = _record_sleeps(monkeypatch)
|
|
|
|
async def run() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
|
|
out: list[StreamPiece | ToolCallPiece | RetryPiece] = []
|
|
with pytest.raises(LLMError, match="mid-stream drop"):
|
|
async for p in chat_stream_retried(
|
|
client, _RETRY_MSGS, retries=3, delay=5.0
|
|
):
|
|
out.append(p)
|
|
return out
|
|
|
|
out = asyncio.run(run())
|
|
assert out == [StreamPiece("content", "partial ")]
|
|
assert not any(isinstance(p, RetryPiece) for p in out)
|
|
assert len(client.request_args) == 1
|
|
assert sleeps == []
|
|
|
|
|
|
def test_retried_zero_retries_is_one_attempt_no_retry_pieces(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""The kill-switch path (retries=0): exactly one attempt, the error
|
|
propagates, no RetryPiece, no sleep — the pre-phase-67 behavior."""
|
|
client = _ScriptedClient([([], LLMError("connection refused"))])
|
|
sleeps = _record_sleeps(monkeypatch)
|
|
with pytest.raises(LLMError, match="connection refused"):
|
|
_collect_retried(client, _RETRY_MSGS, retries=0, delay=5.0)
|
|
assert len(client.request_args) == 1
|
|
assert sleeps == []
|
|
|
|
|
|
def test_retried_healthy_stream_is_untouched(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""No failure → exactly one attempt, every piece kind (thinking / tool
|
|
call / content) passes through unchanged, no RetryPiece, no sleep —
|
|
a healthy turn is byte-identical to the plain chat_stream."""
|
|
answer = [
|
|
StreamPiece("thinking", "hmm"),
|
|
ToolCallPiece(id="call_1", name="ls", arguments={}),
|
|
StreamPiece("content", "Talos."),
|
|
]
|
|
client = _ScriptedClient([(answer, None)])
|
|
sleeps = _record_sleeps(monkeypatch)
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"function": {"name": "ls", "parameters": {}},
|
|
}
|
|
]
|
|
pieces = _collect_retried(
|
|
client, _RETRY_MSGS, tools=tools, retries=3, delay=5.0
|
|
)
|
|
assert pieces == answer
|
|
assert client.request_args == [(_RETRY_MSGS, tools)]
|
|
assert sleeps == []
|
|
|
|
|
|
def test_retried_zero_delay_still_notifies(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""The e2e fast path (BOR_LLM_RETRY_DELAY=0): the RetryPiece is still
|
|
emitted and the (zero) sleep is still awaited."""
|
|
client = _ScriptedClient(
|
|
[([], LLMError("down")), ([StreamPiece("content", "ok")], None)]
|
|
)
|
|
sleeps = _record_sleeps(monkeypatch)
|
|
pieces = _collect_retried(client, _RETRY_MSGS, retries=1, delay=0)
|
|
assert pieces == [RetryPiece(2, 2), StreamPiece("content", "ok")]
|
|
assert sleeps == [0]
|
|
|
|
|
|
def test_retried_forwards_the_filter_to_every_attempt(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Phase 71: the caller-owned filter reaches EVERY attempt's
|
|
``chat_stream`` — the same object on the dead attempt and the
|
|
surviving one."""
|
|
client = _ScriptedClient(
|
|
[([], LLMError("down")), ([StreamPiece("content", "ok")], None)]
|
|
)
|
|
f = ScaffoldingFilter()
|
|
_record_sleeps(monkeypatch)
|
|
pieces = _collect_retried(client, _RETRY_MSGS, retries=1, delay=0, scaffolding=f)
|
|
assert pieces == [RetryPiece(2, 2), StreamPiece("content", "ok")]
|
|
assert client.scaffoldings == [f, f]
|
|
|
|
|
|
def test_retried_reuses_the_unfed_filter_after_a_dead_attempt(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Phase 71: reusing the SAME filter across attempts is safe by
|
|
construction — the dead attempt emitted no piece, so the filter was
|
|
never fed; the surviving attempt filters through it as if fresh."""
|
|
span = _INCIDENT_SPAN
|
|
client = _ScriptedClient(
|
|
[([], LLMError("down")), ([StreamPiece("content", span + " clean")], None)]
|
|
)
|
|
f = ScaffoldingFilter()
|
|
_record_sleeps(monkeypatch)
|
|
pieces = _collect_retried(client, _RETRY_MSGS, retries=1, delay=0, scaffolding=f)
|
|
assert pieces == [RetryPiece(2, 2), StreamPiece("content", " clean")]
|
|
assert f.stripped_chars == len(span)
|
|
assert client.scaffoldings == [f, f]
|
|
|
|
|
|
def test_retried_abandon_mid_attempt_closes_the_attempt_stream() -> None:
|
|
"""Consumer abandon at a mid-attempt piece (the stop-generation path,
|
|
phase 48): no exception leaks and the attempt's stream is torn down
|
|
through the wrapper's explicit close."""
|
|
client = _ScriptedClient(
|
|
[([StreamPiece("content", "A "), StreamPiece("content", "B ")], None)]
|
|
)
|
|
|
|
async def run() -> None:
|
|
gen = chat_stream_retried(client, _RETRY_MSGS, retries=3, delay=1.0)
|
|
first = await gen.__anext__()
|
|
assert first == StreamPiece("content", "A ")
|
|
await gen.aclose() # the consumer stops after the first piece
|
|
|
|
asyncio.run(run())
|
|
assert client.closed == [0] # attempt 1's stream was closed
|
|
assert len(client.request_args) == 1 # no second attempt
|
|
|
|
|
|
def test_retried_abandon_during_retry_sleep_leaks_nothing(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Consumer abandon while the generator is parked in the pre-retry
|
|
sleep (client disconnect): the driving task is cancelled cleanly,
|
|
the phase-48 teardown ``aclose()`` on the generator does not raise,
|
|
and the retry never starts."""
|
|
entered = asyncio.Event()
|
|
|
|
async def parking_sleep(seconds: float) -> None:
|
|
entered.set()
|
|
await asyncio.Event().wait() # park until the abandon arrives
|
|
|
|
monkeypatch.setattr(asyncio, "sleep", parking_sleep)
|
|
client = _ScriptedClient(
|
|
[([], LLMError("endpoint down")), ([StreamPiece("content", "never")], None)]
|
|
)
|
|
|
|
async def run() -> None:
|
|
gen = chat_stream_retried(client, _RETRY_MSGS, retries=3, delay=1.5)
|
|
|
|
async def consumer() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
|
|
return [p async for p in gen]
|
|
|
|
task = asyncio.ensure_future(consumer())
|
|
await entered.wait() # the generator is inside the pre-retry sleep
|
|
assert not task.done()
|
|
task.cancel() # client disconnect: the driving task is cancelled
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
# Production teardown (phase 48 pattern): the endpoint's finally
|
|
# closes the stream generator — must not raise.
|
|
await gen.aclose()
|
|
|
|
asyncio.run(run())
|
|
assert len(client.request_args) == 1 # the retry never started
|
|
assert client.closed == [0] # attempt 1's stream was torn down
|