feat(rag): steering notes — tune how Brain answers, stored in Postgres and injected into every system prompt
This commit is contained in:
+42
-10
@@ -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).
|
||||
- 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.
|
||||
|
||||
``max_tokens`` is honored deterministically (token ≈ whitespace word),
|
||||
like a real endpoint: an answer longer than the cap is truncated. This
|
||||
@@ -89,25 +92,54 @@ def long_answer() -> str:
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
#: First numbered note line of a ``<tuning>`` section (phase 15).
|
||||
_TUNING_BLOCK_RE = re.compile(r"<tuning>\n(.*?)\n</tuning>", re.S)
|
||||
_NOTE_LINE_RE = re.compile(r"^\d+\.\s*(.+)$")
|
||||
|
||||
|
||||
def first_tuning_note(system: str) -> str | None:
|
||||
"""The first steering note in the system prompt, or ``None``.
|
||||
|
||||
The prompt numbers notes 1..N oldest-first (see
|
||||
``app.rag.prompts.build_steering_section``); the mock echoes the first
|
||||
one into its answer so prompt injection is observable in the UI.
|
||||
"""
|
||||
block = _TUNING_BLOCK_RE.search(system)
|
||||
if not block:
|
||||
return None
|
||||
for line in block.group(1).splitlines():
|
||||
m = _NOTE_LINE_RE.match(line.strip())
|
||||
if m:
|
||||
return m.group(1).strip()
|
||||
return None
|
||||
|
||||
|
||||
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 (
|
||||
answer = long_answer()
|
||||
elif "DEFLECT_MODE" in system:
|
||||
answer = (
|
||||
"Ah — I haven't done anything like that, so I don't want to make stuff up! "
|
||||
"You're thinking bigger than my notes for a second. Try asking about "
|
||||
"kubernetes, backups, or deploying a new service — I know those inside out. "
|
||||
"You've got this!"
|
||||
)
|
||||
ctx = _context(body)
|
||||
snippet = ctx[:220].replace("\n", " ").strip()
|
||||
return (
|
||||
f"Great question — you've absolutely got this! Here's what my notes say about "
|
||||
f"“{user.strip()[:80]}”: {snippet}… That's the gist from the docs; happy to "
|
||||
"dig into any of it. (Deterministic mock answer for E2E.)"
|
||||
)
|
||||
else:
|
||||
ctx = _context(body)
|
||||
snippet = ctx[:220].replace("\n", " ").strip()
|
||||
answer = (
|
||||
f"Great question — you've absolutely got this! Here's what my notes say about "
|
||||
f"“{user.strip()[:80]}”: {snippet}… That's the gist from the docs; happy to "
|
||||
"dig into any of it. (Deterministic mock answer for E2E.)"
|
||||
)
|
||||
# Steering (phase 15): when the system prompt carries <tuning>, the
|
||||
# answer ends with the first note — deterministically observable.
|
||||
note = first_tuning_note(system)
|
||||
if note:
|
||||
answer = f"{answer} (tuning: {note})"
|
||||
return answer
|
||||
|
||||
|
||||
@app.post("/__shutdown__")
|
||||
|
||||
@@ -0,0 +1,293 @@
|
||||
"""Phase 15 E2E (Playwright): tune how Brain answers (steering notes).
|
||||
|
||||
Story: ``.agent/user_stories/steering-notes.md``
|
||||
Run in isolation (DB must be up: ``podman compose up -d db``):
|
||||
|
||||
uv run pytest tests/e2e/test_steering.py -v --no-cov
|
||||
|
||||
The steering loop: "Tune" under a completed answer → short instruction →
|
||||
stored in Postgres (``steering_notes``) → injected into the system prompt
|
||||
of every subsequent turn as the ``<tuning>`` section. The mock LLM
|
||||
echoes the first tuning note into its answer
|
||||
(`` (tuning: <first note line>)``), so prompt injection is observable in
|
||||
the UI deterministically. Notes are listed newest-first in the header
|
||||
"Tuning" panel, where each can be deleted.
|
||||
|
||||
Test → story mapping (Playwright Mapping Rule):
|
||||
1. ``test_tune_under_answer_persists_and_steers``
|
||||
2. ``test_delete_note_stops_steering``
|
||||
3. ``test_note_rendered_as_text_xss_safe``
|
||||
4. ``test_tuning_panel_a11y``
|
||||
"""
|
||||
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 select, text
|
||||
|
||||
from app.config import Settings
|
||||
from app.db import SessionLocal
|
||||
from app.models import SteeringNote
|
||||
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"
|
||||
QUESTION = "How is my Kubernetes cluster set up?"
|
||||
MOCK_ANSWER_MARKER = "Deterministic mock answer for E2E"
|
||||
NOTE = "STEEER-MARKER be concise"
|
||||
XSS_NOTE = "<script>window.__xss = true; alert('xss')</script>"
|
||||
#: index.html ships exactly two classic/module script tags.
|
||||
BASE_SCRIPT_COUNT = 2
|
||||
|
||||
|
||||
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 + steering notes), re-import fixtures."""
|
||||
with SessionLocal() as db:
|
||||
db.execute(text("TRUNCATE chunks, documents, query_log, steering_notes"))
|
||||
db.commit()
|
||||
if not seed:
|
||||
return None
|
||||
return _run_in_thread(_import_fixtures(mock_port))
|
||||
|
||||
|
||||
def _ask(page: Page, question: str) -> None:
|
||||
"""Send one turn and wait until the grounded answer has fully landed."""
|
||||
page.fill("#message-input", question)
|
||||
page.click("#send-btn")
|
||||
expect(page.locator(".msg.user .bubble").last).to_contain_text(question)
|
||||
expect(page.locator(".msg.brain .bubble").last).to_contain_text(
|
||||
MOCK_ANSWER_MARKER, timeout=30_000
|
||||
)
|
||||
expect(page.locator("#send-btn")).to_be_enabled()
|
||||
expect(page.locator("#send-label")).to_have_text("Send")
|
||||
|
||||
|
||||
def _tune_and_save(page: Page, note: str) -> None:
|
||||
"""Tune the last completed brain bubble and save *note*."""
|
||||
tune = page.locator(".msg.brain .tune-btn").last
|
||||
expect(tune).to_be_visible()
|
||||
tune.click()
|
||||
form = page.locator(".msg.brain .tune-form").last
|
||||
expect(form).to_be_visible()
|
||||
form.locator("textarea").fill(note)
|
||||
form.locator(".tune-save").click()
|
||||
saved = page.locator(".msg.brain .tune-saved").last
|
||||
expect(saved).to_contain_text("Saved — future answers will follow this.", timeout=15_000)
|
||||
|
||||
|
||||
def _open_panel(page: Page) -> None:
|
||||
page.click("#steering-toggle")
|
||||
expect(page.locator("#steering-panel")).to_be_visible()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. Tune under an answer → persisted → next answer carries the note
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_tune_under_answer_persists_and_steers(
|
||||
page: Page, app_url: str, mock_llm: int, db_ready: None
|
||||
) -> None:
|
||||
summary = _reset_db(mock_llm, seed=True)
|
||||
assert summary is not None and summary.added == 8 # A9 formats
|
||||
page.set_default_timeout(30_000)
|
||||
page.goto(app_url)
|
||||
_ask(page, QUESTION)
|
||||
|
||||
# The Tune control: ghost button in the answer's meta row, ≥44px.
|
||||
tune = page.locator(".msg.brain .tune-btn").last
|
||||
expect(tune).to_have_count(1)
|
||||
expect(tune).to_have_attribute("type", "button")
|
||||
box = tune.bounding_box()
|
||||
assert box is not None and box["height"] >= 44
|
||||
|
||||
_tune_and_save(page, NOTE)
|
||||
|
||||
# Persisted in Postgres.
|
||||
with SessionLocal() as db:
|
||||
rows = db.scalars(select(SteeringNote)).all()
|
||||
assert [r.note for r in rows] == [NOTE]
|
||||
|
||||
# The header panel shows the note with an updated count badge.
|
||||
_open_panel(page)
|
||||
expect(page.locator("#steering-count")).to_have_text("1")
|
||||
expect(page.locator("#steering-list .steering-note")).to_have_count(1)
|
||||
expect(page.locator("#steering-list .steering-note-text")).to_have_text(NOTE)
|
||||
page.click("#steering-toggle") # close again
|
||||
|
||||
# The NEXT answer carries the note — it reached the system prompt.
|
||||
_ask(page, QUESTION)
|
||||
bubble = page.locator(".msg.brain .bubble").last
|
||||
expect(bubble).to_contain_text(f"(tuning: {NOTE})")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. Delete from the panel → count 0 → steering stops
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delete_note_stops_steering(
|
||||
page: Page, app_url: str, mock_llm: int, db_ready: None
|
||||
) -> None:
|
||||
_reset_db(mock_llm, seed=True)
|
||||
page.set_default_timeout(30_000)
|
||||
page.goto(app_url)
|
||||
_ask(page, QUESTION)
|
||||
_tune_and_save(page, NOTE)
|
||||
|
||||
# Steering is live: one more answer carries the marker.
|
||||
_ask(page, QUESTION)
|
||||
expect(page.locator(".msg.brain .bubble").last).to_contain_text(f"(tuning: {NOTE})")
|
||||
|
||||
# Delete the note from the panel.
|
||||
_open_panel(page)
|
||||
expect(page.locator("#steering-count")).to_have_text("1")
|
||||
page.locator("#steering-list .steering-delete").click()
|
||||
expect(page.locator("#steering-list .steering-note")).to_have_count(0)
|
||||
expect(page.locator("#steering-count")).to_have_text("0")
|
||||
expect(page.locator("#steering-empty")).to_be_visible()
|
||||
expect(page.locator("#steering-announcer")).to_contain_text("deleted")
|
||||
|
||||
with SessionLocal() as db:
|
||||
assert db.scalars(select(SteeringNote)).all() == []
|
||||
|
||||
# The next answer no longer carries the marker.
|
||||
_ask(page, QUESTION)
|
||||
bubble = page.locator(".msg.brain .bubble").last
|
||||
expect(bubble).to_contain_text(MOCK_ANSWER_MARKER)
|
||||
expect(bubble).not_to_contain_text("STEEER-MARKER")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Notes render as text (XSS-safe)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_note_rendered_as_text_xss_safe(
|
||||
page: Page, app_url: str, mock_llm: int, db_ready: None
|
||||
) -> None:
|
||||
_reset_db(mock_llm, seed=True)
|
||||
page.set_default_timeout(30_000)
|
||||
page.goto(app_url)
|
||||
|
||||
dialogs: list[str] = []
|
||||
|
||||
def _handle_dialog(d) -> None:
|
||||
dialogs.append(d.message)
|
||||
d.dismiss()
|
||||
|
||||
page.on("dialog", _handle_dialog)
|
||||
|
||||
_ask(page, QUESTION)
|
||||
_tune_and_save(page, XSS_NOTE)
|
||||
|
||||
# Panel: the payload is visible as LITERAL text…
|
||||
_open_panel(page)
|
||||
expect(page.locator("#steering-list .steering-note-text")).to_have_text(XSS_NOTE)
|
||||
|
||||
# …never as an executed element: no script tag anywhere, no dialog.
|
||||
assert page.locator("#steering-panel script").count() == 0
|
||||
expect(page.locator("script")).to_have_count(BASE_SCRIPT_COUNT)
|
||||
assert dialogs == [], f"the note must never execute as script: {dialogs}"
|
||||
assert page.evaluate("() => window.__xss === undefined") is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Tuning panel accessibility
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_tuning_panel_a11y(page: Page, app_url: str, db_ready: None) -> None:
|
||||
_reset_db(mock_port=0, seed=False) # no KB seeding needed for the panel a11y
|
||||
page.set_default_timeout(30_000)
|
||||
page.goto(app_url)
|
||||
|
||||
toggle = page.locator("#steering-toggle")
|
||||
panel = page.locator("#steering-panel")
|
||||
announcer = page.locator("#steering-announcer")
|
||||
|
||||
# Initial: closed, correctly wired, polite live region present.
|
||||
expect(toggle).to_have_attribute("aria-expanded", "false")
|
||||
expect(toggle).to_have_attribute("aria-controls", "steering-panel")
|
||||
expect(panel).to_have_attribute("role", "region")
|
||||
assert "Tuning notes" in (panel.get_attribute("aria-label") or "")
|
||||
expect(panel).to_be_hidden()
|
||||
assert announcer.get_attribute("role") == "status"
|
||||
assert announcer.get_attribute("aria-live") == "polite"
|
||||
# Accessible name comes from its visible text (icon is aria-hidden).
|
||||
assert "Tuning" in toggle.inner_text()
|
||||
|
||||
# Open: expanded + the designed empty state.
|
||||
toggle.click()
|
||||
expect(toggle).to_have_attribute("aria-expanded", "true")
|
||||
expect(panel).to_be_visible()
|
||||
expect(page.locator("#steering-empty")).to_be_visible()
|
||||
expect(page.locator("#steering-count")).to_have_text("0")
|
||||
|
||||
# Add a note (API), then re-open the panel to refresh it.
|
||||
page.evaluate(
|
||||
"""async () => {
|
||||
const r = await fetch('/api/steering', {
|
||||
method: 'POST',
|
||||
headers: {'Content-Type': 'application/json'},
|
||||
body: JSON.stringify({note: 'a11y note one'}),
|
||||
});
|
||||
if (!r.ok) throw new Error('steering POST failed: ' + r.status);
|
||||
}"""
|
||||
)
|
||||
toggle.click() # close
|
||||
toggle.click() # re-open (refreshes the list)
|
||||
note_item = page.locator("#steering-list .steering-note")
|
||||
expect(note_item).to_have_count(1)
|
||||
expect(note_item.locator(".steering-note-text")).to_have_text("a11y note one")
|
||||
|
||||
# The per-note delete is a real, labeled button (≥44px target).
|
||||
delete = page.locator("#steering-list .steering-delete")
|
||||
expect(delete).to_have_attribute("type", "button")
|
||||
assert (delete.get_attribute("aria-label") or "").startswith("Delete tuning note:")
|
||||
box = delete.bounding_box()
|
||||
assert box is not None and box["height"] >= 44
|
||||
|
||||
# Delete: list empties, count updates, the live region announces it.
|
||||
delete.click()
|
||||
expect(page.locator("#steering-list .steering-note")).to_have_count(0)
|
||||
expect(page.locator("#steering-count")).to_have_text("0")
|
||||
expect(announcer).to_contain_text("deleted")
|
||||
|
||||
# And the toggle closes cleanly again.
|
||||
toggle.click()
|
||||
expect(toggle).to_have_attribute("aria-expanded", "false")
|
||||
expect(panel).to_be_hidden()
|
||||
Reference in New Issue
Block a user