merge: #1120 — size the chat prompt in tokens and refuse silent host-side truncation (refs #1120)

src/llm/config.py

@@ -40,6 +40,11 @@ class LlmConfig:
     chat_num_ctx: int = 8192
     chat_temperature: float = 0.2
     chat_max_tokens: int = 1200
+    #: Characters per token the instruct model's tokenizer yields on the
+    #: indexed material; converts ``chat_num_ctx`` into the character budget
+    #: ``build_messages`` trims excerpts to. Conservative: an estimate that
+    #: still overflows is caught by the host and the prompt is rebuilt.
+    chat_chars_per_token: float = 2.4
     recency_half_life_days: float = 365.0
     recency_weight: float = 0.3
     k_per_kind: dict[str, int] = field(
@@ -122,6 +127,7 @@ def load() -> LlmConfig:
         chat_num_ctx=int(_opt(section, "chat_num_ctx", 8192)),
         chat_temperature=float(_opt(section, "chat_temperature", 0.2)),
         chat_max_tokens=int(_opt(section, "chat_max_tokens", 1200)),
+        chat_chars_per_token=float(_opt(section, "chat_chars_per_token", 2.4)),
         recency_half_life_days=float(_opt(section, "recency_half_life_days", 365.0)),
         recency_weight=float(_opt(section, "recency_weight", 0.3)),
         k_per_kind=k_per_kind,

src/llm/metrics.py

@@ -15,6 +15,7 @@ Instruments (all exported as ``stack_llm_…``)::
     stack_llm_chat_seconds{stage}              retrieve / generate
     stack_llm_chat_total{outcome}              ok / error
     stack_llm_chat_tokens_total{direction}     prompt / completion
+    stack_llm_chat_prompt_overflow_total{resolution}  rebuilt / host_truncated
     stack_llm_indexed_chunks_total{collection} chunks written to pgvector
     stack_llm_indexed_items_total{collection,outcome}
 
@@ -89,6 +90,11 @@ class _Instruments:
             unit="1",
             description="Chat tokens reported by Ollama, by direction",
         )
+        self.chat_prompt_overflow_total = m.create_counter(
+            "stack_llm_chat_prompt_overflow_total",
+            unit="1",
+            description="Chat prompts the host refused as longer than the context, by how they were resolved",
+        )
         self.indexed_chunks_total = m.create_counter(
             "stack_llm_indexed_chunks_total",
             unit="1",
@@ -166,6 +172,13 @@ def chat_tokens(prompt: int, completion: int) -> None:
         i.chat_tokens_total.add(completion, {"direction": "completion"})
 
 
+def chat_prompt_overflow(resolution: str) -> None:
+    """Count one prompt the host refused as longer than its context —
+    ``rebuilt`` (trimmed tighter and resent) or ``host_truncated`` (the
+    last resort: resent with host-side truncation allowed)."""
+    instruments().chat_prompt_overflow_total.add(1, {"resolution": resolution})
+
+
 def indexed_item(collection: str, outcome: str, chunks: int = 0) -> None:
     """Count one item the indexer finished (``indexed``/``skipped``)."""
     i = instruments()

src/llm/rag.py

@@ -401,6 +401,77 @@ _BUDGET_CITED_KEEP = 8
 _BUDGET_MANUAL_KEEP = 1
 _BUDGET_VALUATION_ROWS_KEEP = 12
 
+# The model counts its context in tokens; the excerpts are trimmed in
+# characters before any model sees them. ``prompt_budget_chars`` converts
+# one into the other, and the chat refuses host-side truncation
+# (``truncate: false``) so a prompt that still overflows comes back as a
+# 400 with the token counts and is rebuilt tighter — never silently cut
+# to its tail half (which drops the system prompt and the best excerpts).
+_PROMPT_TEMPLATE_RESERVE_TOKENS = 64  # chat-template markup around the messages
+_PROMPT_ATTEMPTS = 3  # estimate, rebuild from the host's counts, then let the host cut
+_OVERFLOW_MARGIN = 0.9  # headroom on the rebuild when the host reports counts
+_OVERFLOW_SHRINK = 0.6  # rebuild factor when it reports only a message
+
+
+def prompt_budget_chars(cfg: LlmConfig) -> int:
+    """Characters the user message may hold so system prompt + user
+    message + reply fit ``chat_num_ctx``: the context minus the reply
+    reservation (``chat_max_tokens``) and template overhead, at
+    ``chat_chars_per_token``, minus the system prompt's own characters."""
+    tokens = cfg.chat_num_ctx - cfg.chat_max_tokens - _PROMPT_TEMPLATE_RESERVE_TOKENS
+    return max(int(tokens * cfg.chat_chars_per_token) - len(_SYSTEM), 0)
+
+
+def _overflow_counts(resp: httpx.Response) -> tuple[int, int] | None:
+    """``(prompt_tokens, num_ctx)`` when *resp* is Ollama refusing a prompt
+    longer than the context (the ``truncate: false`` 400), else None.
+
+    Current hosts nest a JSON error with ``n_prompt_tokens``/``n_ctx`` in
+    the ``error`` string; older ones send only a message naming the
+    context length — those yield ``(0, 0)``."""
+    if resp.status_code != 400:
+        return None
+    try:
+        payload = json.loads(resp.read())
+    except ValueError:
+        return None
+    err = payload.get("error", "") if isinstance(payload, dict) else ""
+    detail: dict = {}
+    if isinstance(err, dict):
+        detail = err
+    elif isinstance(err, str):
+        try:
+            inner = json.loads(err)
+            detail = inner.get("error", {}) if isinstance(inner, dict) else {}
+        except ValueError:
+            detail = {}
+    message = str(detail.get("message") or err)
+    if (
+        detail.get("type") != "exceed_context_size_error"
+        and "context" not in message.lower()
+    ):
+        return None
+    return int(detail.get("n_prompt_tokens") or 0), int(detail.get("n_ctx") or 0)
+
+
+def _relay_tokens(resp: httpx.Response) -> Iterator[dict]:
+    """Yield token events from one streaming ``/api/chat`` response and
+    record the token counts Ollama reports on its final frame."""
+    for line in resp.iter_lines():
+        if not line:
+            continue
+        data = json.loads(line)
+        chunk = data.get("message", {}).get("content", "")
+        if chunk:
+            yield {"type": "token", "text": chunk}
+        if data.get("done"):
+            # Older hosts/models omit the counts (chat_tokens skips 0).
+            metrics.chat_tokens(
+                int(data.get("prompt_eval_count") or 0),
+                int(data.get("eval_count") or 0),
+            )
+            break
+
 
 def _user_message(
     question: str,
@@ -782,50 +853,72 @@ def _stream_answer(
         yield evidence.payload()
     metrics.chat_stage("retrieve", time.monotonic() - began)
     pool.check(cfg.instruct_model)
-    messages = build_messages(
-        question,
-        sources,
-        evidence,
-        lineage,
-        n_retrieved=n_retrieved,
-        n_cited=n_cited,
-        n_lineage_sources=n_lineage_sources,
-        n_manual=n_manual,
-        budget_chars=cfg.chat_num_ctx * 3,
-    )
+    budget = prompt_budget_chars(cfg)
+    truncate = False
+    sent_len = 0
+    prompt_tokens = num_ctx = 0
     began = time.monotonic()
     with pool.acquire_generation() as host, httpx.Client(timeout=_TIMEOUT) as client:
         model = pick_model(cfg, pool, host)
-        with client.stream(
-            "POST",
-            f"{host}/api/chat",
-            json={
-                "model": model,
-                "messages": messages,
-                "stream": True,
-                "options": {
-                    "num_ctx": cfg.chat_num_ctx,
-                    "temperature": cfg.chat_temperature,
-                    "num_predict": cfg.chat_max_tokens,
+        for attempt in range(_PROMPT_ATTEMPTS):
+            messages = build_messages(
+                question,
+                sources,
+                evidence,
+                lineage,
+                n_retrieved=n_retrieved,
+                n_cited=n_cited,
+                n_lineage_sources=n_lineage_sources,
+                n_manual=n_manual,
+                budget_chars=budget,
+            )
+            user_len = len(messages[-1]["content"])
+            if attempt:
+                # Last resort when this is the final attempt or the trimming
+                # floors would not let the prompt get any shorter — a tighter
+                # budget would only bounce off the host again.
+                truncate = attempt == _PROMPT_ATTEMPTS - 1 or user_len >= sent_len
+                metrics.chat_prompt_overflow(
+                    "host_truncated" if truncate else "rebuilt"
+                )
+                log.warning(
+                    "prompt overflowed the context (%s tokens of %s) on %s; %s",
+                    prompt_tokens or "?",
+                    num_ctx or cfg.chat_num_ctx,
+                    host,
+                    (
+                        f"letting the host truncate ({user_len} chars of excerpts)"
+                        if truncate
+                        else f"rebuilt with a {budget}-char excerpt budget"
+                    ),
+                )
+            sent_len = user_len
+            with client.stream(
+                "POST",
+                f"{host}/api/chat",
+                json={
+                    "model": model,
+                    "messages": messages,
+                    "stream": True,
+                    "truncate": truncate,
+                    "options": {
+                        "num_ctx": cfg.chat_num_ctx,
+                        "temperature": cfg.chat_temperature,
+                        "num_predict": cfg.chat_max_tokens,
+                    },
                 },
-            },
-        ) as resp:
-            resp.raise_for_status()
-            for line in resp.iter_lines():
-                if not line:
-                    continue
-                data = json.loads(line)
-                chunk = data.get("message", {}).get("content", "")
-                if chunk:
-                    yield {"type": "token", "text": chunk}
-                if data.get("done"):
-                    # Ollama reports its token counts on the final frame;
-                    # older hosts/models omit them (chat_tokens skips 0).
-                    metrics.chat_tokens(
-                        int(data.get("prompt_eval_count") or 0),
-                        int(data.get("eval_count") or 0),
-                    )
+            ) as resp:
+                overflow = None if truncate else _overflow_counts(resp)
+                if overflow is None:
+                    resp.raise_for_status()
+                    yield from _relay_tokens(resp)
                     break
+            prompt_tokens, num_ctx = overflow
+            if prompt_tokens and num_ctx:
+                fit = num_ctx - cfg.chat_max_tokens - _PROMPT_TEMPLATE_RESERVE_TOKENS
+                budget = int(budget * fit / prompt_tokens * _OVERFLOW_MARGIN)
+            else:
+                budget = int(budget * _OVERFLOW_SHRINK)
     metrics.chat_stage("generate", time.monotonic() - began)
     yield {
         "type": "sources",

stack.toml

@@ -113,6 +113,7 @@ large_min_vram_gb = 20
 chat_num_ctx = 8192                    # passed per request; no Modelfile ctx variants
 chat_temperature = 0.2                 # grounded answers cite labels verbatim; low temperature keeps them literal
 chat_max_tokens = 1200                 # num_predict: a history answer ran to 5,000 tokens (3 min) without a cap
+chat_chars_per_token = 2.4             # measured 2.5–2.9 on FR/regulatory text; sizes the excerpt budget so prompt + reply fit chat_num_ctx (overflow is caught and rebuilt)
 embed_dim = 768
 # HNSW index build. Safe since the postgres image ships an AVX-512-free
 # pgvector rebuild (infra/images/postgresql.Dockerfile, #580); before that it

tests/llm/test_config.py

@@ -74,6 +74,7 @@ class TestNewKnobs:
         cfg = llm_config.load()
         assert cfg.large_min_vram_gb == 20.0
         assert cfg.chat_num_ctx == 8192
+        assert cfg.chat_chars_per_token == 2.4
         assert cfg.recency_half_life_days == 365.0
         assert cfg.recency_weight == 0.3
         assert cfg.k_per_kind == {"comment": 8, "rule": 4, "corpus": 4}

tests/llm/test_metrics.py

@@ -107,6 +107,7 @@ class TestInstrumentSet:
             "stack_llm_chat_seconds": "histogram",
             "stack_llm_chat_total": "counter",
             "stack_llm_chat_tokens_total": "counter",
+            "stack_llm_chat_prompt_overflow_total": "counter",
             "stack_llm_indexed_chunks_total": "counter",
             "stack_llm_indexed_items_total": "counter",
         }
@@ -117,7 +118,7 @@ class TestInstrumentSet:
     def test_instruments_built_once(self, fake_meter):
         first = metrics.instruments()
         assert metrics.instruments() is first
-        assert len(fake_meter.instruments) == 10
+        assert len(fake_meter.instruments) == 11
 
     def test_reset_rebuilds(self, fake_meter):
         first = metrics.instruments()
@@ -406,3 +407,13 @@ class TestNoopMeter:
         metrics.chat_outcome("ok")
         metrics.chat_tokens(1, 2)
         metrics.indexed_item("rules", "indexed", 4)
+
+
+class TestChatPromptOverflow:
+    def test_counts_by_resolution(self, fake_meter):
+        metrics.chat_prompt_overflow("rebuilt")
+        metrics.chat_prompt_overflow("host_truncated")
+        assert _calls(fake_meter, "stack_llm_chat_prompt_overflow_total") == [
+            (1, {"resolution": "rebuilt"}),
+            (1, {"resolution": "host_truncated"}),
+        ]

tests/llm/test_rag.py

@@ -13,10 +13,12 @@ import llm.rag as rag
 from llm.config import LlmConfig
 from llm.lineage import LineageEvent, LineageEvidence
 from llm.rag import (
+    _SYSTEM,
     build_messages,
     era_balance,
     era_of,
     is_history_question,
+    prompt_budget_chars,
     retrieve,
     stream_answer,
 )
@@ -1448,3 +1450,255 @@ class TestAcronymConflicts:
         assert [s["label"] for s in out] and all(
             "contractility" not in s["snippet"].lower() for s in out
         )
+
+
+class TestPromptBudget:
+    """The excerpt budget is derived from the token context, not a flat
+    multiple of it: the reply reservation, the chat template and the
+    system prompt all come out first, so system + user + reply fit
+    ``chat_num_ctx`` and the host never has to cut the prompt."""
+
+    def test_budget_leaves_room_for_system_prompt_and_reply(self):
+        from llm import rag
+
+        tokens = (
+            CFG.chat_num_ctx - CFG.chat_max_tokens - rag._PROMPT_TEMPLATE_RESERVE_TOKENS
+        )
+        assert prompt_budget_chars(CFG) == int(tokens * CFG.chat_chars_per_token) - len(
+            _SYSTEM
+        )
+        # strictly below the old flat 3× multiple that overflowed an 8k context
+        assert prompt_budget_chars(CFG) < CFG.chat_num_ctx * 3
+
+    def test_budget_scales_with_context_and_ratio(self):
+        from dataclasses import replace
+
+        wider = replace(CFG, chat_num_ctx=CFG.chat_num_ctx * 2)
+        assert prompt_budget_chars(wider) > prompt_budget_chars(CFG)
+        denser = replace(CFG, chat_chars_per_token=1.0)
+        assert prompt_budget_chars(denser) < prompt_budget_chars(CFG)
+
+    def test_budget_never_negative(self):
+        from dataclasses import replace
+
+        tiny = replace(CFG, chat_num_ctx=16, chat_max_tokens=8)
+        assert prompt_budget_chars(tiny) == 0
+
+
+def _overflow_resp(prompt_tokens: int | None, num_ctx: int = 8192):
+    """A mocked ``/api/chat`` 400 as Ollama returns it with ``truncate:
+    false`` — the current nested form with token counts, or (None) the
+    older bare message."""
+    resp = MagicMock()
+    resp.status_code = 400
+    if prompt_tokens is None:
+        body = {
+            "error": "the prompt is longer than the context length currently available to the model"
+        }
+    else:
+        inner = {
+            "error": {
+                "code": 400,
+                "message": f"request ({prompt_tokens} tokens) exceeds the available context size ({num_ctx} tokens), try increasing it",
+                "type": "exceed_context_size_error",
+                "n_prompt_tokens": prompt_tokens,
+                "n_ctx": num_ctx,
+            }
+        }
+        body = {"error": json.dumps(inner)}
+    resp.read.return_value = json.dumps(body).encode()
+    return resp
+
+
+def _ok_resp(
+    lines=(
+        '{"message":{"content":"ok"},"done":false}',
+        '{"message":{"content":""},"done":true,"prompt_eval_count":5000,"eval_count":1}',
+    ),
+):
+    resp = MagicMock()
+    resp.status_code = 200
+    resp.iter_lines.return_value = iter(lines)
+    return resp
+
+
+class TestStreamAnswerPromptOverflow:
+    """The chat sends ``truncate: false`` and sizes the prompt itself; a
+    host 400 for an over-long prompt rebuilds the messages under a
+    tighter budget (from the host's own token counts when it reports
+    them) and resends — host-side truncation is the last resort only."""
+
+    def _pool(self):
+        pool = MagicMock()
+        pool.acquire_generation.return_value.__enter__.return_value = "http://h1:11434"
+        pool.vram.return_value = 24.0
+        pool.serves.return_value = True
+        return pool
+
+    def _src(self, i):
+        return {
+            "id": f"C{i}",
+            "label": f"C{i}",
+            "kind": "comment",
+            "snippet": "x" * 400,
+            "score": 0.1,
+        }
+
+    @patch("llm.rag.metrics.chat_prompt_overflow")
+    @patch("llm.rag.valuation_evidence", return_value=None)
+    @patch("llm.rag.lineage_evidence", return_value=None)
+    @patch("llm.rag.httpx.Client")
+    @patch("llm.rag.retrieve")
+    def test_first_attempt_uses_token_budget_and_refuses_host_truncation(
+        self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
+    ):
+        mock_retrieve.return_value = [self._src(0)]
+        client = MockClient.return_value.__enter__.return_value
+        client.stream.return_value.__enter__.return_value = _ok_resp()
+
+        with patch("llm.rag.build_messages", wraps=build_messages) as bm:
+            events = list(stream_answer("q", cfg=CFG, pool=self._pool()))
+
+        assert events[0] == {"type": "token", "text": "ok"}
+        assert bm.call_args.kwargs["budget_chars"] == prompt_budget_chars(CFG)
+        body = client.stream.call_args.kwargs["json"]
+        assert body["truncate"] is False
+        assert client.stream.call_count == 1
+        mock_overflow.assert_not_called()
+
+    @patch("llm.rag.metrics.chat_prompt_overflow")
+    @patch("llm.rag.valuation_evidence", return_value=None)
+    @patch("llm.rag.lineage_evidence", return_value=None)
+    @patch("llm.rag.httpx.Client")
+    @patch("llm.rag.retrieve")
+    def test_overflow_rebuilds_tighter_from_host_counts_then_streams(
+        self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
+    ):
+        from llm import rag
+
+        mock_retrieve.return_value = [self._src(i) for i in range(30)]
+        client = MockClient.return_value.__enter__.return_value
+        client.stream.return_value.__enter__.side_effect = [
+            _overflow_resp(prompt_tokens=9500, num_ctx=8192),
+            _ok_resp(),
+        ]
+
+        with patch("llm.rag.build_messages", wraps=build_messages) as bm:
+            events = list(stream_answer("q", cfg=CFG, pool=self._pool()))
+
+        assert events[0] == {"type": "token", "text": "ok"}
+        assert events[-1] == {"type": "done"}
+        assert client.stream.call_count == 2
+        first, second = (c.kwargs["budget_chars"] for c in bm.call_args_list)
+        fit = 8192 - CFG.chat_max_tokens - rag._PROMPT_TEMPLATE_RESERVE_TOKENS
+        assert first == prompt_budget_chars(CFG)
+        assert second == int(first * fit / 9500 * rag._OVERFLOW_MARGIN)
+        assert second < first
+        # the second request still refuses host-side truncation
+        assert client.stream.call_args_list[1].kwargs["json"]["truncate"] is False
+        # the resent prompt really is smaller
+        sent = [
+            c.kwargs["json"]["messages"][1]["content"]
+            for c in client.stream.call_args_list
+        ]
+        assert len(sent[1]) < len(sent[0])
+        mock_overflow.assert_called_once_with("rebuilt")
+
+    @patch("llm.rag.metrics.chat_prompt_overflow")
+    @patch("llm.rag.valuation_evidence", return_value=None)
+    @patch("llm.rag.lineage_evidence", return_value=None)
+    @patch("llm.rag.httpx.Client")
+    @patch("llm.rag.retrieve")
+    def test_bare_message_overflow_shrinks_by_fixed_factor(
+        self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
+    ):
+        from llm import rag
+
+        mock_retrieve.return_value = [self._src(i) for i in range(30)]
+        client = MockClient.return_value.__enter__.return_value
+        client.stream.return_value.__enter__.side_effect = [
+            _overflow_resp(prompt_tokens=None),
+            _ok_resp(),
+        ]
+
+        with patch("llm.rag.build_messages", wraps=build_messages) as bm:
+            list(stream_answer("q", cfg=CFG, pool=self._pool()))
+
+        first, second = (c.kwargs["budget_chars"] for c in bm.call_args_list)
+        assert second == int(first * rag._OVERFLOW_SHRINK)
+        mock_overflow.assert_called_once_with("rebuilt")
+
+    @patch("llm.rag.metrics.chat_prompt_overflow")
+    @patch("llm.rag.valuation_evidence", return_value=None)
+    @patch("llm.rag.lineage_evidence", return_value=None)
+    @patch("llm.rag.httpx.Client")
+    @patch("llm.rag.retrieve")
+    def test_repeated_overflow_falls_back_to_host_truncation(
+        self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
+    ):
+        mock_retrieve.return_value = [self._src(i) for i in range(30)]
+        client = MockClient.return_value.__enter__.return_value
+        client.stream.return_value.__enter__.side_effect = [
+            _overflow_resp(9500),
+            _overflow_resp(9000),
+            _ok_resp(),
+        ]
+
+        events = list(stream_answer("q", cfg=CFG, pool=self._pool()))
+
+        assert events[-1] == {"type": "done"}
+        flags = [c.kwargs["json"]["truncate"] for c in client.stream.call_args_list]
+        assert flags == [False, False, True]
+        assert [c.args[0] for c in mock_overflow.call_args_list] == [
+            "rebuilt",
+            "host_truncated",
+        ]
+
+    @patch("llm.rag.metrics.chat_prompt_overflow")
+    @patch("llm.rag.valuation_evidence", return_value=None)
+    @patch("llm.rag.lineage_evidence", return_value=None)
+    @patch("llm.rag.httpx.Client")
+    @patch("llm.rag.retrieve")
+    def test_overflow_at_the_trimming_floors_goes_straight_to_host_truncation(
+        self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
+    ):
+        # One excerpt: the floors keep it however small the budget gets, so a
+        # second truncate:false attempt would only bounce again.
+        mock_retrieve.return_value = [self._src(0)]
+        client = MockClient.return_value.__enter__.return_value
+        client.stream.return_value.__enter__.side_effect = [
+            _overflow_resp(9500),
+            _ok_resp(),
+        ]
+
+        events = list(stream_answer("q", cfg=CFG, pool=self._pool()))
+
+        assert events[-1] == {"type": "done"}
+        flags = [c.kwargs["json"]["truncate"] for c in client.stream.call_args_list]
+        assert flags == [False, True]
+        mock_overflow.assert_called_once_with("host_truncated")
+
+    @patch("llm.rag.metrics.chat_prompt_overflow")
+    @patch("llm.rag.valuation_evidence", return_value=None)
+    @patch("llm.rag.lineage_evidence", return_value=None)
+    @patch("llm.rag.httpx.Client")
+    @patch("llm.rag.retrieve")
+    def test_unrelated_400_still_raises(
+        self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
+    ):
+        import httpx
+
+        mock_retrieve.return_value = [self._src(0)]
+        resp = MagicMock()
+        resp.status_code = 400
+        resp.read.return_value = json.dumps({"error": "model 'big' not found"}).encode()
+        resp.raise_for_status.side_effect = httpx.HTTPStatusError(
+            "400", request=MagicMock(), response=MagicMock()
+        )
+        client = MockClient.return_value.__enter__.return_value
+        client.stream.return_value.__enter__.return_value = resp
+
+        with pytest.raises(httpx.HTTPStatusError):
+            list(stream_answer("q", cfg=CFG, pool=self._pool()))
+        assert client.stream.call_count == 1
+        mock_overflow.assert_not_called()