fix(rag): lift chat output cap to 32768 tokens — long answers no longer cut off

This commit is contained in:
2026-08-22 11:30:19 -04:00
parent 6ec6181c7b
commit 0da5275eeb
9 changed files with 295 additions and 5 deletions
+52
View File
@@ -0,0 +1,52 @@
# Phase 11 — Long Answers (No Truncation)
**Story:** `.agent/user_stories/long-answers.md`
**Context:** owner report 2026-08-22 — "responses keep getting cut off.
It should be allowed to respond up to 32768 tokens."
## Goal
Remove the hard 700-token output cap on chat answers; the model may
respond up to **32 768** tokens (`BOR_MAX_OUTPUT_TOKENS`, default
32 768).
## Diagnosis
`app/rag/llm.py::chat_stream` calls `chat.completions.create(...,
max_tokens=700, ...)`. Long answers die mid-sentence at ~700 tokens.
## Implementation steps
1. **Config** (`app/config.py`): `max_output_tokens: int = 32_768` in the
RAG-tuning section (env `BOR_MAX_OUTPUT_TOKENS`); document in
`.env.example`.
2. **LLM client** (`app/rag/llm.py`): `chat_stream` passes
`max_tokens=self.settings.max_output_tokens`.
3. **E2E mock** (`tests/e2e/mock_llm.py`):
- Honor `max_tokens` deterministically: token ≈ whitespace word; if
the composed answer is longer, truncate to the first N words.
(With the old 700 cap a long answer loses its tail — the mock now
behaves like the real endpoint.)
- New trigger: user message containing `write a long answer` →
deterministic ~4 000-word numbered answer ending in a unique final
line (`LONG-ANSWER-END`).
- No behavior change for existing (short) answers: they fit under any
sane cap.
4. **Tests:**
- Unit: settings default + env override (`test_config.py`);
`chat_stream` forwards the configured `max_tokens` (fake client in
`test_llm_client.py`).
- E2E: `tests/e2e/test_long_answers.py` per the story mapping.
## Locked decisions
None touched. A5 (aipi endpoint) unchanged; `turbo` accepts the larger
cap per owner instruction.
## Testing & Quality
- Unit + integration green; `uv run pytest --cov=app --cov-report=term-missing`
**>90%**; E2E in isolation:
`uv run pytest tests/e2e/test_long_answers.py -v --no-cov`.
- No regressions: `test_chat_rag.py` + `test_chat_api.py` green
(short answers unaffected by the mock's new `max_tokens` honoring).
## Commit
```bash
git add -A .agent/ app/ tests/ .env.example && git commit --no-gpg-sign -m "fix(rag): lift chat output cap to 32768 tokens — long answers no longer cut off"
```
+34
View File
@@ -0,0 +1,34 @@
# Story: Long Answers (No Truncation)
**Phase:** `11_long_answers.md` · **E2E:** `tests/e2e/test_long_answers.py`
## Narrative
As **a user**, I want Brain to be able to answer at full length (up to
32 768 output tokens) so complex questions ("walk me through the whole
setup", "list every service and its config") get a **complete** answer
instead of one that stops mid-sentence.
- **Given** any question that deserves a long answer
- **When** Brain streams its reply
- **Then** the reply runs to its natural end — the model is allowed up to
32 768 output tokens, not a hard 700-token cap.
## Acceptance criteria
1. `LLMClient.chat_stream` sends `max_tokens` from settings
(`BOR_MAX_OUTPUT_TOKENS`, default **32 768**) — the hard-coded 700 is
gone.
2. A genuinely long streamed answer (several thousand words) arrives
**complete** in the browser — final line intact (E2E).
3. Setting is overridable via env; unit-tested.
4. Unit + integration green, `app/` coverage >90%, story E2E green in
isolation, one `--no-gpg-sign` commit.
## Playwright Mapping Rule
**Test Scenario → `tests/e2e/test_long_answers.py`** (mock LLM, seeded KB):
1. `test_long_answer_streams_to_completion` — question with the
"write a long answer" trigger → mock emits a ~4 000-word deterministic
answer and **honors `max_tokens`** (word-based) → the browser shows the
final line of the answer; under the old 700 cap the tail is missing.
2. `test_normal_answer_unaffected` — a normal question still streams a
complete, short answer.
+1
View File
@@ -21,6 +21,7 @@ BOR_EMBEDDING_DIM=768 # verified 2026-08 via scripts/llm_probe.py
BOR_TOP_N_DOCS=2 BOR_TOP_N_DOCS=2
BOR_RELEVANCE_THRESHOLD=0.62 # answer when best cosine >= this OR an FTS hit; else honest deflection BOR_RELEVANCE_THRESHOLD=0.62 # answer when best cosine >= this OR an FTS hit; else honest deflection
BOR_MAX_CONTEXT_CHARS=24000 # cap on total document text sent to the LLM BOR_MAX_CONTEXT_CHARS=24000 # cap on total document text sent to the LLM
BOR_MAX_OUTPUT_TOKENS=32768 # max answer length in tokens (answers must not be cut off)
BOR_CHUNK_TARGET_CHARS=2000 BOR_CHUNK_TARGET_CHARS=2000
BOR_CHUNK_OVERLAP_CHARS=200 BOR_CHUNK_OVERLAP_CHARS=200
BOR_EMBED_BATCH_SIZE=16 BOR_EMBED_BATCH_SIZE=16
+4
View File
@@ -51,6 +51,10 @@ class Settings(BaseSettings):
# AND no candidate chunk matches the question lexically (see A8). # AND no candidate chunk matches the question lexically (see A8).
relevance_threshold: float = 0.62 relevance_threshold: float = 0.62
max_context_chars: int = 24_000 max_context_chars: int = 24_000
#: Maximum output tokens a chat answer may use (owner instruction
#: 2026-08-22: answers must run to their natural end — the old hard
#: 700-token cap cut long answers off mid-sentence).
max_output_tokens: int = 32_768
chunk_target_chars: int = 2_000 chunk_target_chars: int = 2_000
chunk_overlap_chars: int = 200 chunk_overlap_chars: int = 200
embed_batch_size: int = 16 embed_batch_size: int = 16
+5 -1
View File
@@ -175,6 +175,10 @@ class LLMClient:
non-empty ``delta.content`` pieces. Any failure (network, HTTP, non-empty ``delta.content`` pieces. Any failure (network, HTTP,
malformed stream) surfaces as :class:`LLMError` so the API layer can malformed stream) surfaces as :class:`LLMError` so the API layer can
turn it into an SSE ``error`` event instead of a hung request. turn it into an SSE ``error`` event instead of a hung request.
Answers are allowed up to ``BOR_MAX_OUTPUT_TOKENS`` (default 32 768)
output tokens — the old hard 700-token cap cut long answers off
mid-sentence (owner report 2026-08-22).
""" """
try: try:
# ``{role, content}`` dicts are exactly what the message params # ``{role, content}`` dicts are exactly what the message params
@@ -183,7 +187,7 @@ class LLMClient:
model=self.settings.llm_chat_model, model=self.settings.llm_chat_model,
messages=cast("list[ChatCompletionMessageParam]", messages), messages=cast("list[ChatCompletionMessageParam]", messages),
temperature=0.4, temperature=0.4,
max_tokens=700, max_tokens=self.settings.max_output_tokens,
stream=True, stream=True,
) )
async for chunk in stream: async for chunk in stream:
+41 -1
View File
@@ -9,10 +9,16 @@ Implements just enough of the aipi surface:
unrelated ones score low and trigger honest deflection. unrelated ones score low and trigger honest deflection.
* ``POST /v1/chat/completions`` — streaming (SSE) or not. The content keys * ``POST /v1/chat/completions`` — streaming (SSE) or not. The content keys
off markers in the system prompt: off markers in the system prompt:
- user message containing ``write a long answer`` -> a ~900-word
deterministic numbered answer (long-answers story, phase 11)
- ``DEFLECT_MODE`` -> honest "I haven't done anything like that" answer - ``DEFLECT_MODE`` -> honest "I haven't done anything like that" answer
- otherwise -> upbeat answer quoting the provided document context - otherwise -> upbeat answer quoting the provided document context
- user message containing ``pretend to think slowly`` -> 3s warm-up delay - user message containing ``pretend to think slowly`` -> 3s warm-up delay
(used by the loading-feedback story). (used by the loading-feedback story).
``max_tokens`` is honored deterministically (token ≈ whitespace word),
like a real endpoint: an answer longer than the cap is truncated. This
is what makes the phase-11 truncation regression observable.
""" """
from __future__ import annotations from __future__ import annotations
@@ -63,9 +69,31 @@ def _context(body: dict[str, Any]) -> str:
return max((m.get("content", "") for m in msgs), key=len) return max((m.get("content", "") for m in msgs), key=len)
LONG_ANSWER_TRIGGER = "write a long answer"
#: ~920 words — comfortably past the old hard 700-token cap (where the
#: tail would be cut) yet short enough to stream in ~8s at the mock's
#: per-chunk pacing.
LONG_ANSWER_LINES = 40
LONG_ANSWER_END = "LONG-ANSWER-END"
def long_answer() -> str:
"""~900-word deterministic walkthrough (phase 11): numbered steps plus
a unique final line that must survive the stream untruncated."""
lines = [
f"{i}. Step {i}: configure node-{i} with the homelab defaults and "
f"verify that step {i} of the long walkthrough is complete before moving on."
for i in range(1, LONG_ANSWER_LINES + 1)
]
lines.append(LONG_ANSWER_END)
return "\n".join(lines)
def compose_answer(body: dict[str, Any]) -> str: def compose_answer(body: dict[str, Any]) -> str:
system = _system(body) system = _system(body)
user = _user(body) user = _user(body)
if LONG_ANSWER_TRIGGER in user.lower():
return long_answer()
if "DEFLECT_MODE" in system: if "DEFLECT_MODE" in system:
return ( return (
"Ah — I haven't done anything like that, so I don't want to make stuff up! " "Ah — I haven't done anything like that, so I don't want to make stuff up! "
@@ -163,9 +191,21 @@ def json_dumps(obj: dict[str, Any]) -> str:
return json.dumps(obj) return json.dumps(obj)
def _apply_max_tokens(answer: str, max_tokens: Any) -> str:
"""Deterministic stand-in for the endpoint's output cap: one token ≈
one whitespace-separated word. Answers within the cap pass through
byte-identical, so existing (short) answers are unaffected."""
if not isinstance(max_tokens, int) or max_tokens <= 0:
return answer
words = answer.split()
if len(words) <= max_tokens:
return answer
return " ".join(words[:max_tokens])
@app.post("/v1/chat/completions") @app.post("/v1/chat/completions")
def chat_completions(body: dict[str, Any]) -> Any: def chat_completions(body: dict[str, Any]) -> Any:
answer = compose_answer(body) answer = _apply_max_tokens(compose_answer(body), body.get("max_tokens"))
delay = 3.0 if "pretend to think slowly" in _user(body) else 0.0 delay = 3.0 if "pretend to think slowly" in _user(body) else 0.0
if not body.get("stream"): if not body.get("stream"):
+133
View File
@@ -0,0 +1,133 @@
"""Phase 11 E2E (Playwright): long answers stream to completion.
Story: ``.agent/user_stories/long-answers.md``
Run in isolation (DB must be up: ``podman compose up -d db``):
uv run pytest tests/e2e/test_long_answers.py -v --no-cov
The mock LLM honors ``max_tokens`` (token ≈ word) like a real endpoint,
and emits a ~900-word deterministic answer for the "write a long answer"
trigger. Under the old hard 700-token cap the answer loses its tail
(the final line never arrives); with ``BOR_MAX_OUTPUT_TOKENS`` defaulting
to 32 768 the full answer streams to completion.
Test → story mapping (Playwright Mapping Rule):
1. ``test_long_answer_streams_to_completion`` — trigger question →
full ~900-word answer, final line intact, >700 words rendered.
2. ``test_normal_answer_unaffected`` — a regular question still streams
a complete short answer.
"""
from __future__ import annotations
import asyncio
from pathlib import Path
from threading import Thread
from typing import Any
from playwright.sync_api import Page, expect
from sqlalchemy import text
from app.config import Settings
from app.db import SessionLocal
from app.rag.importer import ImportSummary, import_sources
from app.rag.llm import LLMClient
REPO = Path(__file__).resolve().parents[2]
FIXTURES = REPO / "tests" / "fixtures" / "docs"
LONG_QUESTION = "How is my Kubernetes cluster set up? write a long answer"
NORMAL_QUESTION = "How is my Kubernetes cluster set up?"
async def _import_fixtures(mock_port: int) -> ImportSummary:
kwargs: dict[str, Any] = {"_env_file": None, "llm_base_url": f"http://127.0.0.1:{mock_port}/v1"}
settings = Settings(**kwargs) # pyright: ignore[reportCallIssue]
return await import_sources([FIXTURES], LLMClient(settings))
def _run_in_thread(coro: Any) -> Any:
"""Run a coroutine on a worker thread (Playwright owns the test loop)."""
box: dict[str, Any] = {}
def runner() -> None:
try:
box["value"] = asyncio.run(coro)
except BaseException as e: # noqa: BLE001 — re-raised on the test thread
box["error"] = e
t = Thread(target=runner)
t.start()
t.join()
if "error" in box:
raise box["error"]
return box["value"]
def _reset_db(mock_port: int) -> ImportSummary:
with SessionLocal() as db:
db.execute(text("TRUNCATE chunks, documents, query_log"))
db.commit()
return _run_in_thread(_import_fixtures(mock_port))
def _last_brain_text(page: Page) -> str:
return page.locator(".msg.brain .bubble").last.inner_text()
# ---------------------------------------------------------------------------
# 1. Long answer: the final line must survive the stream
# ---------------------------------------------------------------------------
def test_long_answer_streams_to_completion(
page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
_reset_db(mock_llm)
page.set_default_timeout(45_000)
page.goto(app_url)
page.fill("#message-input", LONG_QUESTION)
page.click("#send-btn")
# The mock streams ~900 words in ~8s; wait for the unique final line —
# under the old 700-token cap it was cut off and never arrived.
expect(page.locator(".msg.brain .bubble").last).to_contain_text(
"LONG-ANSWER-END", timeout=60_000
)
full = _last_brain_text(page)
# The old cap would have stopped the answer at 700 words — prove the
# rendered answer ran well past it.
assert len(full.split()) > 700, (
f"answer looks truncated at {len(full.split())} words"
)
# First and last step both rendered (the markdown list strips the
# "1." prefix — no mid-sentence cut between them either).
assert "Step 1:" in full
assert "Step 40:" in full
# Turn settled: send button re-enabled (never-stale contract).
expect(page.locator("#send-btn")).to_be_enabled()
# ---------------------------------------------------------------------------
# 2. Normal (short) answers are unaffected by the raised cap
# ---------------------------------------------------------------------------
def test_normal_answer_unaffected(
page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
_reset_db(mock_llm)
page.set_default_timeout(30_000)
page.goto(app_url)
page.fill("#message-input", NORMAL_QUESTION)
page.click("#send-btn")
bubble = page.locator(".msg.brain .bubble").last
expect(bubble).to_contain_text("Deterministic mock answer for E2E", timeout=30_000)
expect(bubble).not_to_contain_text("LONG-ANSWER-END")
# Grounded: the question's own document is cited as a chip.
expect(
page.locator(".msg.brain .source-chip", has_text="kubernetes.md")
).to_have_count(1, timeout=30_000)
expect(page.locator("#send-btn")).to_be_enabled()
+8
View File
@@ -34,6 +34,8 @@ def test_defaults_match_locked_decisions(monkeypatch: pytest.MonkeyPatch) -> Non
assert s.hybrid_lexical_candidates >= 1 assert s.hybrid_lexical_candidates >= 1
assert s.rrf_k >= 1 assert s.rrf_k >= 1
assert s.top_n_docs >= 1 assert s.top_n_docs >= 1
# Owner instruction 2026-08-22: answers may run up to 32 768 tokens.
assert s.max_output_tokens == 32_768
assert len(s.suggestions) >= 3 assert len(s.suggestions) >= 3
# A9 (revised): the import scope covers the seven A9 formats. # A9 (revised): the import scope covers the seven A9 formats.
assert s.import_extension_set == { assert s.import_extension_set == {
@@ -49,6 +51,12 @@ def test_env_override(monkeypatch) -> None:
assert s.llm_chat_model == "juggernaut" assert s.llm_chat_model == "juggernaut"
def test_max_output_tokens_env_override(monkeypatch) -> None:
monkeypatch.setenv("BOR_MAX_OUTPUT_TOKENS", "1234")
s = _settings()
assert s.max_output_tokens == 1234
def test_import_extensions_env_override_is_a_csv_list(monkeypatch) -> None: def test_import_extensions_env_override_is_a_csv_list(monkeypatch) -> None:
monkeypatch.setenv("BOR_IMPORT_EXTENSIONS", "md,yml") monkeypatch.setenv("BOR_IMPORT_EXTENSIONS", "md,yml")
s = _settings() s = _settings()
+17 -3
View File
@@ -273,11 +273,13 @@ class _FakeCompletions:
def _make_stream_client( def _make_stream_client(
chunks: list | None = None, fail: Exception | None = None chunks: list | None = None,
fail: Exception | None = None,
**settings_kwargs: Any,
) -> tuple[LLMClient, _FakeCompletions]: ) -> tuple[LLMClient, _FakeCompletions]:
completions = _FakeCompletions(chunks, fail) completions = _FakeCompletions(chunks, fail)
fake_openai = SimpleNamespace(chat=SimpleNamespace(completions=completions)) fake_openai = SimpleNamespace(chat=SimpleNamespace(completions=completions))
llm = LLMClient(_settings()) llm = LLMClient(_settings(**settings_kwargs))
llm._client = fake_openai # pyright: ignore[reportAttributeAccessIssue] llm._client = fake_openai # pyright: ignore[reportAttributeAccessIssue]
return llm, completions return llm, completions
@@ -302,10 +304,22 @@ def test_chat_stream_uses_locked_generation_params() -> None:
assert completions.kwargs["model"] == "turbo" assert completions.kwargs["model"] == "turbo"
assert completions.kwargs["stream"] is True assert completions.kwargs["stream"] is True
assert completions.kwargs["temperature"] == 0.4 assert completions.kwargs["temperature"] == 0.4
assert completions.kwargs["max_tokens"] == 700 # Phase 11: the old hard 700-token cap is gone — answers may run up to
# BOR_MAX_OUTPUT_TOKENS (default 32 768) so they are not cut off.
assert completions.kwargs["max_tokens"] == 32_768
assert completions.kwargs["messages"] == messages assert completions.kwargs["messages"] == messages
def test_chat_stream_max_tokens_comes_from_settings() -> None:
"""The output cap is operator-configurable, not a client constant."""
llm, completions = _make_stream_client(
[_chunk("x")], max_output_tokens=1234 # pyright: ignore[reportArgumentType]
)
asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))
assert completions.kwargs is not None
assert completions.kwargs["max_tokens"] == 1234
def test_chat_stream_skips_empty_deltas_and_choiceless_chunks() -> None: def test_chat_stream_skips_empty_deltas_and_choiceless_chunks() -> None:
llm, _ = _make_stream_client([_chunk("a"), _chunk(empty=True), _chunk(None), _chunk("b")]) llm, _ = _make_stream_client([_chunk("a"), _chunk(empty=True), _chunk(None), _chunk("b")])
assert asyncio.run(_collect(llm, [{"role": "user", "content": "q"}])) == ["a", "b"] assert asyncio.run(_collect(llm, [{"role": "user", "content": "q"}])) == ["a", "b"]