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

2026-08-22 11:30:19 -04:00
parent 6ec6181c7b
commit 0da5275eeb
9 changed files with 295 additions and 5 deletions
@@ -9,10 +9,16 @@ Implements just enough of the aipi surface:
  unrelated ones score low and trigger honest deflection.
 * ``POST /v1/chat/completions`` — streaming (SSE) or not. The content keys
  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
  - otherwise         -> upbeat answer quoting the provided document context
  - user message containing ``pretend to think slowly`` -> 3s warm-up delay
    (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

@@ -63,9 +69,31 @@ def _context(body: dict[str, Any]) -> str:
    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:
    system = _system(body)
    user = _user(body)
+    if LONG_ANSWER_TRIGGER in user.lower():
+        return long_answer()
    if "DEFLECT_MODE" in system:
        return (
            "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)


+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")
 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

    if not body.get("stream"):
@@ -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()