brain-of-reese/tests/integration/test_agent_tools.py

"""Integration: the agent DB accessors against real Postgres (phase 37).

``list_catalog`` must order rows by ``(source, path)`` — the same order as
``GET /api/docs`` — and ``find_document`` must resolve a hit to the full
document row (content included, for the never-truncated read) and return
``None`` for unknown ``source``/``path`` pairs. Phase 68: the
``search_documents`` tool is pinned here too — its locked parameter
shape in ``AGENT_TOOLS``, and a scripted ``ToolCallPiece`` executed
through ``run_agent`` against the real DB (``all_documents`` for a
whole-KB search, ``find_document`` for a scoped one). The
combined-form self-correction (a ``source`` argument carrying
``source/path``) is pinned here as well, through ``run_agent``:
the split read executes against the real table, and a still-unknown
split gets the educational refusal.

Requires:  podman compose up -d db
"""
from __future__ import annotations

import asyncio
import uuid
from collections.abc import AsyncIterator, Iterator
from copy import deepcopy
from typing import Any, cast

import pytest
from sqlalchemy import text
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, RetryPiece, StreamPiece, ToolCallPiece


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


@pytest.fixture()
def kb(db) -> Iterator[None]:
    """Fresh documents table (chunks first — the FK) for these accessors."""
    db.execute(text("TRUNCATE chunks, documents"))
    db.commit()
    yield
    db.execute(text("TRUNCATE chunks, documents"))
    db.commit()


def test_list_catalog_orders_by_source_then_path(kb, db) -> None:
    _doc(db, "Zeta", "b/second.md", "Zeta B", "ZB")
    _doc(db, "Zeta", "a/first.md", "Zeta A", "ZA")
    _doc(db, "Alpha", "c/third.md", "Alpha C", "AC")
    db.commit()

    assert agent.list_catalog(db) == [
        ("Alpha", "c/third.md", "Alpha C"),
        ("Zeta", "a/first.md", "Zeta A"),
        ("Zeta", "b/second.md", "Zeta B"),
    ]


def test_list_catalog_is_empty_without_rows(kb, db) -> None:
    assert agent.list_catalog(db) == []


def test_find_document_hit_returns_full_row(kb, db) -> None:
    created = _doc(db, "Alpha", "deep/nested/doc.md", "The Doc", "FULL-TEXT")
    db.commit()

    found = agent.find_document(db, "Alpha", "deep/nested/doc.md")
    assert found is not None
    assert found.id == created.id
    assert found.source == "Alpha"
    assert found.path == "deep/nested/doc.md"
    assert found.title == "The Doc"
    assert found.content == "FULL-TEXT"  # the read tool feeds this, untruncated


def test_find_document_none_for_unknown_pairs(kb, db) -> None:
    _doc(db, "Alpha", "x.md", "X", "X-CONTENT")
    db.commit()

    assert agent.find_document(db, "Alpha", "nope.md") is None  # wrong path
    assert agent.find_document(db, "Beta", "x.md") is None  # wrong source
    assert agent.find_document(db, "nope", "nope.md") is None  # nothing at all


# ---------- search_documents (phase 68) ----------


def test_all_documents_orders_by_source_then_path(kb, db) -> None:
    _doc(db, "Zeta", "b/second.md", "Zeta B", "ZB")
    _doc(db, "Zeta", "a/first.md", "Zeta A", "ZA")
    _doc(db, "Alpha", "c/third.md", "Alpha C", "AC")
    db.commit()

    docs = agent.all_documents(db)
    assert [(d.source, d.path) for d in docs] == [
        ("Alpha", "c/third.md"),
        ("Zeta", "a/first.md"),
        ("Zeta", "b/second.md"),
    ]
    assert [d.content for d in docs] == ["AC", "ZA", "ZB"]  # full rows


def test_agent_tools_offers_search_documents_with_locked_shape() -> None:
    by_name = {t["function"]["name"]: t for t in AGENT_TOOLS}
    assert list(by_name) == [  # the third tool, in order
        "list_documents",
        "read_document",
        "search_documents",
    ]
    search = by_name["search_documents"]["function"]["parameters"]
    assert search["type"] == "object"
    assert search["required"] == ["pattern"]
    assert set(search["properties"]) == {"pattern", "source", "path"}
    assert all(p["type"] == "string" for p in search["properties"].values())


class ScriptedToolLLM:
    """One scripted tool-call stream, then one canned answer stream.
    Records every ``chat_stream`` request's messages and tools."""

    def __init__(self, call: ToolCallPiece) -> None:
        self.call = call
        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), deepcopy(tools)))
        if len(self.requests) == 1:
            yield self.call
        else:
            yield StreamPiece("content", "ans")


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


def _run_search(
    db: Session, arguments: dict[str, Any]
) -> tuple[AgentHolder, ScriptedToolLLM]:
    """Drive one scripted ``search_documents`` call through ``run_agent``."""
    holder = AgentHolder()
    llm = ScriptedToolLLM(
        ToolCallPiece(id="call_1", name="search_documents", arguments=arguments)
    )
    asyncio.run(_consume(llm, db, holder))
    return holder, llm


def _run_read(
    db: Session, arguments: dict[str, Any]
) -> tuple[AgentHolder, ScriptedToolLLM]:
    """Drive one scripted ``read_document`` call through ``run_agent``."""
    holder = AgentHolder()
    llm = ScriptedToolLLM(
        ToolCallPiece(id="call_1", name="read_document", arguments=arguments)
    )
    asyncio.run(_consume(llm, db, holder))
    return holder, llm


async def _consume(
    llm: ScriptedToolLLM, db: Session, holder: AgentHolder
) -> list[StreamPiece | ToolCallPiece | RetryPiece]:
    out: list[StreamPiece | ToolCallPiece | RetryPiece] = []
    async for piece in run_agent(
        cast("LLMClient", llm),
        db,
        system_prompt="SYSTEM_PROMPT",
        user_message="QUESTION",
        seed_docs=[],
        settings=_settings(),
        holder=holder,
    ):
        out.append(piece)
    return out


def test_search_whole_kb_through_run_agent(kb, db) -> None:
    _doc(db, "Beta", "b/two.md", "Two", "no hit\nNEEDLE in two\nlast")
    _doc(db, "Alpha", "a/one.md", "One", "first\nneedle in one\nthird")
    db.commit()

    holder, llm = _run_search(db, {"pattern": "needle"})

    # Offered: the first request carries AGENT_TOOLS (the 3-tool list).
    assert llm.requests[0][1] == AGENT_TOOLS
    # Executed against the real DB: catalog order, grep-style lines.
    assert llm.requests[1][0][3]["content"] == (
        "Alpha/a/one.md:2: needle in one\n"
        "Beta/b/two.md:2: NEEDLE in two"
    )
    assert holder.tool_calls == 1
    assert holder.read_docs == []  # locked A5: search adds no context


def test_search_scoped_through_run_agent(kb, db) -> None:
    _doc(db, "Alpha", "a/one.md", "One", "first\nNeedle here\nthird")
    _doc(db, "Beta", "b/two.md", "Two", "NEEDLE too")
    db.commit()

    holder, llm = _run_search(
        db, {"pattern": "needle", "source": "Alpha", "path": "a/one.md"}
    )

    # Only the named document is searched — the other one's hit is absent.
    assert llm.requests[1][0][3]["content"] == "Alpha/a/one.md:2: Needle here"
    assert holder.tool_calls == 1
    assert holder.read_docs == []


def test_search_scoped_missing_doc_refused_through_run_agent(kb, db) -> None:
    _doc(db, "Alpha", "a/one.md", "One", "nothing")
    db.commit()

    holder, llm = _run_search(
        db, {"pattern": "needle", "source": "Alpha", "path": "ghost.md"}
    )

    assert (
        llm.requests[1][0][3]["content"]
        == "No document at Alpha/ghost.md — check the list_documents output."
    )
    assert holder.tool_calls == 0 and holder.read_docs == []


# ---------- combined 'source/path' self-correction (read_document) ----------


def test_read_combined_source_self_corrects_through_run_agent(kb, db) -> None:
    """The model's combined 'source' ('Alpha/deep/nested/doc.md') resolves
    through the first-slash split against the REAL table: the read
    executes, the holder records the row, the result header carries the
    true source/path."""
    created = _doc(db, "Alpha", "deep/nested/doc.md", "The Doc", "FULL-TEXT")
    db.commit()

    holder, llm = _run_read(
        db,
        {
            "source": "Alpha/deep/nested/doc.md",
            "path": "deep/nested/doc.md",
        },
    )

    assert llm.requests[1][0][3]["content"] == (
        "Document Alpha/deep/nested/doc.md:\nFULL-TEXT"
    )
    assert holder.tool_calls == 1
    assert holder.read_docs == [created]


def test_read_combined_source_later_slash_split_through_run_agent(kb, db) -> None:
    """source='Alpha/deep' + path='nested/doc.md' (a split at a LATER
    slash) resolves via the continuation candidate against the real
    table."""
    created = _doc(db, "Alpha", "deep/nested/doc.md", "The Doc", "FULL-TEXT")
    db.commit()

    holder, llm = _run_read(
        db, {"source": "Alpha/deep", "path": "nested/doc.md"}
    )

    assert llm.requests[1][0][3]["content"] == (
        "Document Alpha/deep/nested/doc.md:\nFULL-TEXT"
    )
    assert holder.tool_calls == 1
    assert holder.read_docs == [created]


def test_read_combined_source_refusal_teaches_split(kb, db) -> None:
    """A combined source that matches nothing (even split) gets the
    educational refusal naming the corrected arguments."""
    _doc(db, "Alpha", "x.md", "X", "X-CONTENT")
    db.commit()

    holder, llm = _run_read(
        db, {"source": "Alpha/nope/deep.md", "path": "nope/deep.md"}
    )

    assert llm.requests[1][0][3]["content"] == (
        "source must not contain '/': for 'Alpha/nope/deep.md' call "
        "read_document(source='Alpha', path='nope/deep.md')."
    )
    assert holder.tool_calls == 0 and holder.read_docs == []


def test_search_no_matches_through_run_agent(kb, db) -> None:
    _doc(db, "Alpha", "a/one.md", "One", "nothing matching")
    db.commit()

    holder, llm = _run_search(db, {"pattern": "zebra"})

    assert llm.requests[1][0][3]["content"] == (
        "No matches for 'zebra' in the knowledge base."
    )
    assert holder.tool_calls == 1  # an executed search with zero hits
    assert holder.read_docs == []