204 lines
7.5 KiB
Python
204 lines
7.5 KiB
Python
"""Unit: deterministic model-stream teardown (phase 48, task 01).
|
|
|
|
The openai SDK stream must be closed on **every** exit of
|
|
``LLMClient.chat_stream`` after a successful ``create()`` — normal
|
|
exhaustion, a wrapped mid-stream failure, and consumer abandon
|
|
(``aclose()`` on the ``chat_stream`` generator — the stop/cancel path,
|
|
2026-08-29, ``TODO.md`` L3). The fakes stand in at the SDK boundary: a
|
|
recording async stream (small sleeps between chunks so an abandon can
|
|
land mid-iteration) behind an ``AsyncOpenAI``-shaped client, following
|
|
the fake patterns of ``tests/unit/test_llm_client.py``.
|
|
|
|
Note on the close method: the phase text says ``aclose()`` — the
|
|
installed openai SDK's ``AsyncStream`` exposes the async ``close()``,
|
|
which awaits the underlying httpx response's ``aclose()``; that is the
|
|
method under test (a quiet no-op on an already-ended SDK stream, so the
|
|
completed path stays byte-identical).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
from app.config import Settings
|
|
from app.rag.llm import LLMClient, LLMError, StreamPiece
|
|
|
|
|
|
def _settings(**kwargs: Any) -> Settings:
|
|
kwargs.setdefault("_env_file", None)
|
|
return Settings(**kwargs) # pyright: ignore[reportCallIssue]
|
|
|
|
|
|
def _chunk(content: str) -> SimpleNamespace:
|
|
"""One fake ChatCompletionChunk (``choices[].delta.content`` shape)."""
|
|
return SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content=content))])
|
|
|
|
|
|
def _tool_chunk() -> SimpleNamespace:
|
|
"""One chunk carrying a malformed-arguments tool call (index 0)."""
|
|
fn = SimpleNamespace(name="read", arguments='{"path": "Homelab",')
|
|
tc = SimpleNamespace(index=0, id="call_x", function=fn)
|
|
delta = SimpleNamespace(content=None, tool_calls=[tc])
|
|
return SimpleNamespace(choices=[SimpleNamespace(delta=delta)])
|
|
|
|
|
|
class _RecordingStream:
|
|
"""A fake SDK stream: yields *chunks* (small sleeps between them, so
|
|
an abandon can land mid-iteration) and records ``close()`` calls."""
|
|
|
|
def __init__(self, chunks: list, delay: float = 0.005) -> None:
|
|
self._chunks = list(chunks)
|
|
self._delay = delay
|
|
self._i = 0
|
|
self.close_calls = 0
|
|
|
|
def __aiter__(self) -> _RecordingStream:
|
|
return self
|
|
|
|
async def __anext__(self) -> SimpleNamespace:
|
|
if self._i >= len(self._chunks):
|
|
raise StopAsyncIteration
|
|
await asyncio.sleep(self._delay)
|
|
chunk = self._chunks[self._i]
|
|
self._i += 1
|
|
return chunk
|
|
|
|
async def close(self) -> None:
|
|
self.close_calls += 1
|
|
|
|
|
|
class _FailingStream(_RecordingStream):
|
|
"""Same shape, but ``__anext__`` raises a transport error after
|
|
*fail_after* chunks (the mid-iteration failure path)."""
|
|
|
|
def __init__(self, chunks: list, fail_after: int, delay: float = 0.005) -> None:
|
|
super().__init__(chunks, delay)
|
|
self._fail_after = fail_after
|
|
|
|
async def __anext__(self) -> SimpleNamespace:
|
|
self._i += 1
|
|
if self._i > self._fail_after:
|
|
raise ConnectionError("simulated mid-stream drop")
|
|
if self._i > len(self._chunks):
|
|
raise StopAsyncIteration
|
|
await asyncio.sleep(self._delay)
|
|
chunk = self._chunks[self._i - 1]
|
|
return chunk
|
|
|
|
|
|
class _FakeCompletions:
|
|
"""``chat.completions.create(stream=True)`` → the fake stream (or a
|
|
create-level failure)."""
|
|
|
|
def __init__(self, stream: _RecordingStream | Exception) -> None:
|
|
self._stream = stream
|
|
self.kwargs: dict | None = None
|
|
|
|
async def create(self, **kwargs: Any):
|
|
self.kwargs = kwargs
|
|
assert kwargs.get("stream") is True
|
|
if isinstance(self._stream, Exception):
|
|
raise self._stream
|
|
return self._stream
|
|
|
|
|
|
def _make_client(
|
|
stream: _RecordingStream | Exception,
|
|
) -> tuple[LLMClient, _FakeCompletions]:
|
|
completions = _FakeCompletions(stream)
|
|
llm = LLMClient(_settings())
|
|
llm._client = SimpleNamespace( # pyright: ignore[reportAttributeAccessIssue]
|
|
chat=SimpleNamespace(completions=completions)
|
|
)
|
|
return llm, completions
|
|
|
|
|
|
def _content_chunks(n: int) -> list:
|
|
return [_chunk(f"piece {i} ") for i in range(1, n + 1)]
|
|
|
|
|
|
async def _drain(llm: LLMClient) -> list[StreamPiece]:
|
|
"""Tools-less drain: without a ``tools`` list no ToolCallPiece can
|
|
appear (the phase-37 contract)."""
|
|
pieces = [p async for p in llm.chat_stream([{"role": "user", "content": "q"}])]
|
|
assert all(isinstance(p, StreamPiece) for p in pieces)
|
|
return cast("list[StreamPiece]", pieces)
|
|
|
|
|
|
def test_full_consumption_closes_stream_exactly_once() -> None:
|
|
"""(a) A fully consumed stream still gets the explicit close — the
|
|
completed path keeps its behavior (a quiet no-op on the real SDK
|
|
stream) while the teardown is pinned."""
|
|
stream = _RecordingStream(_content_chunks(3))
|
|
llm, _ = _make_client(stream)
|
|
pieces = asyncio.run(_drain(llm))
|
|
assert [p.text for p in pieces] == ["piece 1 ", "piece 2 ", "piece 3 "]
|
|
assert stream.close_calls == 1
|
|
|
|
|
|
def test_mid_iteration_abandon_closes_stream_before_close_completes() -> None:
|
|
"""(b) Abandoning the ``chat_stream`` generator after the first
|
|
piece (``await gen.aclose()`` — the consumer-stop path) must await
|
|
the SDK stream's close before the generator's close completes."""
|
|
stream = _RecordingStream(_content_chunks(10))
|
|
llm, _ = _make_client(stream)
|
|
|
|
async def abandon_after_first() -> None:
|
|
gen = llm.chat_stream([{"role": "user", "content": "q"}])
|
|
first = await gen.__anext__()
|
|
assert isinstance(first, StreamPiece)
|
|
assert first.text == "piece 1 "
|
|
# The generator is suspended at its first yield; aclose must run
|
|
# the finally (the SDK stream's close) before it returns.
|
|
await gen.aclose()
|
|
|
|
asyncio.run(abandon_after_first())
|
|
assert stream.close_calls == 1
|
|
|
|
|
|
def test_create_failure_wraps_and_never_closes() -> None:
|
|
"""A ``create()`` failure keeps today's wrap — generic exception →
|
|
``LLMError`` with the base URL — and no stream exists to close."""
|
|
llm, _ = _make_client(ConnectionError("connection reset by peer"))
|
|
with pytest.raises(LLMError, match="chat stream from .* failed") as exc:
|
|
asyncio.run(_drain(llm))
|
|
assert "connection reset by peer" in str(exc.value)
|
|
|
|
|
|
def test_mid_iteration_failure_wraps_and_closes() -> None:
|
|
"""A failure inside the ``async for`` wraps exactly as before
|
|
(``LLMError`` with the original message) — and the stream is closed
|
|
on the exception path."""
|
|
stream = _FailingStream(_content_chunks(10), fail_after=2)
|
|
llm, _ = _make_client(stream)
|
|
|
|
async def drain() -> None:
|
|
async for _ in llm.chat_stream([{"role": "user", "content": "q"}]):
|
|
pass
|
|
|
|
with pytest.raises(LLMError, match="simulated mid-stream drop"):
|
|
asyncio.run(drain())
|
|
assert stream.close_calls == 1
|
|
|
|
|
|
def test_llm_error_materialization_passes_through_and_closes() -> None:
|
|
"""An ``LLMError`` from tool-call materialization (after the loop,
|
|
before normal exhaustion) re-raises unwrapped — and the stream is
|
|
still closed on the way out."""
|
|
stream = _RecordingStream([_tool_chunk()])
|
|
llm, _ = _make_client(stream)
|
|
|
|
async def drain() -> None:
|
|
async for _ in llm.chat_stream(
|
|
[{"role": "user", "content": "q"}],
|
|
tools=[{"type": "function", "function": {"name": "read"}}],
|
|
):
|
|
pass
|
|
|
|
with pytest.raises(LLMError, match="malformed tool-call arguments"):
|
|
asyncio.run(drain())
|
|
assert stream.close_calls == 1
|