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

"""Unit: the phase-106 D5 date surfaces (task 06) — the date rides every
document the model sees.

Three surfaces, pinned here against pure inputs (no database):

* the HIGH prompt's ``<document>`` block carries
  ``date="YYYY-MM-DD"`` — the row's ``created_at`` UTC date part,
  APPENDED after ``title`` (the only position; always present,
  ``created_at`` is NOT NULL) — while the DEFLECTION prompt stays
  byte-identical to the pre-phase text (the A8 contract — titles only,
  no documents, no dates);
* the ``read`` tool result carries ``date: YYYY-MM-DD`` as its SECOND
  line — the FIRST line stays ``Document {source}/{path}:``
  BYTE-IDENTICAL (the E2E mock's ``_READ_RESULT_PREFIX`` header
  contract) — in BOTH the plain and the truncated shapes (marker +
  notice still follow the cut content);
* the ``ls`` FILE line ENDS with the appended `` | date: YYYY-MM-DD``
  field (never inserted before ``title``) and the 50-line cap note is
  unchanged; source/folder lines carry no date.
"""
from __future__ import annotations

import asyncio
import uuid
from collections.abc import AsyncIterator
from datetime import UTC, datetime
from typing import Any, cast

import pytest
from sqlalchemy.orm import Session

from app.config import Settings
from app.models import Document
from app.rag import agent
from app.rag.agent import READ_TRUNCATION_NOTICE, AgentHolder, run_agent
from app.rag.llm import (
    LLMClient,
    StreamPiece,
    ToolCallPiece,
)
from app.rag.prompts import (
    _base,
    build_deflect_prompt,
    build_high_prompt,
    build_steering_section,
)
from app.rag.retriever import TRUNCATION_MARKER

#: A fixed creation date; the UTC date part is ``2024-06-15``.
CREATED_AT = datetime(2024, 6, 15, 12, 30, 45, tzinfo=UTC)
DATE = "2024-06-15"


def _doc(
    source: str = "S",
    path: str = "P",
    title: str = "T",
    content: str = "CONTENT",
    created_at: datetime = CREATED_AT,
) -> Document:
    return Document(
        id=uuid.uuid4(),
        source=source,
        path=path,
        full_path=f"/tmp/{path}",
        title=title,
        content=content,
        content_hash="0" * 64,
        created_at=created_at,
    )


def _settings(**kwargs: Any) -> Settings:
    kwargs.setdefault("_env_file", None)
    return Settings(**kwargs)  # pyright: ignore[reportCallIssue]


class ScriptedLLM:
    """Canned stream sequences; records every ``chat_stream`` request."""

    def __init__(self, *streams: list[StreamPiece | ToolCallPiece]) -> None:
        self.streams: list[list[StreamPiece | ToolCallPiece]] = list(streams)
        self.requests: list[tuple[list[dict[str, Any]], list[dict[str, Any]] | None]] = []

    async def chat_stream(
        self,
        messages: list[dict[str, str]],
        tools: list[dict[str, Any]] | None = None,
        scaffolding: object = None,
    ) -> AsyncIterator[StreamPiece | ToolCallPiece]:
        self.requests.append((messages, tools))
        if not self.streams:
            return
        for piece in self.streams.pop(0):
            yield piece


async def _run(
    llm: ScriptedLLM,
    holder: AgentHolder,
    settings: Settings,
) -> None:
    async for _piece in run_agent(
        cast("LLMClient", llm),
        cast("Session", None),
        system_prompt="SYSTEM_PROMPT",
        user_message="QUESTION",
        seed_docs=[],
        settings=settings,
        holder=holder,
    ):
        pass


# --------------------------------------------------------------------
# The HIGH prompt's <document> block — date="YYYY-MM-DD" after title
# --------------------------------------------------------------------


def test_high_block_carries_date_attribute_after_title() -> None:
    """The block's identity attributes are ``source`` / ``path`` /
    ``title`` — and, phase 106 (D5), ``date`` APPENDED after ``title``
    (the only position), always present."""
    doc = _doc(source="S", path="P", title="T", content="THE BODY")
    prompt = build_high_prompt([doc])
    block = (
        f'<document source="S" path="P" title="T" date="{DATE}">\n'
        "THE BODY\n"
        "</document>"
    )
    assert block in prompt
    # The attribute order is pinned: date directly after title, then
    # the closing bracket (nothing may land between title and date or
    # after date before the block opens).
    assert f'title="T" date="{DATE}">' in prompt
    assert '<document source="S" path="P" title="T">\n' not in prompt  # pre-phase shape is gone


def test_high_block_date_is_the_utc_date_part() -> None:
    """The date is the row's ``created_at`` UTC DATE part: a late UTC
    instant (23:59:59 on 2024-06-15) renders ``2024-06-15`` — and the
    time-of-day is never shown (full precision is stored; only the
    date part rides the block). Stored rows are always UTC-aware
    (``timestamptz`` read-back + :func:`normalize_doc_date`), so the
    UTC date part is the stored date part."""
    doc = _doc(created_at=datetime(2024, 6, 15, 23, 59, 59, tzinfo=UTC))
    prompt = build_high_prompt([doc])
    assert f'title="T" date="{DATE}">' in prompt
    assert "2024-06-16" not in prompt
    # The time-of-day is never rendered on the block.
    assert "23:59" not in prompt and "12:30" not in prompt


def test_high_block_date_always_present_for_every_document() -> None:
    """Every document block carries the date (``created_at`` is NOT
    NULL — the attribute is never omitted), each with its OWN row's
    date."""
    a = _doc(source="A", path="a.md", title="A", content="CA", created_at=CREATED_AT)
    b = _doc(
        source="B",
        path="b.md",
        title="B",
        content="CB",
        created_at=datetime(2019, 6, 15, 5, 0, 0, tzinfo=UTC),
    )
    prompt = build_high_prompt([a, b])
    assert '<document source="A" path="a.md" title="A" date="2024-06-15">' in prompt
    assert '<document source="B" path="b.md" title="B" date="2019-06-15">' in prompt
    # No date-less block anywhere.
    assert prompt.count("<document ") == prompt.count(' date="')


# --------------------------------------------------------------------
# The deflection prompt — byte-identical to the pre-phase text (A8)
# --------------------------------------------------------------------


def test_deflect_prompt_is_byte_identical_to_pre_phase() -> None:
    """The deflection path is titles-only (no documents involved): the
    A8 byte-identity contract holds — same inputs, pre-phase bytes,
    with or without steering notes."""
    # No notes: the pre-phase LOW build verbatim.
    expected_plain = (
        _base("LOW")
        + "\n"
        + "DEFLECT_MODE: retrieval was weak — the titles below are the closest "
        "your notes come to the question. They are titles only; do not pretend "
        "they answer it. Use them to propose 2-3 alternative questions.\n"
        + "Reply in plain text only — you have no tools in this mode.\n"
        + "- T1\n- T2"
    )
    assert build_deflect_prompt(["T1", "T2"]) == expected_plain
    # With steering notes: the pre-phase text + the <tuning> section
    # (phase 15) — still no date anywhere.
    expected_steered = (
        _base("LOW")
        + "\n"
        + build_steering_section(["be concise"])
        + "\n"
        + "DEFLECT_MODE: retrieval was weak — the titles below are the closest "
        "your notes come to the question. They are titles only; do not pretend "
        "they answer it. Use them to propose 2-3 alternative questions.\n"
        + "Reply in plain text only — you have no tools in this mode.\n"
        + "- T1\n- T2"
    )
    assert build_deflect_prompt(["T1", "T2"], notes=["be concise"]) == expected_steered
    for prompt in (expected_plain, expected_steered):
        assert "<document" not in prompt
        assert 'date="' not in prompt


# --------------------------------------------------------------------
# The read result — the date is the SECOND line (first line
# byte-identical — the mock's _READ_RESULT_PREFIX header contract)
# --------------------------------------------------------------------


def test_read_result_plain_carries_date_second_line(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """At or under the cap: ``Document {source}/{path}:`` (byte-
    identical FIRST line) + ``date: YYYY-MM-DD`` (the D5 SECOND line)
    + the full content — nothing else."""
    doc = _doc(source="S", path="a.md", title="A", content="A-CONTENT")
    monkeypatch.setattr(agent, "find_document", lambda db, source, path: doc)
    holder = AgentHolder()
    llm = ScriptedLLM(
        [ToolCallPiece(id="call_1", name="read", arguments={"path": "S/a.md"})],
        [StreamPiece("content", "ans")],
    )
    asyncio.run(_run(llm, holder, _settings()))
    assert holder.read_docs == [doc] and holder.tool_calls == 1
    content = llm.requests[1][0][3]["content"]
    lines = content.splitlines()
    assert lines[0] == "Document S/a.md:"  # byte-identical first line
    assert lines[1] == f"date: {DATE}"  # the D5 second line
    assert lines[2:] == ["A-CONTENT"]


def test_read_result_truncated_carries_date_then_marker_and_notice(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Over the cap: the SAME two header lines, then the first ``cap``
    chars, the shared marker, and the pinned notice — the truncation
    contract (phase 95) is unchanged by the date line."""
    cap = 10
    content_body = "x" * 25
    doc = _doc(source="S", path="big.md", title="Big", content=content_body)
    monkeypatch.setattr(agent, "find_document", lambda db, source, path: doc)
    holder = AgentHolder()
    llm = ScriptedLLM(
        [ToolCallPiece(id="call_1", name="read", arguments={"path": "S/big.md"})],
        [StreamPiece("content", "ans")],
    )
    asyncio.run(_run(llm, holder, _settings(read_max_chars=cap)))
    expected = (
        "Document S/big.md:\n"
        f"date: {DATE}\n"
        + content_body[:cap]
        + "\n"
        + TRUNCATION_MARKER
        + "\n"
        + READ_TRUNCATION_NOTICE.format(shown=cap, total=len(content_body))
    )
    assert llm.requests[1][0][3]["content"] == expected
    # The truncation still records on the holder (the ToolResultPiece
    # signal) — the date line adds nothing to the counters.
    assert holder.read_truncations == [("S/big.md", cap, len(content_body))]
    assert holder.tool_calls == 1


# --------------------------------------------------------------------
# The ls FILE line — the appended " | date: YYYY-MM-DD" field
# --------------------------------------------------------------------


def test_ls_file_line_ends_with_appended_date() -> None:
    """The ``ls`` FILE line is ``source: X | path: Y | title: Z`` with
    the phase-106 D5 field APPENDED at the END (never inserted before
    ``title`` — the mock's non-greedy ``path`` capture would swallow an
    inserted field)."""
    rows = [
        ("a.md", "A", "2019-06-15"),
        ("b.md", "B", "2024-06-15"),
    ]
    subfolders, files, total = agent.group_folder_listing("S", "", rows, {})
    rendered = agent.render_folder_listing("S", subfolders, files, total)
    assert total == 2
    lines = rendered.splitlines()
    assert lines[0] == "S — 2 documents, 0 folders:"
    # The file lines end with the appended date field — per-row dates.
    assert lines[2] == "source: S | path: a.md | title: A | date: 2019-06-15"
    assert lines[3] == "source: S | path: b.md | title: B | date: 2024-06-15"
    # The date field is the LAST field (appended, never inserted before
    # ``title`` — where the mock's non-greedy path capture would
    # swallow it).
    for line in (lines[2], lines[3]):
        assert line.index("date:") > line.index("title:")
        assert line.rsplit(" | ", 1)[-1].startswith("date: ")  # date is final


def test_ls_folder_and_source_lines_carry_no_date() -> None:
    """Only FILE lines are documents: the folder header, the indented
    subfolder lines, and the top-level source lines carry NO date."""
    rows = [
        ("a/x.md", "X", "2019-06-15"),
        ("a/y.md", "Y", "2020-01-02"),
        ("top.md", "Top", "2024-06-15"),
    ]
    subfolders, files, total = agent.group_folder_listing("S", "", rows, {"a": "A stuff."})
    rendered = agent.render_folder_listing("S", subfolders, files, total)
    lines = rendered.splitlines()
    assert lines[0] == "S — 1 documents, 1 folders:"  # header: no date
    assert lines[2] == "  a/ — 2 documents: A stuff."  # subfolder: no date
    # Only the DIRECT root file lists (a/ is a subfolder line — the
    # drill-down is one level per call), and it alone carries a date.
    file_lines = [line for line in lines if line.startswith("source: ")]
    assert file_lines == [
        "source: S | path: top.md | title: Top | date: 2024-06-15"
    ]
    # The top level (source lines) is byte-identical to the pre-phase
    # template — no date on source lines.
    assert agent.render_ls_top([("S", 3, "Source stuff.")]) == (
        "1 sources:\n\nS — 3 documents\n  Source stuff."
    )
    assert "date" not in agent.render_ls_top([("S", 3, None)])


def test_ls_cap_note_unchanged_past_fifty() -> None:
    """The 50-line cap note is UNCHANGED by the date field: 51 files →
    50 lines (each with its appended date) + the exact pre-phase note;
    the line COUNT is unaffected by the wider lines."""
    rows = [(f"f{i:03d}.md", f"T{i}", "2024-06-15") for i in range(51)]
    subfolders, files, total = agent.group_folder_listing("S", "", rows, {})
    rendered = agent.render_folder_listing("S", subfolders, files, total)
    lines = rendered.splitlines()
    assert total == 51 and len(lines) == 1 + 1 + 50 + 1  # header, blank, 50, note
    assert lines[-1] == (
        "…and 1 more documents in this folder — use grep (pattern) to "
        "find a specific one."
    )
    assert lines[1] == ""  # the blank after the header is untouched
    assert all(
        line.endswith(" | date: 2024-06-15")
        for line in lines[2:52]
    )