diff --git a/src/llm/config.py b/src/llm/config.py index bf12db4..59077c2 100644 --- a/src/llm/config.py +++ b/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, diff --git a/src/llm/metrics.py b/src/llm/metrics.py index bf3036d..549152f 100644 --- a/src/llm/metrics.py +++ b/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() diff --git a/src/llm/rag.py b/src/llm/rag.py index e13be63..1653221 100644 --- a/src/llm/rag.py +++ b/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", diff --git a/stack.toml b/stack.toml index 44196bf..bf5ac17 100644 --- a/stack.toml +++ b/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 diff --git a/tests/llm/test_config.py b/tests/llm/test_config.py index 970e1fc..caa0024 100644 --- a/tests/llm/test_config.py +++ b/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} diff --git a/tests/llm/test_metrics.py b/tests/llm/test_metrics.py index 3a4c7c2..c04daa0 100644 --- a/tests/llm/test_metrics.py +++ b/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"}), + ] diff --git a/tests/llm/test_rag.py b/tests/llm/test_rag.py index 121c610..791be10 100644 --- a/tests/llm/test_rag.py +++ b/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()