brain-of-reese/tests/unit/test_steering_prompts.py

"""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
from datetime import UTC, datetime

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,
        # Phase 106, D5: the HIGH block formats the row's created_at
        # UTC date part — the detached fixture carries it (the NOT NULL
        # DB column guarantees it for real rows).
        created_at=datetime(2024, 6, 15, 12, 0, 0, tzinfo=UTC),
    )


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