feat(chat): stream model thinking over SSE and show it in a collapsible block
This commit is contained in:
+65
-7
@@ -15,6 +15,9 @@ Implements just enough of the aipi surface:
|
||||
- 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).
|
||||
- user message containing ``think out loud`` -> the answer is preceded by
|
||||
~800 chars of deterministic ``reasoning_content`` chunks (the
|
||||
thinking-display story, phase 17).
|
||||
- system prompt containing ``<tuning>`` (phase 15, steering notes) ->
|
||||
the composed answer ends with `` (tuning: <first note line>)`` —
|
||||
makes prompt injection observable in the UI deterministically.
|
||||
@@ -79,6 +82,13 @@ LONG_ANSWER_TRIGGER = "write a long answer"
|
||||
LONG_ANSWER_LINES = 40
|
||||
LONG_ANSWER_END = "LONG-ANSWER-END"
|
||||
|
||||
#: Phase 17 (thinking-display story): a user message containing this
|
||||
#: substring (case-insensitive) is answered with a deterministic
|
||||
#: ``reasoning_content`` stream ahead of the content — same convention as
|
||||
#: the other user-message triggers above. Existing E2E questions do not
|
||||
#: contain the substring, so every other suite is unaffected.
|
||||
THINKING_TRIGGER = "think out loud"
|
||||
|
||||
|
||||
def long_answer() -> str:
|
||||
"""~900-word deterministic walkthrough (phase 11): numbered steps plus
|
||||
@@ -142,6 +152,32 @@ def compose_answer(body: dict[str, Any]) -> str:
|
||||
return answer
|
||||
|
||||
|
||||
def compose_thinking(body: dict[str, Any]) -> str:
|
||||
"""Deterministic reasoning scratchpad (thinking-display story, phase 17).
|
||||
|
||||
A fixed 4-line "Step 1… Step 4" template quoting the first ~60 chars
|
||||
of the user question: unique per question, byte-stable across runs,
|
||||
~700–900 chars total (≈ 60–75 frames at the mock's 12-char/0.02s
|
||||
pacing). The ``Step 2: Check my notes`` line fragment is what the E2E
|
||||
assertions key off.
|
||||
"""
|
||||
q = _user(body).strip()[:60]
|
||||
return (
|
||||
f"Step 1: Read the question carefully — “{q}” — and figure out what kind of "
|
||||
"answer it wants (a how-to, a lookup, or a design decision) before touching "
|
||||
"the docs, so I don't over- or under-answer.\n"
|
||||
"Step 2: Check my notes for the closest match. The homelab kubernetes file "
|
||||
"is the obvious candidate, but I should also consider whether a deployments "
|
||||
"note covers the same ground better.\n"
|
||||
"Step 3: Re-read the relevant sections top to bottom so every specific — "
|
||||
"hosts, versions, ports, schedules — is exact as written rather than "
|
||||
"remembered, and note which document each fact comes from.\n"
|
||||
"Step 4: Draft the answer around those specifics, keep it tight with short "
|
||||
"paragraphs and bullets where it helps, cite the documents by path, and "
|
||||
"double-check that nothing is invented."
|
||||
)
|
||||
|
||||
|
||||
@app.post("/__shutdown__")
|
||||
def shutdown() -> dict[str, Any]:
|
||||
"""Test hook (loading-feedback story): terminate this mock process to
|
||||
@@ -186,11 +222,31 @@ def embeddings(body: dict[str, Any]) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _sse_stream(answer: str, delay: float) -> Any:
|
||||
def _sse_stream(answer: str, delay: float, thinking: str = "") -> Any:
|
||||
"""SSE frames for one chat completion (phase 17: + reasoning).
|
||||
|
||||
When ``thinking`` is non-empty its 12-char slices go out FIRST as
|
||||
``delta.reasoning_content`` frames — same 0.02s cadence and envelope
|
||||
as the content frames, the aipi wire convention (reasoning before
|
||||
content). Without ``thinking`` the output is byte-identical to the
|
||||
content-only stream, so the other story suites are unaffected.
|
||||
"""
|
||||
model = "turbo"
|
||||
chunk_id = f"chatcmpl-{uuid.uuid4()}"
|
||||
if delay:
|
||||
time.sleep(delay)
|
||||
for piece in re.findall(r".{1,12}", thinking, re.S):
|
||||
payload = {
|
||||
"id": chunk_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [
|
||||
{"index": 0, "delta": {"reasoning_content": piece}, "finish_reason": None}
|
||||
],
|
||||
}
|
||||
yield f"data: {json_dumps(payload)}\n\n"
|
||||
time.sleep(0.02)
|
||||
for piece in re.findall(r".{1,12}", answer, re.S):
|
||||
payload = {
|
||||
"id": chunk_id,
|
||||
@@ -239,25 +295,27 @@ def _apply_max_tokens(answer: str, max_tokens: Any) -> str:
|
||||
def chat_completions(body: dict[str, Any]) -> Any:
|
||||
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
|
||||
thinking = compose_thinking(body) if THINKING_TRIGGER in _user(body).lower() else ""
|
||||
|
||||
if not body.get("stream"):
|
||||
message: dict[str, Any] = {"role": "assistant", "content": answer}
|
||||
if thinking:
|
||||
# Harmless future-proofing: the app only uses streaming, but a
|
||||
# non-streaming client that reads the field gets the reasoning.
|
||||
message["reasoning_content"] = thinking
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4()}",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": body.get("model", "turbo"),
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": answer},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
{"index": 0, "message": message, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
|
||||
return StreamingResponse(
|
||||
_sse_stream(answer, delay),
|
||||
_sse_stream(answer, delay, thinking=thinking),
|
||||
media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
||||
)
|
||||
|
||||
@@ -0,0 +1,281 @@
|
||||
"""Phase 17 E2E (Playwright, mock-only): the model's "thinking" display.
|
||||
|
||||
Story: ``.agent/user_stories/thinking-display.md``
|
||||
Run in isolation (DB must be up: ``podman compose up -d db``):
|
||||
|
||||
uv run pytest tests/e2e/test_thinking_display.py -v --no-cov
|
||||
|
||||
MOCK-ONLY suite: ``E2E_REAL_LLM=1`` is not supported here. The real
|
||||
``turbo`` thinks on *every* turn, which would break the no-thinking
|
||||
regression test (scenario 4) — the deterministic mock's ``think out loud``
|
||||
trigger (mock_llm.py) keeps all five scenarios reproducible.
|
||||
|
||||
Test → story mapping (Playwright Mapping Rule):
|
||||
1. ``test_thinking_block_streams_open_then_collapses``
|
||||
2. ``test_thinking_toggle_after_done``
|
||||
3. ``test_thinking_restored_after_reload``
|
||||
4. ``test_no_thinking_block_without_trigger``
|
||||
5. ``test_thinking_with_deflection``
|
||||
|
||||
Determinism note: the mock paces every SSE frame at 0.02s and the thinking
|
||||
text is ~700–900 chars (≈ 60–75 frames ≈ 1.2–1.5s) before the first
|
||||
content frame, so "attach → assert open" runs well inside the open window
|
||||
on headless Chromium; all other assertions are made after the send button
|
||||
re-enables (fully settled state).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from threading import Thread
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
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"
|
||||
THINK_QUESTION = "think out loud — how is my kubernetes cluster set up?"
|
||||
PLAIN_QUESTION = "How is my Kubernetes cluster set up?"
|
||||
THINK_DEFLECT_QUESTION = "think out loud — tell me about quantum wormhole cooling"
|
||||
MOCK_ANSWER_MARKER = "Deterministic mock answer for E2E"
|
||||
DEFLECT_PHRASE = r"haven't done anything like that"
|
||||
#: Line fragment the mock's deterministic scratchpad must carry — the
|
||||
#: suite keys off it (mock_llm.compose_thinking).
|
||||
THINKING_FRAGMENT = "Step 2: Check my notes"
|
||||
STORAGE_KEY = "bor.chat.v1"
|
||||
#: Phase-10 viewer URL + phase-13 back=/ (byte-identical to the chip the
|
||||
#: persistence suite pins — grounded-turn sources are unchanged by 17).
|
||||
CHIP_HREF = "/document.html?source=docs&path=homelab%2Fkubernetes.md&back=%2F"
|
||||
|
||||
|
||||
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's sync API keeps an asyncio loop running on the test thread,
|
||||
so ``asyncio.run`` cannot be called directly from a test body.
|
||||
"""
|
||||
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, seed: bool) -> ImportSummary | None:
|
||||
"""Truncate the KB (and query log), then optionally re-import fixtures."""
|
||||
with SessionLocal() as db:
|
||||
db.execute(text("TRUNCATE chunks, documents, query_log"))
|
||||
db.commit()
|
||||
if not seed:
|
||||
return None
|
||||
return _run_in_thread(_import_fixtures(mock_port))
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def seeded_kb(mock_llm: int, db_ready: None) -> Iterator[None]:
|
||||
"""A fresh KB seeded from ``tests/fixtures/docs`` (8 docs, A9 formats),
|
||||
truncated again on teardown. ``db_ready`` (conftest) skips with clear
|
||||
instructions when Postgres is down."""
|
||||
summary = _reset_db(mock_llm, seed=True)
|
||||
assert summary is not None and summary.added == 8
|
||||
yield
|
||||
_reset_db(mock_llm, seed=False)
|
||||
|
||||
|
||||
def send_and_wait(page: Page, question: str) -> None:
|
||||
"""Type into #message-input, submit via #composer, then wait until the
|
||||
last brain message settles (send button re-enabled, label "Send")."""
|
||||
page.fill("#message-input", question)
|
||||
page.evaluate("() => document.querySelector('#composer').requestSubmit()")
|
||||
expect(page.locator(".msg.user .bubble").last).to_contain_text(question)
|
||||
# The mock streams at 0.02s/chunk, so thinking + answer land in a few
|
||||
# seconds — 30s is generous on headless Chromium.
|
||||
expect(page.locator(".msg.brain .bubble").last).not_to_have_text("", timeout=30_000)
|
||||
expect(page.locator("#send-btn")).to_be_enabled(timeout=30_000)
|
||||
expect(page.locator("#send-label")).to_have_text("Send")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. Streaming: the block attaches OPEN at the first thinking event, then
|
||||
# auto-collapses when the first answer token lands
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_thinking_block_streams_open_then_collapses(
|
||||
page: Page, app_url: str, seeded_kb: None
|
||||
) -> None:
|
||||
page.set_default_timeout(30_000)
|
||||
page.goto(app_url)
|
||||
page.fill("#message-input", THINK_QUESTION)
|
||||
page.evaluate("() => document.querySelector('#composer').requestSubmit()")
|
||||
expect(page.locator(".msg.user .bubble").last).to_contain_text(THINK_QUESTION)
|
||||
|
||||
# The block attaches at the FIRST thinking event — before any answer
|
||||
# token — and is created OPEN.
|
||||
details = page.locator(".msg.brain").last.locator("details.thinking")
|
||||
details.wait_for(state="attached", timeout=10_000)
|
||||
# The ~800-char thinking stream (≈1.3s) keeps the block open right
|
||||
# after attach — assert while it is still streaming.
|
||||
expect(details).to_have_attribute("open", "")
|
||||
expect(details.locator(".thinking-text")).not_to_have_text("")
|
||||
|
||||
# First answer token: the block auto-collapses and stays closed.
|
||||
bubble = page.locator(".msg.brain .bubble").last
|
||||
expect(bubble).not_to_have_text("", timeout=30_000)
|
||||
expect(details).not_to_have_attribute("open")
|
||||
|
||||
# Settled: full scratchpad, grounded mock answer, source chip(s),
|
||||
# and the re-enabled send button.
|
||||
expect(details.locator(".thinking-text")).to_contain_text(THINKING_FRAGMENT)
|
||||
expect(bubble).to_contain_text(MOCK_ANSWER_MARKER)
|
||||
chip = page.locator(".msg.brain .source-chip", has_text="kubernetes.md")
|
||||
expect(chip.first).to_be_visible()
|
||||
expect(chip.first).to_have_attribute("href", CHIP_HREF)
|
||||
expect(page.locator("#send-btn")).to_be_enabled()
|
||||
expect(page.locator("#send-label")).to_have_text("Send")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. Toggle: after settle the block is closed; the summary re-opens it
|
||||
# (a real keyboard-focusable control) and closes it again
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_thinking_toggle_after_done(page: Page, app_url: str, seeded_kb: None) -> None:
|
||||
page.set_default_timeout(30_000)
|
||||
page.goto(app_url)
|
||||
send_and_wait(page, THINK_QUESTION)
|
||||
|
||||
last = page.locator(".msg.brain").last
|
||||
details = last.locator("details.thinking")
|
||||
expect(details).to_have_count(1)
|
||||
expect(details).not_to_have_attribute("open") # auto-collapsed at first token
|
||||
|
||||
# The summary is a real, keyboard-focusable control.
|
||||
details.locator("summary").focus()
|
||||
assert page.evaluate("() => document.activeElement.tagName") == "SUMMARY"
|
||||
|
||||
# Open: the full scratchpad is visible.
|
||||
details.locator("summary").click()
|
||||
expect(details).to_have_attribute("open", "")
|
||||
text_el = details.locator(".thinking-text")
|
||||
expect(text_el).to_be_visible()
|
||||
expect(text_el).to_contain_text(THINKING_FRAGMENT)
|
||||
expect(text_el).to_contain_text("nothing is invented")
|
||||
|
||||
# Closed again — user control in both directions.
|
||||
details.locator("summary").click()
|
||||
expect(details).not_to_have_attribute("open")
|
||||
expect(text_el).not_to_be_visible()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Persistence: the thinking block (and its text) survives a reload,
|
||||
# restored COLLAPSED — phase-14 restore path + phase-17 field
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_thinking_restored_after_reload(page: Page, app_url: str, seeded_kb: None) -> None:
|
||||
page.set_default_timeout(30_000)
|
||||
page.goto(app_url)
|
||||
send_and_wait(page, THINK_QUESTION)
|
||||
|
||||
# The live block is collapsed; capture what it shows and what the
|
||||
# turn persisted (raw text, same as what was rendered).
|
||||
details = page.locator(".msg.brain").last.locator("details.thinking")
|
||||
expect(details).not_to_have_attribute("open")
|
||||
captured = details.locator(".thinking-text").text_content()
|
||||
assert captured
|
||||
# The persisted raw text is the same scratchpad (renderMarkdown turns
|
||||
# the line breaks into <br>, which textContent drops — compare without
|
||||
# whitespace).
|
||||
raw = json.loads(
|
||||
page.evaluate(f"() => localStorage.getItem('{STORAGE_KEY}')")
|
||||
)["messages"][1]["thinking"]
|
||||
assert re.sub(r"\s+", "", raw) == re.sub(r"\s+", "", captured)
|
||||
|
||||
page.reload()
|
||||
expect(page.locator("#empty-state")).to_be_hidden()
|
||||
|
||||
restored = page.locator(".msg.brain").last.locator("details.thinking")
|
||||
expect(restored).to_have_count(1)
|
||||
expect(restored).not_to_have_attribute("open") # restored COLLAPSED
|
||||
expect(restored.locator(".thinking-text")).to_have_text(captured)
|
||||
|
||||
# Answer bubble + source chip are intact (phase-14 restore path).
|
||||
expect(page.locator(".msg.brain .bubble").last).to_contain_text(MOCK_ANSWER_MARKER)
|
||||
chip = page.locator(".msg.brain .source-chip", has_text="kubernetes.md")
|
||||
expect(chip.first).to_have_attribute("href", CHIP_HREF)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. No thinking, no block: a model/turn that emits no reasoning renders
|
||||
# exactly as before (no layout regression)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_no_thinking_block_without_trigger(page: Page, app_url: str, seeded_kb: None) -> None:
|
||||
page.set_default_timeout(30_000)
|
||||
page.goto(app_url)
|
||||
send_and_wait(page, PLAIN_QUESTION)
|
||||
|
||||
# No trigger → no thinking events → no block anywhere on the page.
|
||||
expect(page.locator("details.thinking")).to_have_count(0)
|
||||
|
||||
# The turn itself is complete and grounded, exactly as before phase 17.
|
||||
chip = page.locator(".msg.brain .source-chip", has_text="kubernetes.md")
|
||||
expect(chip.first).to_be_visible()
|
||||
expect(chip.first).to_have_attribute("href", CHIP_HREF)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. Coexistence: the honesty gate (deflection) and the thinking block
|
||||
# on the same turn
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_thinking_with_deflection(page: Page, app_url: str, seeded_kb: None) -> None:
|
||||
page.set_default_timeout(30_000)
|
||||
page.goto(app_url)
|
||||
send_and_wait(page, THINK_DEFLECT_QUESTION)
|
||||
|
||||
last = page.locator(".msg.brain").last
|
||||
# The honesty gate fired: amber deflected bubble + "Maybe try" chips.
|
||||
expect(last).to_have_class(re.compile(r"is-deflected"))
|
||||
expect(last.locator(".bubble")).to_contain_text(
|
||||
re.compile(DEFLECT_PHRASE, re.IGNORECASE)
|
||||
)
|
||||
chips = last.locator(".maybe-try .suggestion-chip")
|
||||
expect(chips.first).to_be_visible()
|
||||
assert chips.count() >= 2
|
||||
|
||||
# And the thinking block came along, closed, with its scratchpad.
|
||||
details = last.locator("details.thinking")
|
||||
expect(details).to_have_count(1)
|
||||
expect(details).not_to_have_attribute("open")
|
||||
expect(details.locator(".thinking-text")).to_contain_text(THINKING_FRAGMENT)
|
||||
Reference in New Issue
Block a user