feat(rag): agent document tools — list/read tools with env-tuned budgets, SSE tool events + "calling tool" UI

Grounded chat turns now run the agent loop (app/rag/agent.py) instead
of a bare chat_stream: while the per-turn budgets last
(BOR_AGENT_LIST_CALLS / BOR_AGENT_READ_CALLS, default 1 each) the model
gets list_documents (the indexed catalog, /api/docs order) and
read_document (full text, never truncated — A7-revised contract); once
both budgets are spent the tools key is dropped from the request and
the model must answer. Rejected calls (unknown tool, unknown/missing
path, document already in context, spent budget) consume no budget.
Budgets 0/0 make exactly one tools=None request — byte-identical to
the pre-phase path (budgets-as-kill-switch). Deflected turns keep the
direct chat_stream (A8 unchanged; the LOW prompt never carries the
<tools> section).

SSE contract gains {"type":"tool","name":...,"argument":
"source/path"|null} frames ahead of the answer deltas (PLAN §4
extension, owner permission 2026-08-26); done.sources, query_log.sources
and the per-turn log line (gains tool_calls=N) report the retrieval
docs + read docs, deduped. The UI shows a "calling tool"
button/label state and one visible .tool-call line per call above the
answer; the lines persist with the chat record and re-render on
reload. chat_stream passes tools through and accumulates streaming
tool_calls deltas into ToolCallPiece (tools=None stays byte-identical).

E2E: deterministic mock tool flow ("use your tools" + <tools> marker:
list -> read first catalog line -> quoted answer) plus the story suite
(marker flow, reload re-render, plain/deflected no-tool regressions).
Docs: .env.example + README (the two tools, the budgets, the SSE tool
frame, the "calling tool" UI state).

probe: turbo tool_calls=supported 2026-08-26 (uv run python -m
scripts.llm_probe --tools — non-streaming + streaming
finish_reason=tool_calls, indexed delta.tool_calls partials)
This commit is contained in:
2026-08-26 22:39:14 -04:00
parent 9efffcb428
commit 15c1272828
30 changed files with 3594 additions and 67 deletions
+83
View File
@@ -0,0 +1,83 @@
"""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.
Requires: podman compose up -d db
"""
from __future__ import annotations
import uuid
from collections.abc import Iterator
import pytest
from sqlalchemy import text
from sqlalchemy.orm import Session
from app.models import Document
from app.rag import agent
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
+13 -2
View File
@@ -192,10 +192,21 @@ def _find_emoji(text: str) -> list[str]:
)
def test_ui_chrome_has_no_emoji(client, path: str) -> None:
"""Permanent regression guard (phase 08): the UI chrome — all pages,
the JS that renders it, and the stylesheet — is emoji-free."""
the JS that renders it, and the stylesheet — is emoji-free.
Phase 37 revision (owner permission 2026-08-26, PLAN §4): the agent's
``.tool-call`` line carries two CONTENT marks — 🔎 (list) and 📄
(read) — the only emoji in the whole frontend, and only as the exact
tool-line template strings in app.js. The guard strips precisely
those two literals; any other emoji, or those marks anywhere else,
still fails."""
r = client.get(path)
assert r.status_code == 200
assert _find_emoji(r.text) == [], f"emoji found in {path}: {_find_emoji(r.text)!r}"
text = r.text
if path == "/assets/app.js":
text = text.replace('"🔎 Listing documents"', "")
text = text.replace('"📄 Reading "', "")
assert _find_emoji(text) == [], f"emoji found in {path}: {_find_emoji(text)!r}"
def test_chat_requires_message(client) -> None:
+256 -3
View File
@@ -13,6 +13,7 @@ from __future__ import annotations
import asyncio
import hashlib
import json
import logging
import math
import re
from collections.abc import Iterator
@@ -27,8 +28,10 @@ from app.api import chat as chat_api
from app.config import Settings, get_settings
from app.main import app as fastapi_app
from app.models import Chunk, QueryLog
from app.rag import agent
from app.rag.agent import AGENT_TOOLS
from app.rag.importer import import_sources
from app.rag.llm import EmbeddingError, LLMError, StreamPiece
from app.rag.llm import EmbeddingError, LLMError, StreamPiece, ToolCallPiece
FIXTURES = Path(__file__).resolve().parents[1] / "fixtures" / "docs"
QUESTION = "How is my Kubernetes cluster set up?"
@@ -57,6 +60,7 @@ class FakeRagLLM:
embed_error: Exception | None = None,
stream_error: Exception | None = None,
fail_mid_stream: bool = False,
tool_script: list[list[StreamPiece | ToolCallPiece]] | None = None,
) -> None:
self.settings = Settings(_env_file=None) # pyright: ignore[reportCallIssue]
self.embed_batches = 0
@@ -67,6 +71,18 @@ class FakeRagLLM:
self.fail_mid_stream = fail_mid_stream
self.question_embeds: list[str] = []
self.seen_messages: list[list[dict[str, str]]] = []
#: Every request's ``tools`` value (phase 37) — ``None`` is the
#: pre-phase request shape (the key is absent from the payload).
self.seen_tools: list[list[dict[str, Any]] | None] = []
#: Canned per-agent-round piece lists (phase 37): ``tool_script[i]``
#: is yielded for the *i*-th request that carries a non-None
#: ``tools`` parameter (a request the agent loop is offering tools
#: on). A request without tools — the deflected direct path, the
#: post-budget answer request, or the 0/0 single-request path —
#: always yields the thinking + answer stream below, so a
#: deflected turn through this fake is byte-identical to the
#: plain fake's output.
self.tool_script: list[list[StreamPiece | ToolCallPiece]] = list(tool_script or [])
async def embed(self, texts: list[str]) -> list[list[float]]:
self.embed_batches += 1
@@ -87,14 +103,25 @@ class FakeRagLLM:
self.question_embeds.append(text)
return _token_vec(text)
async def chat_stream(self, messages: list[dict[str, str]]):
async def chat_stream(
self,
messages: list[dict[str, str]],
tools: list[dict[str, Any]] | None = None,
):
"""Typed stream (phase 17): ``thinking`` slices (same 12-char
cadence as content) **before** the content pieces. With the
default ``thinking=""`` this yields content-only pieces — today's
behavior, new yield type."""
behavior, new yield type. Phase 37: *tools* is the agent loop's
``tools=…`` passthrough (recorded in ``seen_tools``); a request
with tools consumes the next ``tool_script`` entry, if any."""
self.seen_messages.append(messages)
self.seen_tools.append(tools)
if self.stream_error is not None:
raise self.stream_error
if tools is not None and self.tool_script:
for piece in self.tool_script.pop(0):
yield piece
return
if self.fail_mid_stream:
yield StreamPiece("content", "partial ")
raise LLMError("mid-stream dropout")
@@ -464,3 +491,229 @@ def test_chat_query_log_failure_still_sends_done(client, db, seeded_kb: FakeRagL
assert [f["type"] for f in frames if f["type"] == "delta"]
assert frames[-1]["type"] == "done"
assert frames[-1]["deflected"] is False
# ---------- phase 37: agent document tools on grounded turns ----------
def test_grounded_turn_streams_tool_frames_and_cites_read_doc(
client, db, seeded_kb: FakeRagLLM, caplog: pytest.LogCaptureFixture
) -> None:
"""(a) Grounded turn with tool calls: the event sequence is
``thinking?/tool/tool/delta…/done``; ``done.sources`` and the
``query_log`` row include the read document (deduped, order
preserved); the per-turn log line carries ``tool_calls=2``.
The agent loop offers tools while budgets last and drops them
(``tools=None``) once both are spent."""
scripted = FakeRagLLM(
tool_script=[
[
StreamPiece("thinking", "Let me list what is indexed…"),
ToolCallPiece(id="call_1", name="list_documents", arguments={}),
],
[
ToolCallPiece(
id="call_2",
name="read_document",
arguments={"source": "docs", "path": "homelab/backups.md"},
)
],
# the post-budget answer request (tools=None) falls back to the
# fake's thinking + answer stream
]
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: scripted
try:
caplog.set_level(logging.INFO, logger="app.chat")
_, _, frames = _stream_chat(client, QUESTION)
finally:
fastapi_app.dependency_overrides.clear()
types = [f["type"] for f in frames]
assert types[0] == "thinking"
assert types[1] == "tool" and types[2] == "tool" # the two executed calls
assert "error" not in types
assert types[3:-1] == ["delta"] * (len(types) - 4) # deltas, then done last
assert frames[-1]["type"] == "done"
list_frame, read_frame = frames[1], frames[2]
assert set(list_frame) == {"type", "name", "argument"}
assert list_frame["name"] == "list_documents"
assert list_frame["argument"] is None # the tool takes no parameters
assert set(read_frame) == {"type", "name", "argument"}
assert read_frame["name"] == "read_document"
assert read_frame["argument"] == "docs/homelab/backups.md"
deltas = [f for f in frames if f["type"] == "delta"]
assert len(deltas) >= 2 # genuinely streamed
assert "".join(d["text"] for d in deltas) == scripted.answer
done = frames[-1]
assert done["deflected"] is False
# done.sources = the retrieval docs + the read doc, deduped, order kept.
sources = [(s["source"], s["path"]) for s in done["sources"]]
assert sources[-1] == ("docs", "homelab/backups.md") # the read doc is cited
assert ("docs", "homelab/kubernetes.md") in sources # …after the retrieval docs
assert len(sources) == len(set(sources)) # deduped by (source, path)
assert done["sources"][-1]["title"] == "Backup Strategy"
# The agent loop offered the tools while any budget remained and
# dropped them once both were spent (single post-budget request).
assert len(scripted.seen_messages) == 3
assert scripted.seen_tools[0] == AGENT_TOOLS
assert scripted.seen_tools[1] == AGENT_TOOLS # the read budget was still open
assert scripted.seen_tools[2] is None
# The query_log row carries the same combined source list.
(row,) = db.scalars(select(QueryLog)).all()
assert row.deflected is False
assert "docs/homelab/kubernetes.md" in row.sources
assert row.sources.endswith(", docs/homelab/backups.md") # the read doc, last
# The required per-turn log line (PLAN §9 extension) counts both calls
# and lists the combined sources (retrieval + read).
lines = [r.getMessage() for r in caplog.records if "question=" in r.getMessage()]
assert lines and "tool_calls=2" in lines[-1]
assert "'docs/homelab/kubernetes.md'" in lines[-1]
assert "'docs/homelab/backups.md'" in lines[-1]
def test_deflected_turn_stays_byte_identical_without_tools(
client, db, seeded_kb: FakeRagLLM
) -> None:
"""(b) Deflected turn: the agent loop never runs — no ``tool``
frames, and the frame sequence is byte-identical to the plain fake's
direct-``chat_stream`` output even for a fake scripted to call tools
(its script is never consumed). The LLM was called once, without a
``tools`` key."""
scripted = FakeRagLLM(
tool_script=[
[ToolCallPiece(id="call_1", name="list_documents", arguments={})],
[
ToolCallPiece(
id="call_2",
name="read_document",
arguments={"source": "docs", "path": "homelab/backups.md"},
)
],
[StreamPiece("content", "never used — the agent never runs")],
]
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: seeded_kb
try:
_, _, baseline = _stream_chat(client, OFF_TOPIC)
finally:
fastapi_app.dependency_overrides.clear()
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: scripted
try:
_, _, frames = _stream_chat(client, OFF_TOPIC)
finally:
fastapi_app.dependency_overrides.clear()
assert frames == baseline # byte-identical to the direct path
assert not any(f["type"] == "tool" for f in frames)
assert frames[-1]["type"] == "done" and frames[-1]["deflected"] is True
assert len(scripted.tool_script) == 3 # the script was never consumed
assert len(scripted.seen_messages) == 1
assert scripted.seen_tools == [None] # one request, no tools key
# The read document never sneaks into the deflected turn's record.
(row,) = [
r
for r in db.scalars(select(QueryLog)).all()
if r.question == OFF_TOPIC
][-1:]
assert row.deflected is True
assert "backups.md" not in row.sources
def test_zero_agent_budgets_reproduce_pre_phase_single_request(
client,
db,
seeded_kb: FakeRagLLM,
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None:
"""(c) ``BOR_AGENT_LIST_CALLS=0 BOR_AGENT_READ_CALLS=0``: no ``tool``
frames, exactly one request **without** a ``tools`` key (the
pre-phase request shape), ``done.sources`` unchanged, and
``tool_calls=0`` in the log line — budgets-as-kill-switch."""
scripted = FakeRagLLM(
tool_script=[
[ToolCallPiece(id="call_1", name="list_documents", arguments={})],
[
ToolCallPiece(
id="call_2",
name="read_document",
arguments={"source": "docs", "path": "homelab/backups.md"},
)
],
]
)
live = get_settings()
monkeypatch.setattr(
chat_api,
"get_settings",
lambda: Settings(
_env_file=None, # pyright: ignore[reportCallIssue]
relevance_threshold=live.relevance_threshold,
agent_list_calls=0,
agent_read_calls=0,
),
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: scripted
try:
caplog.set_level(logging.INFO, logger="app.chat")
_, _, frames = _stream_chat(client, QUESTION)
finally:
fastapi_app.dependency_overrides.clear()
assert not any(f["type"] == "tool" for f in frames)
assert "error" not in [f["type"] for f in frames]
done = frames[-1]
assert done["type"] == "done"
assert done["deflected"] is False
paths = [s["path"] for s in done["sources"]]
assert "homelab/kubernetes.md" in paths # retrieval docs, unchanged
assert "homelab/backups.md" not in paths # nothing was read
# Exactly one request, and it carried no ``tools`` key at all — the
# scripted tool calls were never even offered a chance.
assert len(scripted.seen_messages) == 1
assert scripted.seen_tools == [None]
assert len(scripted.tool_script) == 2 # never consumed
(row,) = db.scalars(select(QueryLog)).all()
assert "docs/homelab/kubernetes.md" in row.sources
assert "backups.md" not in row.sources
lines = [r.getMessage() for r in caplog.records if "question=" in r.getMessage()]
assert lines and "tool_calls=0" in lines[-1]
def test_tool_execution_db_failure_yields_error_event(
client, db, seeded_kb: FakeRagLLM, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A tool call that hits a dead DB mid-stream gets the same structured
``error`` event as the pre-stream retrieval path — never a severed
stream (the "never stale" contract, PLAN §7.4)."""
scripted = FakeRagLLM(
tool_script=[[ToolCallPiece(id="call_1", name="list_documents", arguments={})]]
)
def boom(*_a: Any, **_k: Any) -> Any:
raise RuntimeError("db exploded mid tool call")
monkeypatch.setattr(agent, "list_catalog", boom)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: scripted
try:
_, _, frames = _stream_chat(client, QUESTION)
finally:
fastapi_app.dependency_overrides.clear()
# The ``tool`` frame went out first (the model requested the call);
# the failed execution ends the turn with the structured error event.
assert [f["type"] for f in frames] == ["tool", "error"]
assert frames[0]["name"] == "list_documents"
assert "offline mid-question" in frames[1]["detail"]
assert db.scalars(select(QueryLog)).all() == [] # no row for a failed turn