diff --git a/lib/stack/ai/client.py b/lib/stack/ai/client.py index 94a59762..8bb55fb1 100644 --- a/lib/stack/ai/client.py +++ b/lib/stack/ai/client.py @@ -141,6 +141,11 @@ class LLMImage: "IRMhEyETIRMhEyEThXzH8QM9OMM6fAAAAABJRU5ErkJggg==" ) +# The probe asks for one word. A success needs no complete answer, so a +# small cap costs nothing and keeps the server from sizing for its +# default output length. +_PROBE_MAX_TOKENS = 16 + # Substrings text-only models emit when rejecting a multimodal request — # used to tell "no vision" (cache it) from "transport flaked" (don't). _NO_VISION_HINTS = ( @@ -306,7 +311,8 @@ async def has_vision(self, *, role: str = "classifier", probe_img = LLMImage(data=base64.b64decode(_PROBE_PNG_B64), mime="image/png") try: await self.complete(role, "Reply with the single word 'ok'.", - images=[probe_img], model_override=model) + images=[probe_img], model_override=model, + max_tokens=_PROBE_MAX_TOKENS) self.capabilities.record_vision(model, True) logger.info("[llm] vision probe: {} -> supported", model) return True diff --git a/stacklets/docs/bot/pipeline.py b/stacklets/docs/bot/pipeline.py index 8ef4416a..28b77f5f 100644 --- a/stacklets/docs/bot/pipeline.py +++ b/stacklets/docs/bot/pipeline.py @@ -101,6 +101,34 @@ def __init__(self, doc_id: int | None, title: str | None): # enrichment functions further down the file consume the same value. DEFAULT_CLASSIFY_MAX_CHARS = 20000 +# Output caps for each LLM call. Every answer has a size known in +# advance, so every call is capped. An uncapped call that repeats itself +# generates until the client timeout. A server that sizes its memory for +# the requested output uses its own default without a cap (on oMLX the +# full context window), and can refuse the request when that does not +# fit. Each cap fits the largest correct answer with margin, because a +# truncated JSON answer files the document unclassified. +# +# classify: about 15 tokens per line-item fact, so 4096 fits a receipt +# with more than 150 items plus the other fields. +CLASSIFY_MAX_TOKENS = 4096 +# capture: a summary of up to 400 words plus facts without a count limit. +CAPTURE_MAX_TOKENS = 4096 +# recall: a prose answer that cites hits by number. +SYNTHESIZE_MAX_TOKENS = 2048 +# reformat returns the whole OCR text with Markdown added, so its cap +# follows the input. Three characters per token is pessimistic for +# German and OCR noise; the headroom covers headings and table pipes. +_REFORMAT_CHARS_PER_TOKEN = 3 +_REFORMAT_HEADROOM = 1.5 +_REFORMAT_TOKEN_FLOOR = 512 + + +def _reformat_budget(text: str) -> int: + """Token cap for reformatting ``text``: its size plus Markdown headroom.""" + estimated = len(text) / _REFORMAT_CHARS_PER_TOKEN * _REFORMAT_HEADROOM + return max(_REFORMAT_TOKEN_FLOOR, int(estimated)) + # ── Enrichment result ──────────────────────────────────────────────────── @@ -695,7 +723,7 @@ async def classify( response = await self._llm.complete( "classifier", prompt, images=attach, json_mode=True, - temperature=0.0, + temperature=0.0, max_tokens=CLASSIFY_MAX_TOKENS, ) if not response: return {} @@ -766,7 +794,7 @@ async def classify_capture( ) response = await self._llm.complete( "classifier", prompt, images=attach, json_mode=True, - temperature=0.0, + temperature=0.0, max_tokens=CAPTURE_MAX_TOKENS, ) if not response: return {} @@ -800,9 +828,12 @@ async def reformat( is the pipeline's job (see `reformat_document`); this returns whatever the LLM produced, trimmed. """ - prompt = _build_reformat_prompt(ocr_text[:max_chars]) + text = ocr_text[:max_chars] + prompt = _build_reformat_prompt(text) try: - result = await self._llm.complete("reformat", prompt) + result = await self._llm.complete( + "reformat", prompt, max_tokens=_reformat_budget(text), + ) except (LLMUnavailableError, LLMModelNotFoundError, LLMTimeoutError): return None return result.strip() if result else None @@ -871,7 +902,10 @@ async def synthesize_answer( today = date.today().isoformat() prompt = _build_synthesize_prompt(question, evidence, lang, today=today) try: - raw = await self._llm.complete("recall", prompt, json_mode=False) + raw = await self._llm.complete( + "recall", prompt, json_mode=False, + max_tokens=SYNTHESIZE_MAX_TOKENS, + ) except (LLMUnavailableError, LLMModelNotFoundError, LLMTimeoutError) as e: logger.warning("[recall] LLM unavailable for synthesis: {}", e) return "" diff --git a/tests/unit/framework/test_ai_client.py b/tests/unit/framework/test_ai_client.py index 9260013b..4d49990a 100644 --- a/tests/unit/framework/test_ai_client.py +++ b/tests/unit/framework/test_ai_client.py @@ -357,6 +357,20 @@ async def test_inconclusive_error_returns_false_uncached(self, httpserver: HTTPS assert ModelCapabilities(path=tmp_path / "caps.json").supports_vision("test-model") is None await llm.aclose() + async def test_probe_asks_for_a_one_word_answer(self, httpserver: HTTPServer, tmp_path): + """The probe needs one word back. A server that plans memory for + the requested output size must not reserve its full default for it.""" + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_json(_completion_payload("ok")) + + llm = _make_llm(httpserver, capabilities=ModelCapabilities(path=tmp_path / "c.json")) + await llm.has_vision() + + request, _ = httpserver.log[0] + assert request.get_json()["max_tokens"] <= 16 + await llm.aclose() + async def test_cached_answer_skips_probe(self, httpserver: HTTPServer, tmp_path): """A pre-recorded answer means zero HTTP traffic.""" cap = ModelCapabilities(path=tmp_path / "caps.json") diff --git a/tests/unit/stacklets/test_pipeline.py b/tests/unit/stacklets/test_pipeline.py index e73f2501..9c8aeb90 100644 --- a/tests/unit/stacklets/test_pipeline.py +++ b/tests/unit/stacklets/test_pipeline.py @@ -1316,7 +1316,7 @@ async def test_new_notes_carry_the_marker(self, seeded_paperless): # — we test only the Classifier-side decision here. -from pipeline import Classifier, ImageAttachment # noqa: E402 +from pipeline import DEFAULT_CLASSIFY_MAX_CHARS, Classifier, ImageAttachment # noqa: E402 from stack.ai.client import ModelCapabilities # noqa: E402 @@ -1338,11 +1338,12 @@ def __init__(self, *, response: str = '{"title": "t"}', self.calls: list[dict] = [] async def complete(self, role, prompt, *, images=None, - json_mode=False, model_override=None, temperature=None): + json_mode=False, model_override=None, temperature=None, + max_tokens=None, timeout=None): self.calls.append({ "role": role, "prompt": prompt, "images": images, "json_mode": json_mode, "model_override": model_override, - "temperature": temperature, + "temperature": temperature, "max_tokens": max_tokens, }) return self._response @@ -1450,6 +1451,109 @@ async def test_non_image_mime_filtered(self): assert c._stub.calls[0]["images"] is None +# Pessimistic on purpose, as for the transcript polish: German compounds +# and OCR noise tokenize denser than English prose. A cap measured with +# this ratio fits the answer under any real tokenizer. +_CHARS_PER_TOKEN = 3 + + +def _itemised_receipt_answer(line_items: int) -> str: + """The classify answer for a long till receipt, one fact per item.""" + import json + return json.dumps({ + "title": "Kwik-E-Mart Kassenbon", + "date": "2026-09-14", + "correspondent": "Kwik-E-Mart", + "document_type": "Kassenbon", + "tags": ["Einkauf", "Lebensmittel"], + "persons": ["Marge"], + "correspondent_facts": [ + "Adresse: 742 Evergreen Terrace, Springfield", + "Telefon: +1 555 0100", + ], + "summary": "Kassenbon von Kwik-E-Mart vom 14.09.2026 für Marge " + "über EUR 412,80. Wocheneinkauf mit Getränken und " + "Haushaltswaren.", + "facts": ["Summe: EUR 412,80", "Bon-Nr.: 0042-1187"] + [ + f"Artikel Nr. {i} Squishee Kirsch groß 2x EUR 3,49" + for i in range(line_items) + ], + "action_items": [], + }, ensure_ascii=False, indent=2) + + +class TestAnswersAreBounded: + """Every answer the archivist asks for has a size known in advance, + so every call carries an output cap. Without one, a model that + repeats itself generates until the client timeout, and the server + plans memory for its own default, which on oMLX is the whole context + window. The cap must still fit the largest correct answer: one set + too low truncates the JSON, and the document is filed unclassified.""" + + @pytest.mark.asyncio + async def test_classify_has_room_for_a_receipt_with_150_line_items(self): + c = _make_classifier() + await c.classify( + ocr_text="Kassenbon", tags={}, doc_types={}, correspondents={}, + ) + answer = _itemised_receipt_answer(line_items=150) + cap = c._stub.calls[0]["max_tokens"] + assert cap is not None + assert cap >= len(answer) / _CHARS_PER_TOKEN + + @pytest.mark.asyncio + async def test_capture_has_room_for_a_full_article_digest(self): + """The capture prompt asks for up to 400 words of summary on a + long article, and facts without a count limit.""" + import json + c = _make_classifier() + await c.classify_capture(text="article", person_names=["Lisa"]) + answer = json.dumps({ + "title": "Saxophon-Wartung: Blätter, Polster und Korken", + "summary": " ".join(["Holzblasinstrument"] * 400), + "facts": [f"Fakt {i}: Rico-Blatt Stärke 2,5 kostet EUR 3,20 " + f"pro Stück im Zehnerpack" for i in range(30)], + "tags": ["saxophon", "instrumentenpflege", "blaetter"], + "persons": ["Lisa"], + }, ensure_ascii=False, indent=2) + cap = c._stub.calls[0]["max_tokens"] + assert cap is not None + assert cap >= len(answer) / _CHARS_PER_TOKEN + + @pytest.mark.asyncio + async def test_reformat_has_room_for_the_longest_document_it_is_given(self): + """Reformat returns the whole text again, with Markdown added.""" + c = _make_classifier(response="text") + ocr = "Rechnung Nr. 4711 Springfield Kraftwerk " * 600 + await c.reformat(ocr) + cap = c._stub.calls[0]["max_tokens"] + assert cap is not None + assert cap >= DEFAULT_CLASSIFY_MAX_CHARS / _CHARS_PER_TOKEN + + @pytest.mark.asyncio + async def test_reformat_cap_follows_the_document_length(self): + """A one-page letter must not reserve room for a 20-page contract.""" + c = _make_classifier(response="text") + await c.reformat("Lieber Homer, " * 20) + await c.reformat("Lieber Homer, " * 1000) + short_cap, long_cap = (call["max_tokens"] for call in c._stub.calls) + assert short_cap < long_cap + + @pytest.mark.asyncio + async def test_a_recall_answer_is_bounded(self): + c = _make_classifier(response="answer") + await c.synthesize_answer( + "Wann ist die Kfz-Versicherung fällig?", + [{"title": "Kfz-Versicherung 2026", "summary": "Fällig 01.10."}], + ) + # The prompt asks for one sentence when that answers it. This is + # a long answer: about 330 words that cite their sources. + answer = "Die Kfz-Versicherung für den Wagen ist am 1. Oktober fällig [1]. " * 30 + cap = c._stub.calls[0]["max_tokens"] + assert cap is not None + assert cap >= len(answer) / _CHARS_PER_TOKEN + + # The recall-mode query rewrite used to be tested here. It now lives in # memory, which owns the vault it searches, and so do its tests: # tests/unit/stacklets/test_memory_query_rewrite.py.