diff --git a/lib/stack/ai/client.py b/lib/stack/ai/client.py index 52bd570..3117d7a 100644 --- a/lib/stack/ai/client.py +++ b/lib/stack/ai/client.py @@ -21,10 +21,12 @@ from __future__ import annotations +import asyncio import base64 import os from dataclasses import dataclass +import httpx import openai from loguru import logger from openai import AsyncOpenAI @@ -148,6 +150,46 @@ class LLMImage: ) +# ── Liveness, not deadlines ────────────────────────────────────────────── +# +# A local model on a small Mac is slow, not broken. A 30K-token prefill on +# a scanned document can run for minutes before the first token, and a +# wall-clock deadline cannot tell that apart from a wedged server: it just +# hangs up on work that was progressing fine, throwing away everything +# spent so far. That is what a non-streaming call forces, because not one +# byte comes back until generation is complete. +# +# So we stream and time the *gaps* instead. Two different questions, two +# different budgets: +# +# FIRST_TOKEN_TIMEOUT — how long we allow silence before anything comes +# back. This covers prefill, which scales with input size and with how +# hard the machine is swapping, so it has to be generous. +# STALL_TIMEOUT — how long we allow silence *between* pieces of an answer +# already underway. Once tokens flow, gaps are short on any machine, so +# a long one means the server stopped rather than slowed. +# +# Measured against local oMLX (Qwen3.5-9B-MLX-4bit, dense 14-page PDF), to +# first token: +# +# 1,000 tokens -> 7.6s 15,000 tokens -> 58.4s +# 3,000 tokens -> 20.4s 28,400 tokens -> 163.6s +# 7,500 tokens -> 36.7s +# +# Roughly linear at ~5.7ms/token, and that is a machine with RAM to spare; +# a 16GB Mac running photos and documents alongside is several times +# slower. Gaps *after* the first token were 0.07-0.13s. Hence the lopsided +# pair: the first-token budget must swallow minutes of nothing, while the +# stall budget can be tight enough to notice a dead server and still sit +# orders of magnitude above any real gap. +# +# The consequence that matters: while output is arriving we never give up, +# no matter how long the document takes. A genuinely hung endpoint is +# caught sooner than the old deadline caught it. +FIRST_TOKEN_TIMEOUT = 900.0 +STALL_TIMEOUT = 90.0 + + class LLM: """An OpenAI-compatible chat client scoped to one stacklet/bot. @@ -158,9 +200,13 @@ class LLM: """ def __init__(self, client: AsyncOpenAI, *, namespace: str | None = None, - capabilities: ModelCapabilities | None = None): + capabilities: ModelCapabilities | None = None, + first_token_timeout: float = FIRST_TOKEN_TIMEOUT, + stall_timeout: float = STALL_TIMEOUT): self._client = client self.namespace = namespace + self.first_token_timeout = first_token_timeout + self.stall_timeout = stall_timeout # Default to in-memory; bots inject a disk-backed cache so the # vision probe survives container restarts. self.capabilities = capabilities or ModelCapabilities() @@ -190,24 +236,17 @@ def from_env(cls, *, namespace: str | None = None, "No AI endpoint configured — set up AI with 'stack up ai'" ) key = os.environ.get("OPENAI_KEY", "") or "not-needed" - # This is a read timeout, and the call is not streamed, so it - # covers the entire silent wait while the model reads a document - # and generates an answer. Prefill alone, measured against local - # oMLX (Qwen3.5-9B-MLX-4bit) on a dense 14-page PDF: - # - # 1,000 tokens -> 7.6s 15,000 tokens -> 58.4s - # 3,000 tokens -> 20.4s 28,400 tokens -> 163.6s - # 7,500 tokens -> 36.7s - # - # That machine had RAM to spare. A 16GB Mac running photos and - # documents alongside is several times slower, so 120s and then - # 300s were both cancelling healthy work partway through a big - # scan and discarding every minute already spent on it. 900s - # covers ~158K tokens here, past any context this serves, and - # still bounds a genuinely dead endpoint. + # No read timeout on purpose: `complete` streams and enforces + # liveness per chunk (see FIRST_TOKEN_TIMEOUT / STALL_TIMEOUT). + # A deadline down here could only ever fire on work that was still + # producing output, which is the bug being fixed. Raising the + # ceiling (120s, then 300s, then 900s) only moved the point at + # which a big enough document got cancelled; timing the silence + # instead removes it. Connect stays short: an endpoint that will + # not accept a socket is down, and that answer arrives at once. client = AsyncOpenAI( base_url=url, api_key=key, max_retries=max_retries, - timeout=900.0, + timeout=httpx.Timeout(None, connect=10.0), ) return cls(client, namespace=namespace, capabilities=capabilities) @@ -242,11 +281,26 @@ async def complete(self, role: str, prompt: str, *, kwargs["temperature"] = temperature try: - resp = await self._client.chat.completions.create( - model=model, - messages=[{"role": "user", "content": content}], - **kwargs, + # The first-token budget has to cover this call, not just the + # chunks after it: a server that withholds its response headers + # until the first token spends the whole prefill in here, where + # no per-chunk timeout can see it. With no read timeout on the + # SDK, that would otherwise be an unbounded wait. + stream = await asyncio.wait_for( + self._client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": content}], + stream=True, + **kwargs, + ), + timeout=self.first_token_timeout, ) + return await self._read_stream(stream, model) + except asyncio.TimeoutError as e: + raise LLMTimeoutError( + f"{model} — went silent before answering " + f"for {self.first_token_timeout:.0f}s" + ) from e # Order matters: APITimeoutError < APIConnectionError, and # Authentication/NotFound < APIStatusError — catch specific first. except openai.APITimeoutError as e: @@ -260,7 +314,53 @@ async def complete(self, role: str, prompt: str, *, except openai.APIStatusError as e: raise LLMUnavailableError(f"HTTP {e.status_code}: {str(e)[:200]}") from e - return resp.choices[0].message.content or "" + async def _read_stream(self, stream, model: str) -> str: + """Reassemble a streamed completion, giving up only on silence. + + Callers still get one string back; streaming is an implementation + detail they should not have to care about. What it buys is the + progress signal: every chunk that arrives proves the server is + still working, so the budget for the next one resets. + + The switch from the first-token budget to the stall budget waits + for actual *content*, not merely the first chunk. Some servers open + with a role-only preamble before prefill has finished, and treating + that as "generation started" would drop us onto the tight budget + with the long silent part still ahead. + """ + parts: list[str] = [] + chunks = stream.__aiter__() + while True: + budget = self.stall_timeout if parts else self.first_token_timeout + try: + chunk = await asyncio.wait_for(chunks.__anext__(), timeout=budget) + except StopAsyncIteration: + break + except asyncio.TimeoutError as e: + await self._close_stream(stream) + where = "mid-answer" if parts else "before answering" + raise LLMTimeoutError( + f"{model} — went silent {where} for {budget:.0f}s" + ) from e + # A chunk with no choices (a usage-only trailer) still counts as + # the server being alive, it just carries no text. + if chunk.choices: + delta = chunk.choices[0].delta + if delta is not None and delta.content: + parts.append(delta.content) + return "".join(parts) + + @staticmethod + async def _close_stream(stream) -> None: + """Release an abandoned stream's connection, best effort. + + We only get here on a timeout, where the useful error is the + timeout itself; a failure to hang up cleanly must not replace it. + """ + try: + await stream.close() + except Exception as e: + logger.debug("[llm] closing abandoned stream failed: {}", e) async def has_vision(self, *, role: str = "classifier", model_override: str | None = None) -> bool: diff --git a/tests/framework/test_ai_client.py b/tests/framework/test_ai_client.py index 3e80f72..9692f2c 100644 --- a/tests/framework/test_ai_client.py +++ b/tests/framework/test_ai_client.py @@ -14,6 +14,7 @@ import base64 import json import sys +import time from pathlib import Path import httpx @@ -49,7 +50,9 @@ def stub_resolver(monkeypatch): def _make_llm(httpserver: HTTPServer, *, timeout: float = 5.0, - capabilities: ModelCapabilities | None = None) -> LLM: + capabilities: ModelCapabilities | None = None, + first_token_timeout: float = 5.0, + stall_timeout: float = 5.0) -> LLM: """Build an LLM pointed at the local httpserver mock. `max_retries=0` keeps tests fast and lets a 401/404 surface on the @@ -60,21 +63,162 @@ def _make_llm(httpserver: HTTPServer, *, timeout: float = 5.0, max_retries=0, timeout=timeout, ) - return LLM(client, namespace="test-bot", capabilities=capabilities) + return LLM(client, namespace="test-bot", capabilities=capabilities, + first_token_timeout=first_token_timeout, + stall_timeout=stall_timeout) -def _completion_payload(content: str, *, model: str = "test-model") -> dict: - return { +def _chunk(delta: dict, *, model: str = "test-model") -> str: + """One `chat.completion.chunk` as an SSE event.""" + body = { "id": "cmpl-test", - "object": "chat.completion", + "object": "chat.completion.chunk", "model": model, - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": content}, - "finish_reason": "stop", - }], - "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + "choices": [{"index": 0, "delta": delta, "finish_reason": None}], } + return f"data: {json.dumps(body)}\n\n" + + +def _sse_body(*contents: str, model: str = "test-model") -> str: + """A complete streamed answer: one chunk per piece, then DONE.""" + events = [_chunk({"content": c}, model=model) for c in contents] + return "".join(events) + "data: [DONE]\n\n" + + +def _sse_response(*contents: str, model: str = "test-model"): + """The whole answer at once. `complete` still returns it joined.""" + from werkzeug.wrappers import Response + return Response(_sse_body(*contents, model=model), + content_type="text/event-stream") + + +def _sse_trickle(*contents: str, gap: float = 0.0, lead: float = 0.0, + tail: float = 0.0, preamble: bool = False): + """A stream paced like a real one, so silence can be tested. + + `lead` is the quiet before anything arrives (prefill). `gap` is the + quiet between pieces. `tail` is a hang after the last piece but before + DONE. `preamble` opens with a role-only chunk carrying no content, the + way some servers announce themselves before generating. + """ + from werkzeug.wrappers import Response + + def generate(): + if preamble: + yield _chunk({"role": "assistant"}) + if lead: + time.sleep(lead) + for i, c in enumerate(contents): + if i and gap: + time.sleep(gap) + yield _chunk({"content": c}) + if tail: + time.sleep(tail) + yield "data: [DONE]\n\n" + + return Response(generate(), content_type="text/event-stream") + + +# ── complete(): liveness ───────────────────────────────────────────────── + +class TestSlowIsNotBroken: + """The distinction the old wall-clock timeout could not make. + + A local model on a small Mac is slow, not broken. The client used to + wait for a complete response with no bytes on the wire, so the only + question it could ask was "has too much time passed", and the answer + was the same for a 30K-token prefill that was progressing fine and a + server that had wedged. It cancelled healthy work mid-flight and threw + away everything spent on it. + + Streaming turns that into a question worth asking: has it gone *quiet*. + """ + + async def test_an_answer_that_keeps_arriving_is_never_cut_off( + self, httpserver: HTTPServer): + """Total time far exceeds the stall budget; no single gap does. + + This is the property that matters for a slow machine: as long as + output keeps coming, there is no deadline at all. + """ + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler( + lambda _r: _sse_trickle("slow ", "but ", "still ", "working", gap=0.3), + ) + + llm = _make_llm(httpserver, stall_timeout=0.6, first_token_timeout=2.0) + result = await llm.complete("classifier", "long document") + assert result == "slow but still working" + await llm.aclose() + + async def test_a_long_silent_prefill_is_allowed(self, httpserver: HTTPServer): + """Nothing comes back while the model reads a long document. That + silence is the work, not a fault, so it gets its own budget.""" + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler(lambda _r: _sse_trickle("done", lead=0.8)) + + llm = _make_llm(httpserver, first_token_timeout=3.0, stall_timeout=0.3) + assert await llm.complete("classifier", "scan") == "done" + await llm.aclose() + + async def test_a_role_only_preamble_does_not_start_the_stall_clock( + self, httpserver: HTTPServer): + """Some servers announce themselves before generating. Counting + that as "generation started" would drop us onto the tight budget + with the whole silent prefill still ahead.""" + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler( + lambda _r: _sse_trickle("here", preamble=True, lead=0.8), + ) + + llm = _make_llm(httpserver, first_token_timeout=3.0, stall_timeout=0.3) + assert await llm.complete("classifier", "scan") == "here" + await llm.aclose() + + +class TestSilenceIsBroken: + """The other half: a server that stops producing must still be caught, + and caught sooner than the old deadline caught it.""" + + async def test_silence_before_any_answer_gives_up(self, httpserver: HTTPServer): + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler(lambda _r: _sse_trickle("never", lead=2.0)) + + llm = _make_llm(httpserver, first_token_timeout=0.3, stall_timeout=0.3) + with pytest.raises(LLMTimeoutError) as excinfo: + await llm.complete("classifier", "hi") + assert "before answering" in str(excinfo.value) + await llm.aclose() + + async def test_silence_partway_through_an_answer_gives_up( + self, httpserver: HTTPServer): + """Generation started and then stopped. The tight budget applies + here precisely because tokens were already flowing.""" + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler(lambda _r: _sse_trickle("started", "stalled", gap=2.0)) + + llm = _make_llm(httpserver, first_token_timeout=3.0, stall_timeout=0.3) + with pytest.raises(LLMTimeoutError) as excinfo: + await llm.complete("classifier", "hi") + assert "mid-answer" in str(excinfo.value) + await llm.aclose() + + +class TestReassembly: + async def test_deltas_are_joined_into_one_answer(self, httpserver: HTTPServer): + """Streaming is an implementation detail; callers get a string.""" + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler(lambda _r: _sse_response("Kwik-", "E-", "Mart")) + + llm = _make_llm(httpserver) + assert await llm.complete("classifier", "shop") == "Kwik-E-Mart" + await llm.aclose() # ── complete(): happy path ─────────────────────────────────────────────── @@ -83,7 +227,7 @@ class TestComplete: async def test_returns_assistant_text(self, httpserver: HTTPServer): httpserver.expect_request( "/v1/chat/completions", method="POST", - ).respond_with_json(_completion_payload("hello world")) + ).respond_with_handler(lambda _r: _sse_response("hello world")) llm = _make_llm(httpserver) result = await llm.complete("classifier", "say hi") @@ -98,11 +242,7 @@ async def test_sends_resolved_model(self, httpserver: HTTPServer, monkeypatch): def handler(request): captured["body"] = json.loads(request.get_data().decode()) - from werkzeug.wrappers import Response - return Response( - json.dumps(_completion_payload("ok")), - content_type="application/json", - ) + return _sse_response("ok") httpserver.expect_request( "/v1/chat/completions", method="POST", @@ -118,11 +258,7 @@ async def test_json_mode_sets_response_format(self, httpserver: HTTPServer): def handler(request): captured["body"] = json.loads(request.get_data().decode()) - from werkzeug.wrappers import Response - return Response( - json.dumps(_completion_payload('{"k":1}')), - content_type="application/json", - ) + return _sse_response('{"k":1}') httpserver.expect_request( "/v1/chat/completions", method="POST", @@ -138,11 +274,7 @@ async def test_images_sent_as_data_url(self, httpserver: HTTPServer): def handler(request): captured["body"] = json.loads(request.get_data().decode()) - from werkzeug.wrappers import Response - return Response( - json.dumps(_completion_payload("ok")), - content_type="application/json", - ) + return _sse_response("ok") httpserver.expect_request( "/v1/chat/completions", method="POST", @@ -164,11 +296,7 @@ async def test_model_override_bypasses_resolver(self, httpserver: HTTPServer): def handler(request): captured["body"] = json.loads(request.get_data().decode()) - from werkzeug.wrappers import Response - return Response( - json.dumps(_completion_payload("ok")), - content_type="application/json", - ) + return _sse_response("ok") httpserver.expect_request( "/v1/chat/completions", method="POST", @@ -322,7 +450,7 @@ async def test_probe_success_caches_true(self, httpserver: HTTPServer, tmp_path) cap = ModelCapabilities(path=tmp_path / "caps.json") httpserver.expect_request( "/v1/chat/completions", method="POST", - ).respond_with_json(_completion_payload("ok")) + ).respond_with_handler(lambda _r: _sse_response("ok")) llm = _make_llm(httpserver, capabilities=cap) assert await llm.has_vision() is True diff --git a/tests/integration/openai_stub.py b/tests/integration/openai_stub.py index c1afc38..601383d 100644 --- a/tests/integration/openai_stub.py +++ b/tests/integration/openai_stub.py @@ -53,19 +53,27 @@ _KINDS = ("classify", "reformat", "rewrite", "synthesize") -def _chat_completion(content: str, model: str = "test-model") -> dict: - """OpenAI chat.completion response envelope.""" - return { +def _chat_completion(content: str, model: str = "test-model") -> Response: + """A streamed chat completion, the way real backends answer. + + The client streams every call so it can tell a slow model from a stopped + one, so a stub that replies with a single JSON body is no longer + speaking the same protocol as the thing it stands in for. Sent as one + content chunk: the stub's job is to be a believable endpoint, not to + reproduce token-by-token pacing (the client's own tests cover pacing). + """ + chunk = { "id": "cmpl-test", - "object": "chat.completion", + "object": "chat.completion.chunk", "model": model, "choices": [{ "index": 0, - "message": {"role": "assistant", "content": content}, - "finish_reason": "stop", + "delta": {"role": "assistant", "content": content}, + "finish_reason": None, }], - "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, } + body = f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n" + return Response(body, content_type="text/event-stream") def _prompt_of(body: dict) -> str: @@ -125,10 +133,7 @@ def _dispatch(self, request: Request) -> Response: body = {} prompt = _prompt_of(body) if prompt.lstrip().startswith(_VISION_PROBE_MARKER): - return Response( - json.dumps(_chat_completion("ok")), - content_type="application/json", - ) + return _chat_completion("ok") head = " ".join(prompt.split())[:120] kind = _route(prompt) with self._lock: @@ -142,10 +147,7 @@ def _dispatch(self, request: Request) -> Response: ) return Response(f"openai stub: {kind} queue empty", status=500) content = queue.popleft() - return Response( - json.dumps(_chat_completion(content)), - content_type="application/json", - ) + return _chat_completion(content) # ── queueing ─────────────────────────────────────────────────────────