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,193 @@
+"""Unit: steering notes (phase 15) — the <tuning> prompt section.
+
+Pure logic, no Postgres and no network: :func:`build_steering_section`
+(empty/one/many/budget-truncation), its placement in the HIGH and LOW
+prompts, and the ``plan_turn`` wiring (notes → prompt + ``tuning_count``).
+"""
+from __future__ import annotations
+
+import uuid
+
+import pytest
+
+from app.api import chat as chat_api
+from app.config import Settings
+from app.models import Document
+from app.rag.prompts import build_deflect_prompt, build_high_prompt, build_steering_section
+from app.rag.retriever import TRUNCATION_MARKER, RetrievedChunk
+
+
+def _settings() -> Settings:
+    return Settings(_env_file=None, relevance_threshold=0.30)  # pyright: ignore[reportCallIssue]
+
+
+def _doc(title: str, content: str) -> Document:
+    return Document(
+        id=uuid.uuid4(),
+        source="Homelab",
+        path=f"{title.lower().replace(' ', '-')}.md",
+        full_path="/tmp/doc.md",
+        title=title,
+        content=content,
+        content_hash="0" * 64,
+    )
+
+
+def _chunk(doc: Document, score: float) -> RetrievedChunk:
+    return RetrievedChunk(
+        chunk_id=uuid.uuid4(),
+        position=0,
+        content=doc.content[:32],
+        score=score,
+        document=doc,
+        cosine=score,
+    )
+
+
+# ---------- build_steering_section ----------
+
+
+def test_steering_section_empty_when_no_notes() -> None:
+    assert build_steering_section([]) == ""
+
+
+def test_steering_section_empty_when_notes_are_blank() -> None:
+    assert build_steering_section(["", "   ", "\n\t"]) == ""
+
+
+def test_steering_section_single_note_numbered() -> None:
+    section = build_steering_section(["be more concise"])
+    assert section.startswith("<tuning>\n")
+    assert section.endswith("\n</tuning>")
+    assert "1. be more concise" in section
+    assert TRUNCATION_MARKER not in section
+
+
+def test_steering_section_trims_note_edges() -> None:
+    section = build_steering_section(["  be more concise  "])
+    assert "1. be more concise" in section
+    assert "1.   be more concise" not in section
+
+
+def test_steering_section_many_notes_numbered_in_order() -> None:
+    section = build_steering_section(["alpha", "beta", "gamma"])
+    assert "1. alpha" in section
+    assert "2. beta" in section
+    assert "3. gamma" in section
+    assert section.index("1. alpha") < section.index("2. beta") < section.index("3. gamma")
+    assert TRUNCATION_MARKER not in section
+
+
+def test_steering_section_budget_truncation_keeps_oldest_prefix_and_marker() -> None:
+    # Each note is 300 chars; with a 600-char budget only note 1 fits, so
+    # the oldest-fitting prefix is kept and the overflow is marked.
+    notes = [f"note-{i} " + "x" * (300 - len(f"note-{i} ")) for i in range(3)]
+    section = build_steering_section(notes, max_chars=600)
+    assert TRUNCATION_MARKER in section
+    assert len(section) <= 600
+    assert "1. note-0" in section
+    assert "note-1" not in section
+    assert "note-2" not in section
+    # The marker comes last, after the kept notes.
+    assert section.index("1. note-0") < section.index(TRUNCATION_MARKER)
+
+
+def test_steering_section_fits_budget_exactly_when_all_notes_fit() -> None:
+    section = build_steering_section(["a", "b", "c"], max_chars=10_000)
+    assert TRUNCATION_MARKER not in section
+    assert len(section) <= 10_000
+
+
+def test_steering_section_default_budget_from_settings(monkeypatch: pytest.MonkeyPatch) -> None:
+    from app.rag import prompts as prompts_mod
+
+    monkeypatch.setattr(
+        prompts_mod, "get_settings", lambda: Settings(_env_file=None)  # pyright: ignore[reportCallIssue]
+    )
+    # 5 notes of 2000 chars (the API max) = 10k+ chars > the 8000 default.
+    notes = [f"note-{i} " + "y" * (2000 - len(f"note-{i} ")) for i in range(5)]
+    section = build_steering_section(notes)
+    assert TRUNCATION_MARKER in section
+    assert len(section) <= 8_000
+
+
+def test_steering_section_nonpositive_budget_is_empty() -> None:
+    assert build_steering_section(["be concise"], max_chars=0) == ""
+    assert build_steering_section(["be concise"], max_chars=-10) == ""
+
+
+def test_steering_section_tiny_budget_never_exceeds_cap() -> None:
+    # Pathological budget: the section must never exceed the cap — bare
+    # marker when it fits, no section at all when even that doesn't.
+    assert len(build_steering_section(["a" * 500], max_chars=10)) <= 10
+    fits_marker = build_steering_section(["a" * 500], max_chars=len(TRUNCATION_MARKER))
+    assert fits_marker == TRUNCATION_MARKER
+
+
+# ---------- prompt placement (both modes) ----------
+
+
+def test_high_prompt_steering_sits_between_relevance_and_documents() -> None:
+    doc = _doc("Kubernetes Homelab Cluster", "TALOS_DOC_CONTENT")
+    prompt = build_high_prompt([doc], notes=["be concise"])
+    i_rel = prompt.index("<relevance>HIGH</relevance>")
+    i_open = prompt.index("<tuning>")
+    i_close = prompt.index("</tuning>")
+    i_docs = prompt.index("<documents>")
+    assert i_rel < i_open < i_close < i_docs
+    assert "1. be concise" in prompt
+    assert "TALOS_DOC_CONTENT" in prompt  # documents still full
+
+
+def test_deflect_prompt_steering_sits_between_relevance_and_deflect_mode() -> None:
+    prompt = build_deflect_prompt(["Title A", "Title B"], notes=["be concise", "cite paths"])
+    i_rel = prompt.index("<relevance>LOW</relevance>")
+    i_open = prompt.index("<tuning>")
+    i_close = prompt.index("</tuning>")
+    i_mode = prompt.index("DEFLECT_MODE")
+    assert i_rel < i_open < i_close < i_mode
+    assert "1. be concise" in prompt
+    assert "2. cite paths" in prompt
+    assert "- Title A" in prompt  # weak-hit titles still carried
+    assert "DEFLECT_MODE" in prompt
+
+
+# ---------- plan_turn wiring (gate + steering, fake retriever rows) ----------
+
+
+def test_plan_turn_high_mode_injects_notes() -> None:
+    doc = _doc("Kubernetes Homelab Cluster", "TALOS_DOC_SENT")
+    plan = chat_api.plan_turn(
+        [_chunk(doc, 0.90)], _settings(), notes=["be concise", "assume NixOS"]
+    )
+    assert plan.deflected is False
+    assert plan.tuning_count == 2
+    assert "<tuning>" in plan.system_prompt
+    assert "1. be concise" in plan.system_prompt
+    assert "2. assume NixOS" in plan.system_prompt
+    assert "<relevance>HIGH</relevance>" in plan.system_prompt
+    assert "TALOS_DOC_SENT" in plan.system_prompt
+
+
+def test_plan_turn_low_mode_injects_notes() -> None:
+    doc = _doc("Kubernetes Homelab Cluster", "TALOS_DOC_NEVER_SENT")
+    plan = chat_api.plan_turn(
+        [_chunk(doc, 0.10)], _settings(), notes=["be concise"]
+    )
+    assert plan.deflected is True
+    assert plan.tuning_count == 1
+    assert "DEFLECT_MODE" in plan.system_prompt
+    assert "<tuning>" in plan.system_prompt
+    assert "1. be concise" in plan.system_prompt
+    assert "TALOS_DOC_NEVER_SENT" not in plan.system_prompt  # titles only, still
+
+
+def test_plan_turn_without_notes_has_no_tuning_section() -> None:
+    doc = _doc("Kubernetes Homelab Cluster", "TALOS_DOC_SENT")
+    for chunks in (
+        [_chunk(doc, 0.90)],  # HIGH
+        [_chunk(doc, 0.10)],  # LOW
+    ):
+        plan = chat_api.plan_turn(chunks, _settings())
+        assert plan.tuning_count == 0
+        assert "<tuning>" not in plan.system_prompt