feat(rag): steering notes — tune how Brain answers, stored in Postgres and injected into every system prompt

2026-08-22 16:44:42 -04:00
parent 19df7df99d
commit fc0d9a2d5c
19 changed files with 1589 additions and 34 deletions
@@ -0,0 +1,253 @@
+"""Integration: steering notes (phase 15) — CRUD + system-prompt injection.
+
+Real Postgres (``podman compose up -d db``); the chat path reuses the
+deterministic fake LLM from ``test_chat_api`` (token-overlap embeddings),
+so the stored note's journey — API → Postgres → ``<tuning>`` section of
+the captured system prompt — is verified end-to-end without a network.
+
+Requires:  podman compose up -d db
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+import uuid
+from collections.abc import Iterator
+from datetime import UTC, datetime, timedelta
+from pathlib import Path
+
+import pytest
+from fastapi.testclient import TestClient
+from sqlalchemy import select, text
+from test_chat_api import FakeRagLLM, _stream_chat
+
+from app.api import chat as chat_api
+from app.main import app as fastapi_app
+from app.models import SteeringNote
+from app.rag.importer import import_sources
+
+FIXTURES = Path(__file__).resolve().parents[1] / "fixtures" / "docs"
+QUESTION = "How is my Kubernetes cluster set up?"
+OFF_TOPIC = "How do I bake sourdough bread?"
+NOTE = "STEEER-MARKER be concise"
+
+
+@pytest.fixture(autouse=True)
+def clean_steering(db) -> Iterator[None]:
+    """Steering notes + query log are global state: reset around every test."""
+    db.execute(text("TRUNCATE steering_notes, query_log"))
+    db.commit()
+    yield
+    db.execute(text("TRUNCATE steering_notes, query_log"))
+    db.commit()
+
+
+@pytest.fixture()
+def seeded_kb(db) -> Iterator[FakeRagLLM]:
+    """Fresh Postgres with the fixture docs imported (real pipeline)."""
+    db.execute(text("TRUNCATE chunks, documents, query_log, steering_notes"))
+    db.commit()
+    llm = FakeRagLLM()
+    summary = asyncio.run(import_sources([FIXTURES], llm, session=db))
+    assert summary.added == 8  # A9 formats; .hidden/ skipped
+    yield llm
+    db.execute(text("TRUNCATE chunks, documents, query_log, steering_notes"))
+    db.commit()
+
+
+def _turn_log_lines(caplog: pytest.LogCaptureFixture) -> list[str]:
+    """The per-turn ``question=…`` log lines (PLAN §9) from this test."""
+    return [r.getMessage() for r in caplog.records if "question=" in r.getMessage()]
+
+
+# ---------- CRUD ----------
+
+
+def test_create_note_returns_201_and_stores_trimmed(client: TestClient, db) -> None:
+    r = client.post("/api/steering", json={"note": f"  {NOTE}  "})
+    assert r.status_code == 201
+    body = r.json()
+    assert body["note"] == NOTE  # trimmed before storage
+    uuid.UUID(body["id"])  # valid UUID
+    assert body["created_at"]
+    rows = db.scalars(select(SteeringNote)).all()
+    assert [row.note for row in rows] == [NOTE]
+
+
+def test_list_notes_empty(client: TestClient) -> None:
+    r = client.get("/api/steering")
+    assert r.status_code == 200
+    assert r.json() == {"notes": []}
+
+
+def test_list_notes_newest_first(client: TestClient, db) -> None:
+    base = datetime.now(UTC)
+    db.add_all(
+        [
+            SteeringNote(note="oldest", created_at=base),
+            SteeringNote(note="newest", created_at=base + timedelta(hours=2)),
+            SteeringNote(note="middle", created_at=base + timedelta(hours=1)),
+        ]
+    )
+    db.commit()
+
+    r = client.get("/api/steering")
+    assert r.status_code == 200
+    body = r.json()
+    assert [n["note"] for n in body["notes"]] == ["newest", "middle", "oldest"]
+    for n in body["notes"]:
+        assert set(n) == {"id", "note", "created_at"}
+        uuid.UUID(n["id"])
+
+
+def test_delete_note_returns_204_and_removes(client: TestClient, db) -> None:
+    created = client.post("/api/steering", json={"note": NOTE}).json()
+
+    assert client.delete(f"/api/steering/{created['id']}").status_code == 204
+    assert client.get("/api/steering").json() == {"notes": []}
+    assert db.scalars(select(SteeringNote)).all() == []
+
+
+def test_delete_unknown_note_returns_404(client: TestClient) -> None:
+    r = client.delete(f"/api/steering/{uuid.uuid4()}")
+    assert r.status_code == 404
+    assert "not found" in r.json()["detail"]
+
+
+def test_delete_invalid_id_returns_422(client: TestClient) -> None:
+    assert client.delete("/api/steering/not-a-uuid").status_code == 422
+
+
+def test_create_rejects_empty_and_blank_notes(client: TestClient) -> None:
+    assert client.post("/api/steering", json={"note": ""}).status_code == 422
+    assert client.post("/api/steering", json={"note": " \t\n "}).status_code == 422
+    assert client.get("/api/steering").json() == {"notes": []}
+
+
+def test_create_enforces_2000_char_limit(client: TestClient) -> None:
+    assert client.post("/api/steering", json={"note": "x" * 2001}).status_code == 422
+    r = client.post("/api/steering", json={"note": "x" * 2000})
+    assert r.status_code == 201
+    assert len(r.json()["note"]) == 2000
+
+
+# ---------- chat turn: note reaches the system prompt ----------
+
+
+def test_chat_turn_high_mode_receives_note_in_system_prompt(
+    client: TestClient, seeded_kb: FakeRagLLM, caplog: pytest.LogCaptureFixture
+) -> None:
+    client.post("/api/steering", json={"note": NOTE})
+    caplog.set_level(logging.INFO, logger="app.chat")
+
+    fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: seeded_kb
+    try:
+        _, _, frames = _stream_chat(client, QUESTION)
+    finally:
+        fastapi_app.dependency_overrides.clear()
+
+    assert frames[-1]["type"] == "done"
+    assert frames[-1]["deflected"] is False
+
+    # The LLM received the HIGH prompt with the <tuning> section.
+    (system, user) = seeded_kb.seen_messages[0][0], seeded_kb.seen_messages[0][1]
+    assert user["content"] == QUESTION
+    assert "<relevance>HIGH</relevance>" in system["content"]
+    assert "<tuning>" in system["content"]
+    assert f"1. {NOTE}" in system["content"]
+    assert "<documents>" in system["content"]
+    # The section sits between the relevance marker and the documents.
+    assert (
+        system["content"].index("<relevance>HIGH</relevance>")
+        < system["content"].index("<tuning>")
+        < system["content"].index("</tuning>")
+        < system["content"].index("<documents>")
+    )
+
+    # The per-turn log line records tuning=N (PLAN §9).
+    lines = _turn_log_lines(caplog)
+    assert lines and "tuning=1" in lines[-1]
+
+
+def test_chat_turn_low_mode_receives_note_in_system_prompt(
+    client: TestClient, seeded_kb: FakeRagLLM, caplog: pytest.LogCaptureFixture
+) -> None:
+    client.post("/api/steering", json={"note": NOTE})
+    caplog.set_level(logging.INFO, logger="app.chat")
+
+    fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: seeded_kb
+    try:
+        _, _, frames = _stream_chat(client, OFF_TOPIC)
+    finally:
+        fastapi_app.dependency_overrides.clear()
+
+    assert frames[-1]["deflected"] is True
+    (system, _user) = seeded_kb.seen_messages[0][0], seeded_kb.seen_messages[0][1]
+    assert "<relevance>LOW</relevance>" in system["content"]
+    assert "DEFLECT_MODE" in system["content"]
+    assert "<tuning>" in system["content"]
+    assert f"1. {NOTE}" in system["content"]
+    lines = _turn_log_lines(caplog)
+    assert lines and "tuning=1" in lines[-1]
+
+
+def test_chat_turn_without_notes_has_no_tuning_section(
+    client: TestClient, seeded_kb: FakeRagLLM, caplog: pytest.LogCaptureFixture
+) -> None:
+    caplog.set_level(logging.INFO, logger="app.chat")
+
+    fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: seeded_kb
+    try:
+        _stream_chat(client, QUESTION)
+    finally:
+        fastapi_app.dependency_overrides.clear()
+
+    (system, _user) = seeded_kb.seen_messages[0][0], seeded_kb.seen_messages[0][1]
+    assert "<tuning>" not in system["content"]
+    lines = _turn_log_lines(caplog)
+    assert lines and "tuning=0" in lines[-1]
+
+
+def test_chat_turn_numbers_notes_oldest_first(client: TestClient, db, seeded_kb) -> None:
+    base = datetime.now(UTC)
+    db.add_all(
+        [
+            SteeringNote(note="older note", created_at=base),
+            SteeringNote(note="newer note", created_at=base + timedelta(hours=1)),
+        ]
+    )
+    db.commit()
+
+    fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: seeded_kb
+    try:
+        _stream_chat(client, QUESTION)
+    finally:
+        fastapi_app.dependency_overrides.clear()
+
+    (system, _user) = seeded_kb.seen_messages[0][0], seeded_kb.seen_messages[0][1]
+    assert "<tuning>" in system["content"]
+    assert "1. older note" in system["content"]
+    assert "2. newer note" in system["content"]
+    assert system["content"].index("1. older note") < system["content"].index("2. newer note")
+
+
+def test_multiple_turns_keep_reading_notes(client: TestClient, seeded_kb) -> None:
+    """The note steers EVERY subsequent turn, not just the next one."""
+    client.post("/api/steering", json={"note": NOTE})
+
+    fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: seeded_kb
+    try:
+        _stream_chat(client, QUESTION)
+        _stream_chat(client, QUESTION)
+        assert len(seeded_kb.seen_messages) == 2
+        for messages in seeded_kb.seen_messages:
+            assert f"1. {NOTE}" in messages[0]["content"]
+
+        # Delete → the following turn is clean again.
+        note_id = client.get("/api/steering").json()["notes"][0]["id"]
+        assert client.delete(f"/api/steering/{note_id}").status_code == 204
+        _stream_chat(client, QUESTION)
+        assert len(seeded_kb.seen_messages) == 3
+        assert "<tuning>" not in seeded_kb.seen_messages[-1][0]["content"]
+    finally:
+        fastapi_app.dependency_overrides.clear()