"""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). 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 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 == [] 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 == []