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

"""Unit: the grounded-turn agent loop (phase 37, ``app.rag.agent``).

A scripted fake LLM (canned stream sequences) + monkeypatched
``list_catalog`` / ``find_document`` — no database, no network. Covers
the loop mechanics: the list → read → answer happy path (event order,
holder state, the ``tools=None`` request after the budgets are spent,
the assistant/tool message history), the 0/0 single-call path, budget
exhaustion, dedupe, unknown tool / missing args / unknown path, the
round cap, and the ``<tools>`` prompt section (HIGH only).
"""
from __future__ import annotations

import asyncio
import json
import uuid
from collections.abc import AsyncIterator
from copy import deepcopy
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 (
    AGENT_TOOLS,
    AgentHolder,
    run_agent,
)
from app.rag.llm import LLMClient, StreamPiece, ToolCallPiece
from app.rag.prompts import TOOLS_SECTION, _base, build_deflect_prompt, build_high_prompt


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


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


class ScriptedLLM:
    """Canned stream sequences; records every ``chat_stream`` request so
    the tests can assert on the messages and the ``tools`` passthrough."""

    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,
    ) -> AsyncIterator[StreamPiece | ToolCallPiece]:
        self.requests.append((deepcopy(messages), tools))
        if not self.streams:
            raise AssertionError("ScriptedLLM ran out of canned streams")
        for piece in self.streams.pop(0):
            yield piece


async def _run(
    llm: ScriptedLLM,
    holder: AgentHolder,
    settings: Settings,
    seed_docs: list[Document] | None = None,
) -> list[StreamPiece | ToolCallPiece]:
    out: list[StreamPiece | ToolCallPiece] = []
    async for piece in run_agent(
        cast("LLMClient", llm),
        cast("Session", None),
        system_prompt="SYSTEM_PROMPT",
        user_message="QUESTION",
        seed_docs=seed_docs or [],
        settings=settings,
        holder=holder,
    ):
        out.append(piece)
    return out


# ---------- AGENT_TOOLS shape ----------


def test_agent_tools_names_and_parameters() -> None:
    by_name = {t["function"]["name"]: t for t in AGENT_TOOLS}
    assert set(by_name) == {"list_documents", "read_document"}
    assert all(t["type"] == "function" for t in AGENT_TOOLS)
    list_params = by_name["list_documents"]["function"]["parameters"]
    assert list_params["type"] == "object"
    assert list_params["properties"] == {}  # no parameters
    read_params = by_name["read_document"]["function"]["parameters"]
    assert read_params["required"] == ["source", "path"]
    assert set(read_params["properties"]) == {"source", "path"}


# ---------- happy path: list → read → answer ----------


def test_list_then_read_then_answer(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    catalog = [
        ("Deployments", "backups.md", "Backup Strategy"),
        ("Homelab", "aws-route53.md", "AWS Route53 Records"),
    ]
    monkeypatch.setattr(agent, "list_catalog", lambda db: catalog)
    target = _doc("Homelab", "aws-route53.md", "AWS Route53 Records", "R53-CONTENT")
    monkeypatch.setattr(agent, "find_document", lambda db, source, path: target)
    seed = [_doc("Homelab", "kubernetes.md", "Kubernetes", "K8S-CONTENT")]
    holder = AgentHolder()
    llm = ScriptedLLM(
        [ToolCallPiece(id="call_1", name="list_documents", arguments={})],
        [
            ToolCallPiece(
                id="call_2",
                name="read_document",
                arguments={"source": "Homelab", "path": "aws-route53.md"},
            )
        ],
        [StreamPiece("thinking", "hmm "), StreamPiece("content", "Done! ")],
    )

    pieces = asyncio.run(_run(llm, holder, _settings(), seed_docs=seed))

    # Event order: tool pieces before the answer content/thinking.
    assert [type(p) for p in pieces] == [
        ToolCallPiece,
        ToolCallPiece,
        StreamPiece,
        StreamPiece,
    ]
    assert pieces[0] == ToolCallPiece(id="call_1", name="list_documents", arguments={})
    assert isinstance(pieces[1], ToolCallPiece)
    assert pieces[1].name == "read_document"
    assert pieces[3] == StreamPiece("content", "Done! ")
    # The read document is recorded for done.sources / query_log (task 04).
    assert holder.read_docs == [target]
    assert holder.tool_calls == 2

    # Default budgets (1/1): tools offered while any budget remains…
    assert llm.requests[0][1] == AGENT_TOOLS
    assert llm.requests[1][1] == AGENT_TOOLS
    # …and dropped (tools=None) once both are spent.
    assert llm.requests[2][1] is None
    assert len(llm.requests) == 3

    # The follow-up request carries the assistant tool-call + tool result.
    msgs = llm.requests[1][0]
    assert msgs[0] == {"role": "system", "content": "SYSTEM_PROMPT"}
    assert msgs[1] == {"role": "user", "content": "QUESTION"}
    assert msgs[2] == {
        "role": "assistant",
        "content": None,
        "tool_calls": [
            {
                "id": "call_1",
                "type": "function",
                "function": {"name": "list_documents", "arguments": "{}"},
            }
        ],
    }
    assert msgs[3] == {
        "role": "tool",
        "tool_call_id": "call_1",
        "content": (
            "2 documents:\n"
            "Deployments/backups.md — Backup Strategy\n"
            "Homelab/aws-route53.md — AWS Route53 Records"
        ),
    }
    # The second follow-up request carries the read call + the FULL text.
    msgs = llm.requests[2][0]
    assert msgs[4]["role"] == "assistant"
    assert msgs[4]["tool_calls"][0]["id"] == "call_2"
    assert json.loads(msgs[4]["tool_calls"][0]["function"]["arguments"]) == {
        "source": "Homelab",
        "path": "aws-route53.md",
    }
    assert msgs[5] == {
        "role": "tool",
        "tool_call_id": "call_2",
        "content": "Document Homelab/aws-route53.md:\nR53-CONTENT",  # full text, no cap
    }


def test_empty_catalog_listing_says_zero_documents(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(agent, "list_catalog", lambda db: [])
    holder = AgentHolder()
    llm = ScriptedLLM(
        [ToolCallPiece(id="call_1", name="list_documents", arguments={})],
        [StreamPiece("content", "ans")],
    )
    asyncio.run(_run(llm, holder, _settings()))
    assert llm.requests[1][0][3]["content"] == "0 documents:\n"
    assert holder.tool_calls == 1


def test_content_and_tool_call_in_one_stream_keeps_both(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Rare stream with content AND a tool call: the content stays (it was
    already emitted) and the tool still runs."""
    monkeypatch.setattr(agent, "list_catalog", lambda db: [])
    holder = AgentHolder()
    llm = ScriptedLLM(
        [
            StreamPiece("content", "Let me check "),
            ToolCallPiece(id="call_1", name="list_documents", arguments={}),
        ],
        [StreamPiece("content", "the answer")],
    )
    pieces = asyncio.run(_run(llm, holder, _settings()))
    assert [type(p) for p in pieces] == [StreamPiece, ToolCallPiece, StreamPiece]
    assert holder.tool_calls == 1  # the tool ran despite the content
    assert llm.requests[1][0][3]["content"] == "0 documents:\n"


# ---------- budgets ----------


def test_zero_budgets_is_one_request_without_tools() -> None:
    """BOR_AGENT_LIST_CALLS=0 BOR_AGENT_READ_CALLS=0 → byte-identical
    single-call path: exactly one request, tools=None, no history growth."""
    holder = AgentHolder()
    llm = ScriptedLLM([StreamPiece("thinking", "t "), StreamPiece("content", "direct answer")])
    pieces = asyncio.run(
        _run(llm, holder, _settings(agent_list_calls=0, agent_read_calls=0))
    )
    assert [type(p) for p in pieces] == [StreamPiece, StreamPiece]
    assert len(llm.requests) == 1
    assert llm.requests[0][1] is None
    assert llm.requests[0][0] == [
        {"role": "system", "content": "SYSTEM_PROMPT"},
        {"role": "user", "content": "QUESTION"},
    ]
    assert holder.read_docs == [] and holder.tool_calls == 0


def test_read_budget_exhausted_refuses_and_appends_nothing(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    a = _doc("S", "a.md", "A", "A-CONTENT")
    monkeypatch.setattr(
        agent, "find_document", lambda db, source, path: a if path == "a.md" else None
    )
    holder = AgentHolder()
    llm = ScriptedLLM(
        [
            ToolCallPiece(
                id="call_1", name="read_document", arguments={"source": "S", "path": "a.md"}
            )
        ],
        [
            ToolCallPiece(
                id="call_2", name="read_document", arguments={"source": "S", "path": "b.md"}
            )
        ],
        [StreamPiece("content", "ans")],
    )
    asyncio.run(_run(llm, holder, _settings(agent_list_calls=1, agent_read_calls=1)))

    assert holder.read_docs == [a]  # the refused read appended nothing
    assert holder.tool_calls == 1  # …and consumed no budget
    refusal = llm.requests[2][0][5]
    assert refusal == {
        "role": "tool",
        "tool_call_id": "call_2",
        "content": agent.READ_EXHAUSTED,
    }
    # The list budget is still open, so tools stay offered after the refusal.
    assert llm.requests[2][1] == AGENT_TOOLS


def test_list_budget_exhausted_refuses_with_its_own_message(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(agent, "list_catalog", lambda db: [])
    holder = AgentHolder()
    llm = ScriptedLLM(
        [ToolCallPiece(id="call_1", name="list_documents", arguments={})],
        [ToolCallPiece(id="call_2", name="list_documents", arguments={})],
        [StreamPiece("content", "ans")],
    )
    asyncio.run(_run(llm, holder, _settings(agent_list_calls=1, agent_read_calls=1)))
    assert holder.tool_calls == 1
    assert llm.requests[2][0][5]["content"] == agent.LIST_EXHAUSTED
    # The read budget is still open, so tools stay offered after the refusal.
    assert llm.requests[2][1] == AGENT_TOOLS


# ---------- rejections (no budget consumed) ----------


def test_reading_a_seed_doc_is_already_in_context(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    seed = [_doc("Homelab", "kubernetes.md", "Kubernetes", "K8S-CONTENT")]

    def _boom(*_a: Any, **_k: Any) -> None:
        raise AssertionError("find_document must not be called for a seeded doc")

    monkeypatch.setattr(agent, "list_catalog", lambda db: [])
    monkeypatch.setattr(agent, "find_document", _boom)
    holder = AgentHolder()
    llm = ScriptedLLM(
        [
            ToolCallPiece(
                id="call_1",
                name="read_document",
                arguments={"source": "Homelab", "path": "kubernetes.md"},
            )
        ],
        [StreamPiece("content", "ans")],
    )
    asyncio.run(_run(llm, holder, _settings(), seed_docs=seed))
    assert holder.read_docs == [] and holder.tool_calls == 0
    assert llm.requests[1][0][3]["content"] == agent.ALREADY_IN_CONTEXT
    # No budget consumed → tools are still offered on the next request.
    assert llm.requests[1][1] == AGENT_TOOLS


def test_reading_an_already_read_doc_is_deduped(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    doc = _doc("S", "a.md", "A", "A-CONTENT")
    monkeypatch.setattr(agent, "find_document", lambda db, source, path: doc)
    holder = AgentHolder()
    llm = ScriptedLLM(
        [
            ToolCallPiece(
                id="call_1", name="read_document", arguments={"source": "S", "path": "a.md"}
            )
        ],
        [
            ToolCallPiece(
                id="call_2", name="read_document", arguments={"source": "S", "path": "a.md"}
            )
        ],
        [StreamPiece("content", "ans")],
    )
    asyncio.run(_run(llm, holder, _settings(agent_list_calls=1, agent_read_calls=1)))
    assert holder.read_docs == [doc]  # appended exactly once
    assert holder.tool_calls == 1
    assert llm.requests[2][0][5]["content"] == agent.ALREADY_IN_CONTEXT
    # The read budget is intact after the deduped refusal…
    assert llm.requests[2][1] == AGENT_TOOLS


def test_unknown_path_refused_without_budget(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(agent, "find_document", lambda db, source, path: None)
    holder = AgentHolder()
    llm = ScriptedLLM(
        [
            ToolCallPiece(
                id="call_1",
                name="read_document",
                arguments={"source": "S", "path": "ghost.md"},
            )
        ],
        [StreamPiece("content", "ans")],
    )
    asyncio.run(_run(llm, holder, _settings()))
    assert holder.read_docs == [] and holder.tool_calls == 0
    assert (
        llm.requests[1][0][3]["content"]
        == "No document at S/ghost.md — check the list_documents output."
    )
    assert llm.requests[1][1] == AGENT_TOOLS  # budget intact


def test_unknown_tool_name_refused(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(agent, "list_catalog", lambda db: [])
    holder = AgentHolder()
    llm = ScriptedLLM(
        [ToolCallPiece(id="call_1", name="delete_universe", arguments={"x": 1})],
        [StreamPiece("content", "ans")],
    )
    asyncio.run(_run(llm, holder, _settings()))
    assert holder.read_docs == [] and holder.tool_calls == 0
    assert llm.requests[1][0][3]["content"] == agent.UNKNOWN_TOOL
    assert llm.requests[1][1] == AGENT_TOOLS  # nothing was consumed


@pytest.mark.parametrize(
    ("arguments", "label"),
    [
        ({}, "no arguments"),
        ({"source": "S"}, "path missing"),
        ({"path": "p.md"}, "source missing"),
        ({"source": "", "path": "p.md"}, "empty source"),
        ({"source": "S", "path": "  "}, "blank path"),
        ({"source": 7, "path": "p.md"}, "non-string source"),
    ],
)
def test_read_document_missing_arguments_refused(
    monkeypatch: pytest.MonkeyPatch, arguments: dict[str, Any], label: str
) -> None:
    def _boom(*_a: Any, **_k: Any) -> None:
        raise AssertionError(f"find_document must not be called ({label})")

    monkeypatch.setattr(agent, "list_catalog", lambda db: [])
    monkeypatch.setattr(agent, "find_document", _boom)
    holder = AgentHolder()
    llm = ScriptedLLM(
        [ToolCallPiece(id="call_1", name="read_document", arguments=arguments)],
        [StreamPiece("content", "ans")],
    )
    asyncio.run(_run(llm, holder, _settings()))
    assert holder.read_docs == [] and holder.tool_calls == 0
    assert llm.requests[1][0][3]["content"] == agent.MISSING_READ_ARGS
    assert llm.requests[1][1] == AGENT_TOOLS


# ---------- round cap (pathological stream) ----------


def test_round_cap_forces_a_final_no_tools_answer(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A model that keeps calling a budget-exhausted tool must be forced
    to answer at ``max_rounds = 2 + list + read`` (= 4 for 1/1)."""
    monkeypatch.setattr(agent, "list_catalog", lambda db: [])
    holder = AgentHolder()
    llm = ScriptedLLM(
        [ToolCallPiece(id="call_1", name="list_documents", arguments={})],
        [ToolCallPiece(id="call_2", name="list_documents", arguments={})],
        [ToolCallPiece(id="call_3", name="list_documents", arguments={})],
        [ToolCallPiece(id="call_4", name="list_documents", arguments={})],
        [StreamPiece("content", "forced answer")],
    )
    pieces = asyncio.run(_run(llm, holder, _settings(agent_list_calls=1, agent_read_calls=1)))
    assert [type(p) for p in pieces] == [
        ToolCallPiece,
        ToolCallPiece,
        ToolCallPiece,
        ToolCallPiece,
        StreamPiece,
    ]
    assert len(llm.requests) == 5
    # The forced final request carries no tools, whatever is left.
    assert llm.requests[4][1] is None
    # Only the first call consumed budget; the three rejections did not.
    assert holder.tool_calls == 1
    # The 4th rejection sits at messages[2 + 4*2 - 1] of the final request.
    assert llm.requests[4][0][9]["content"] == agent.LIST_EXHAUSTED


# ---------- settings ----------


def test_agent_budget_settings_default_to_one_each() -> None:
    s = _settings()
    assert s.agent_list_calls == 1
    assert s.agent_read_calls == 1


def test_agent_budget_settings_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("BOR_AGENT_LIST_CALLS", "0")
    monkeypatch.setenv("BOR_AGENT_READ_CALLS", "2")
    s = _settings()
    assert s.agent_list_calls == 0
    assert s.agent_read_calls == 2


# ---------- prompts: <tools> section (HIGH only) ----------


def test_high_prompt_carries_tools_section_after_documents() -> None:
    prompt = build_high_prompt([_doc("S", "a.md", "A", "A-CONTENT")])
    assert TOOLS_SECTION in prompt
    assert "call `list_documents`" in prompt
    assert "then `read_document` to pull in exactly one more document" in prompt
    assert "do not read more than one extra document" in prompt
    # After the mode body: <tools> follows </documents>.
    assert prompt.index("</documents>") < prompt.index("<tools>")
    assert prompt.rstrip().endswith("</tools>")


def test_high_prompt_tools_section_with_notes_and_kb() -> None:
    prompt = build_high_prompt(
        [_doc("S", "a.md", "A", "A-CONTENT")], notes=["be concise"], kb_overview="- KB"
    )
    assert prompt.index("<knowledge_base>") < prompt.index("<tuning>")
    assert prompt.index("<tuning>") < prompt.index("<documents>")
    assert prompt.index("<documents>") < prompt.index("<tools>")


def test_low_prompt_is_byte_identical_and_tool_free() -> None:
    expected = (
        _base("LOW")
        + "\nDEFLECT_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"
        + "- T1\n- T2"
    )
    assert build_deflect_prompt(["T1", "T2"]) == expected
    for prompt in (
        build_deflect_prompt(["T1"]),
        build_deflect_prompt(["T1"], notes=["be concise"]),
        build_deflect_prompt(["T1"], kb_overview="- KB"),
        build_deflect_prompt(["T1"], notes=["be concise"], kb_overview="- KB"),
    ):
        assert "<tools>" not in prompt
        assert TOOLS_SECTION not in prompt