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
+93 -5
View File
@@ -46,6 +46,28 @@ prompts stay byte-identical to the pre-phase text (phase 15
convention); ``TurnPlan.kb_chars`` records the length of the stored
outline (0 when absent) and the per-turn log line records
``kb_chars=N`` after ``tuning=N`` (PLAN §9 line extension).
Agent document tools (phase 37, PLAN §4 extension, owner permission
2026-08-26): a **grounded** turn (``not plan.deflected``) no longer
streams a bare ``chat_stream`` — it runs the agent loop
(``app.rag.agent.run_agent``), which offers the model the two
server-side tools ``list_documents`` / ``read_document`` while the
per-turn budgets (``BOR_AGENT_LIST_CALLS`` / ``BOR_AGENT_READ_CALLS``,
default 1 each) last; once both budgets are spent the ``tools`` key is
dropped from the request and the model must answer. Each
model-requested call streams as an SSE ``tool`` event —
``{"type": "tool", "name": …, "argument": "source/path" | null}`` —
ahead of the answer's ``delta`` frames. ``done.sources``,
``query_log.sources`` and the per-turn log line all report the same
combined source list (retrieval docs + the agent's read docs, deduped
by ``(source, path)``, order preserved), and the log line records
``tool_calls=N`` after ``thinking_chars=N`` (PLAN §9 line extension —
``N`` counts budget-consuming executions; rejected calls do not
count). **Deflected turns keep the direct ``chat_stream`` — byte-
identical to the pre-phase path (A8):** the LOW prompt never carries
tools, and with **both budgets at 0** ``run_agent`` makes exactly one
``tools=None`` request, reproducing the pre-phase behavior (budgets-
as-kill-switch).
"""
from __future__ import annotations
@@ -64,11 +86,13 @@ from app.api.steering import load_steering_notes
from app.config import Settings, get_settings
from app.db import db_available, get_db
from app.models import Document, QueryLog
from app.rag.agent import AgentHolder, run_agent
from app.rag.llm import (
EmbeddingError,
LLMClient,
LLMError,
StreamPiece, # noqa: F401 (phase 17 typing: chat_stream yields StreamPiece)
StreamPiece, # type of the answer pieces streamed by the agent loop
ToolCallPiece, # phase 37: one model-requested tool call
)
from app.rag.overview import load_kb_overview
from app.rag.prompts import build_deflect_prompt, build_high_prompt
@@ -79,6 +103,7 @@ from app.schemas import (
ChatErrorEvent,
ChatRequest,
ChatThinkingEvent,
ChatToolEvent,
SourceRef,
)
@@ -260,7 +285,6 @@ async def chat(
).model_dump()
)
return
source_paths = [f"{d.source}/{d.path}" for d in plan.docs]
messages = [
{"role": "system", "content": plan.system_prompt},
{"role": "user", "content": request.message},
@@ -271,9 +295,44 @@ async def chat(
# ahead of the ``delta`` events (PLAN §4 extension); the
# kill-switch (``BOR_STREAM_THINKING=0``) suppresses the
# frames, not the counting.
# Phase 37: a grounded turn runs the agent loop instead of a
# bare ``chat_stream`` — its ``ToolCallPiece``s stream as
# ``tool`` events ahead of the answer. A deflected turn keeps
# the direct ``chat_stream`` (byte-identical, A8): the LOW
# prompt never carries tools, and with both budgets at 0
# ``run_agent`` is a single ``tools=None`` request anyway.
holder = AgentHolder()
answer_stream: AsyncIterator[StreamPiece | ToolCallPiece]
if plan.deflected:
answer_stream = llm.chat_stream(messages)
else:
answer_stream = run_agent(
llm,
db,
system_prompt=plan.system_prompt,
user_message=request.message,
seed_docs=plan.docs,
settings=settings,
holder=holder,
)
thinking_chars = 0
try:
async for piece in llm.chat_stream(messages): # StreamPiece (phase 17)
async for piece in answer_stream: # StreamPiece | ToolCallPiece
if isinstance(piece, ToolCallPiece):
# Phase 37 (PLAN §4 extension): one SSE ``tool``
# frame per model-requested call; ``argument`` is the
# read_document "source/path" (null otherwise).
yield sse_event(
ChatToolEvent(
name=piece.name,
argument=(
f"{piece.arguments.get('source')}/{piece.arguments.get('path')}"
if piece.name == "read_document"
else None
),
).model_dump()
)
continue
if piece.kind == "thinking":
thinking_chars += len(piece.text)
if settings.stream_thinking:
@@ -293,8 +352,35 @@ async def chat(
).model_dump()
)
return
except Exception: # noqa: BLE001 — a tool call hit the DB mid-stream
# Phase 37: tool execution (list_catalog / find_document) runs
# inside the stream now; a mid-turn DB failure gets the same
# structured ``error`` event as the pre-stream retrieval path.
logger.exception(
"chat: tool execution failed question=%r total_ms=%d",
request.message,
int((time.monotonic() - started) * 1000),
)
yield sse_event(
ChatErrorEvent(
detail="The knowledge base went offline mid-question — is Postgres up?"
).model_dump()
)
return
# 4. Durable record + required per-turn log line (PLAN §9).
# Phase 37: the agent's read documents join the retrieval's —
# deduped by (source, path), order preserved — and the same
# combined list feeds done.sources, query_log.sources and the
# log line (empty on deflected turns: the agent never runs).
cited_docs: list[Document] = []
seen: set[tuple[str, str]] = set()
for doc in [*plan.docs, *holder.read_docs]:
key = (doc.source, doc.path)
if key not in seen:
seen.add(key)
cited_docs.append(doc)
source_paths = [f"{d.source}/{d.path}" for d in cited_docs]
total_ms = int((time.monotonic() - started) * 1000)
try:
db.add(
@@ -314,7 +400,8 @@ async def chat(
logger.info(
"question=%r embed_ms=%d top_score=%.3f fts_hits=%d summary_hits=%d tuning=%d "
"kb_chars=%d threshold=%.2f deflected=%s sources=%r thinking_chars=%d total_ms=%d",
"kb_chars=%d threshold=%.2f deflected=%s sources=%r thinking_chars=%d "
"tool_calls=%d total_ms=%d",
request.message,
embed_ms,
plan.top_score,
@@ -326,13 +413,14 @@ async def chat(
plan.deflected,
source_paths,
thinking_chars,
holder.tool_calls,
total_ms,
)
yield sse_event(
ChatDoneEvent(
deflected=plan.deflected,
sources=[
SourceRef(source=d.source, path=d.path, title=d.title) for d in plan.docs
SourceRef(source=d.source, path=d.path, title=d.title) for d in cited_docs
],
suggestions=plan.suggestions,
).model_dump()