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:
+93
-5
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user