367 lines
13 KiB
Python
367 lines
13 KiB
Python
"""Integration: the agent DB accessors against real Postgres (phase 37;
|
|
the harness-aligned ``ls``/``read``/``grep`` surface, phase 70).
|
|
|
|
``list_catalog`` must order rows by ``(source, path)`` — the same order
|
|
as ``GET /api/docs`` — ``list_source_names`` must resolve the
|
|
registered source names (the scoped ``ls`` join), and ``find_document``
|
|
must resolve a hit to the full document row (content included, for the
|
|
never-truncated read) and return ``None`` for unknown pairs. Phase 70:
|
|
the ``ls``/``read``/``grep`` tools are pinned here too — the locked
|
|
parameter shape in ``AGENT_TOOLS``, and scripted ``ToolCallPiece``s
|
|
executed through ``run_agent`` against the real DB: ``ls`` scoped to a
|
|
registered source name (unknown name → refusal), ``read`` on the
|
|
canonical combined ``source/path`` form (first-slash split; a bare
|
|
source name and an unknown identity get the no-document refusal), and
|
|
``grep`` (``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 delete, text
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.config import Settings
|
|
from app.models import Document, GitSource
|
|
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()
|
|
|
|
|
|
@pytest.fixture()
|
|
def src(db) -> Iterator[GitSource]:
|
|
"""One registered git source — the scoped ``ls`` source-name check
|
|
reads the real registry, so the row is inserted and deleted around
|
|
the tests (``repo_name`` resolves the URL to ``Homelab``)."""
|
|
row = GitSource(url="https://github.com/reese/Homelab.git", kind="git")
|
|
db.add(row)
|
|
db.commit()
|
|
yield row
|
|
db.execute(delete(GitSource).where(GitSource.id == row.id))
|
|
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_list_source_names_resolves_registry_rows(db) -> None:
|
|
"""The real registry: git names resolve through the import pipeline's
|
|
``repo_name`` (trailing ``.git`` stripped); a second row resolving to
|
|
the same name (the phase-69 sibling case) is listed once."""
|
|
a = GitSource(url="https://github.com/reese/Homelab.git", kind="git")
|
|
b = GitSource(url="https://github.com/reese/Homelab", kind="git") # sibling
|
|
c = GitSource(url="https://e.com/deployments", kind="git")
|
|
db.add_all([a, b, c])
|
|
db.commit()
|
|
try:
|
|
assert agent.list_source_names(db).count("Homelab") == 1 # deduped
|
|
assert "deployments" in agent.list_source_names(db)
|
|
finally:
|
|
db.execute(delete(GitSource).where(GitSource.id.in_([a.id, b.id, c.id])))
|
|
db.commit()
|
|
|
|
|
|
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
|
|
|
|
|
|
# ---------- AGENT_TOOLS surface (phase 70: ls / read / grep) ----------
|
|
|
|
|
|
def test_agent_tools_offers_the_harness_aligned_surface() -> None:
|
|
by_name = {t["function"]["name"]: t for t in AGENT_TOOLS}
|
|
assert list(by_name) == [ # the harness order, phase 70
|
|
"ls",
|
|
"read",
|
|
"grep",
|
|
]
|
|
ls = by_name["ls"]["function"]["parameters"]
|
|
assert ls["type"] == "object"
|
|
assert ls["required"] == [] # path is optional
|
|
assert set(ls["properties"]) == {"path"}
|
|
read = by_name["read"]["function"]["parameters"]
|
|
assert read["type"] == "object"
|
|
assert read["required"] == ["path"]
|
|
assert set(read["properties"]) == {"path"}
|
|
grep = by_name["grep"]["function"]["parameters"]
|
|
assert grep["type"] == "object"
|
|
assert grep["required"] == ["pattern"]
|
|
assert set(grep["properties"]) == {"pattern", "path"}
|
|
assert all(p["type"] == "string" for p in grep["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_call(
|
|
db: Session, name: str, arguments: dict[str, Any]
|
|
) -> tuple[AgentHolder, ScriptedToolLLM]:
|
|
"""Drive one scripted tool call through ``run_agent``."""
|
|
holder = AgentHolder()
|
|
llm = ScriptedToolLLM(ToolCallPiece(id="call_1", name=name, 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
|
|
|
|
|
|
# ---------- ls (scoped through the real registry) ----------
|
|
|
|
|
|
def test_ls_scoped_to_registered_source_through_run_agent(kb, src, db) -> None:
|
|
_doc(db, "Homelab", "a.md", "A", "A-CONTENT")
|
|
_doc(db, "Other", "b.md", "B", "B-CONTENT")
|
|
db.commit()
|
|
|
|
holder, llm = _run_call(db, "ls", {"path": "Homelab"})
|
|
|
|
# Offered: the first request carries AGENT_TOOLS (the 3-tool list).
|
|
assert llm.requests[0][1] == AGENT_TOOLS
|
|
# Executed against the real DB: the listing filtered to the source.
|
|
assert llm.requests[1][0][3]["content"] == (
|
|
"1 documents:\nsource: Homelab | path: a.md | title: A"
|
|
)
|
|
assert holder.tool_calls == 1
|
|
assert holder.read_docs == []
|
|
|
|
|
|
def test_ls_scoped_unknown_source_refused_through_run_agent(kb, src, db) -> None:
|
|
_doc(db, "Homelab", "a.md", "A", "A-CONTENT")
|
|
db.commit()
|
|
|
|
holder, llm = _run_call(db, "ls", {"path": "Ghost"})
|
|
|
|
assert (
|
|
llm.requests[1][0][3]["content"] == "No source named 'Ghost' — check the ls output."
|
|
)
|
|
assert holder.tool_calls == 0 and holder.read_docs == []
|
|
|
|
|
|
# ---------- read (the canonical combined source/path form) ----------
|
|
|
|
|
|
def test_read_combined_path_through_run_agent(kb, db) -> None:
|
|
"""The combined ``source/path`` identity resolves at the FIRST slash
|
|
against the REAL table (a path with further slashes included): 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_call(db, "read", {"path": "Alpha/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_bare_source_name_refused_through_run_agent(kb, db) -> None:
|
|
"""A bare source name (no '/') can never be a document — the
|
|
no-document refusal echoing the argument as passed; the old
|
|
split-teaching refusal is gone (phase 70)."""
|
|
_doc(db, "Alpha", "deep/nested/doc.md", "The Doc", "FULL-TEXT")
|
|
db.commit()
|
|
|
|
holder, llm = _run_call(db, "read", {"path": "Alpha"})
|
|
|
|
assert (
|
|
llm.requests[1][0][3]["content"] == "No document at 'Alpha' — check the ls output."
|
|
)
|
|
assert holder.tool_calls == 0 and holder.read_docs == []
|
|
|
|
|
|
def test_read_unknown_combined_path_refused_through_run_agent(kb, db) -> None:
|
|
"""A combined identity that matches nothing gets the no-document
|
|
refusal (the argument echoed as passed — the model sees its own
|
|
form)."""
|
|
_doc(db, "Alpha", "x.md", "X", "X-CONTENT")
|
|
db.commit()
|
|
|
|
holder, llm = _run_call(db, "read", {"path": "Alpha/nope/deep.md"})
|
|
|
|
assert (
|
|
llm.requests[1][0][3]["content"]
|
|
== "No document at 'Alpha/nope/deep.md' — check the ls output."
|
|
)
|
|
assert holder.tool_calls == 0 and holder.read_docs == []
|
|
|
|
|
|
# ---------- grep (the phase-68 A5 contract under the new name) ----------
|
|
|
|
|
|
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_grep_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_call(db, "grep", {"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: grep adds no context
|
|
|
|
|
|
def test_grep_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_call(db, "grep", {"pattern": "needle", "path": "Alpha/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_grep_scoped_missing_doc_refused_through_run_agent(kb, db) -> None:
|
|
_doc(db, "Alpha", "a/one.md", "One", "nothing")
|
|
db.commit()
|
|
|
|
holder, llm = _run_call(db, "grep", {"pattern": "needle", "path": "Alpha/ghost.md"})
|
|
|
|
assert (
|
|
llm.requests[1][0][3]["content"]
|
|
== "No document at 'Alpha/ghost.md' — check the ls output."
|
|
)
|
|
assert holder.tool_calls == 0 and holder.read_docs == []
|
|
|
|
|
|
def test_grep_no_matches_through_run_agent(kb, db) -> None:
|
|
_doc(db, "Alpha", "a/one.md", "One", "nothing matching")
|
|
db.commit()
|
|
|
|
holder, llm = _run_call(db, "grep", {"pattern": "zebra"})
|
|
|
|
assert llm.requests[1][0][3]["content"] == (
|
|
"No matches for 'zebra' in the knowledge base."
|
|
)
|
|
assert holder.tool_calls == 1 # an executed grep with zero hits
|
|
assert holder.read_docs == []
|