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
+300
View File
@@ -0,0 +1,300 @@
"""Agent loop: the grounded-turn document tools (phase 37, task 03).
Probe verdict (task 01 — ``uv run python -m scripts.llm_probe --tools``
run live against aipi): **``probe: turbo tool_calls=supported 2026-08-26``**
— ``turbo`` answers OpenAI ``tools`` requests with
``finish_reason="tool_calls"`` and streams the calls as indexed
``delta.tool_calls`` partials (id + name on the first partial, arguments
in fragments). This module therefore uses the **native tool-calling
path**: tool calls arrive as :class:`app.rag.llm.ToolCallPiece` values
from ``chat_stream(messages, tools=AGENT_TOOLS)``. The prompt-based
JSON-block fallback (documented in the task file) is *not* implemented —
it exists only for a "not supported"/"intermittent" verdict, and the
probe came back "supported".
Loop contract (one grounded chat turn; the API layer wires this in,
task 04):
1. While budget remains the model is offered the two OpenAI functions in
:data:`AGENT_TOOLS`: up to ``settings.agent_list_calls``
(``BOR_AGENT_LIST_CALLS``, default 1) ``list_documents`` calls and up
to ``settings.agent_read_calls`` (``BOR_AGENT_READ_CALLS``, default 1)
``read_document`` calls. With both budgets at 0 the loop makes exactly
one request with ``tools=None`` — byte-identical to the pre-phase chat
path (budgets-as-kill-switch, phase 37 locked decision).
2. Each tool call the model emits is executed server-side against
Postgres only (no LLM, no network): ``list_documents`` returns the
indexed catalog — one ``source/path — title`` line per document,
``GET /api/docs`` order (uncapped in v1; the UI never shows it, only
the model does) — and ``read_document`` returns the document's **full**
content (A7-revised contract: never truncated).
3. Rejected calls consume **no** budget and get a one-line refusal:
unknown tool name → ``"Unknown tool."``; missing ``source``/``path``
arguments; a document already in context (seed or previously read) →
``"Already in your context."``; an unknown ``source/path`` →
``"No document at …"``; an exhausted list/read budget → the matching
``"No … budget left"`` refusal.
4. Every executed call is appended back to the message history as the
assistant tool-call message + the tool result, and the model is called
again. Once **both** budgets are spent, ``tools`` is dropped from the
request and the model must answer. Belt-and-braces round cap:
``max_rounds = 2 + agent_list_calls + agent_read_calls`` (every tool
round consumes a budget, so the cap only catches pathological streams
that keep calling rejected tools) — at the cap the loop forces one
final ``chat_stream(messages, tools=None)`` and returns.
5. A rare stream that carries both content and a tool call keeps the
content (it was already emitted) **and** still runs the tool.
6. *holder* (an :class:`AgentHolder`) records the read documents and the
number of budget-consuming tool executions; the API layer (task 04)
reads it after the stream to extend ``done.sources`` /
``query_log.sources`` and the per-turn log line (``tool_calls=N``).
The DB accessors (:func:`list_catalog`, :func:`find_document`) are
module-level functions so unit tests can monkeypatch them without a
database.
"""
from __future__ import annotations
import json
import logging
from collections.abc import AsyncIterator, Sequence
from dataclasses import dataclass, field
from typing import Any, cast
from sqlalchemy import select
from sqlalchemy.orm import Session
from app.config import Settings
from app.models import Document
from app.rag.llm import LLMClient, StreamPiece, ToolCallPiece
logger = logging.getLogger("app.agent")
#: The two agent tools (phase 37): OpenAI function definitions passed as
#: ``tools=AGENT_TOOLS`` to ``chat_stream`` while the per-turn budgets
#: (``BOR_AGENT_LIST_CALLS`` / ``BOR_AGENT_READ_CALLS``) remain.
AGENT_TOOLS: list[dict[str, Any]] = [
{
"type": "function",
"function": {
"name": "list_documents",
"description": (
"List every document indexed in the knowledge base, one "
"`source/path — title` line each"
),
"parameters": {"type": "object", "properties": {}, "required": []},
},
},
{
"type": "function",
"function": {
"name": "read_document",
"description": (
"Add the full content of exactly one more indexed document "
"to your context"
),
"parameters": {
"type": "object",
"properties": {
"source": {
"type": "string",
"description": (
"The document's source (a directory basename, "
"e.g. 'Homelab')."
),
},
"path": {
"type": "string",
"description": (
"The document's path relative to its source "
"directory."
),
},
},
"required": ["source", "path"],
},
},
},
]
#: Tool refusal texts (phase 37): rejected calls consume no budget.
LIST_EXHAUSTED = "No listing budget left — answer with what you have."
READ_EXHAUSTED = "No reading budget left — answer with what you have."
ALREADY_IN_CONTEXT = "Already in your context."
UNKNOWN_TOOL = "Unknown tool."
MISSING_READ_ARGS = "read_document requires string arguments 'source' and 'path'."
def list_catalog(db: Session) -> list[tuple[str, str, str]]:
"""Every indexed document as ``(source, path, title)``.
Ordered by ``(source, path)`` — the same order as ``GET /api/docs``.
Module-level (not a method) so unit tests can monkeypatch it.
"""
rows = db.execute(
select(Document.source, Document.path, Document.title).order_by(
Document.source, Document.path
)
).all()
return [(source, path, title) for source, path, title in rows]
def find_document(db: Session, source: str, path: str) -> Document | None:
"""The indexed document at ``(source, path)``, or ``None``.
Module-level (not a method) so unit tests can monkeypatch it.
"""
return db.scalar(
select(Document).where(Document.source == source, Document.path == path)
)
@dataclass
class AgentHolder:
"""Per-turn agent state the API layer reads after the stream (task 04).
``read_docs``: the documents ``read_document`` added to the context,
in read order (deduped — re-reading a document appends nothing).
``tool_calls``: how many tool executions consumed budget; rejected
calls (unknown tool, unknown/missing document, exhausted budget,
already-in-context) do not count. Drives the per-turn log line's
``tool_calls=N`` field (task 04).
"""
read_docs: list[Document] = field(default_factory=list)
tool_calls: int = 0
def _execute_tool(
db: Session,
call: ToolCallPiece,
seed_docs: Sequence[Document],
holder: AgentHolder,
list_left: int,
read_left: int,
) -> tuple[str, int, int]:
"""Execute one tool call server-side (DB only).
Returns ``(result, list_left, read_left)``. Rejected calls consume no
budget; a successful read appends the :class:`Document` to
``holder.read_docs`` and bumps ``holder.tool_calls``.
"""
if call.name == "list_documents":
if list_left <= 0:
return LIST_EXHAUSTED, list_left, read_left
rows = list_catalog(db)
listing = f"{len(rows)} documents:\n" + "\n".join(
f"{source}/{path} — {title}" for source, path, title in rows
)
holder.tool_calls += 1
return listing, list_left - 1, read_left
if call.name == "read_document":
raw_source = call.arguments.get("source")
raw_path = call.arguments.get("path")
source = raw_source.strip() if isinstance(raw_source, str) else ""
path = raw_path.strip() if isinstance(raw_path, str) else ""
if not source or not path:
return MISSING_READ_ARGS, list_left, read_left
known = {(doc.source, doc.path) for doc in (*seed_docs, *holder.read_docs)}
if (source, path) in known:
return ALREADY_IN_CONTEXT, list_left, read_left
if read_left <= 0:
return READ_EXHAUSTED, list_left, read_left
doc = find_document(db, source, path)
if doc is None:
return (
f"No document at {source}/{path} — check the list_documents output.",
list_left,
read_left,
)
holder.read_docs.append(doc)
holder.tool_calls += 1
return f"Document {source}/{path}:\n{doc.content}", list_left, read_left - 1
return UNKNOWN_TOOL, list_left, read_left
async def run_agent(
llm: LLMClient,
db: Session,
*,
system_prompt: str,
user_message: str,
seed_docs: Sequence[Document],
settings: Settings,
holder: AgentHolder,
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
"""Run the grounded-turn tool loop, yielding every stream piece.
Every piece (``thinking`` / ``content`` / tool calls) is yielded as it
arrives; the API layer (task 04) turns tool-call pieces into SSE
``tool`` events. After the loop finishes, *holder* carries the read
documents and the budget-consuming tool count.
``seed_docs`` are the documents the retrieval already put in context
(they shape the *system_prompt* the caller built); re-reading one of
them is rejected as "Already in your context." without spending budget.
"""
messages: list[dict[str, Any]] = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_message},
]
list_left = settings.agent_list_calls
read_left = settings.agent_read_calls
tools: list[dict[str, Any]] | None = AGENT_TOOLS if (list_left or read_left) else None
# Every tool round consumes a budget, so this cap only catches
# pathological streams that keep calling rejected tools (belt and
# braces — the budgets already force the answer after
# list + read rounds).
max_rounds = 2 + settings.agent_list_calls + settings.agent_read_calls
rounds = 0
while True:
calls: list[ToolCallPiece] = []
async for piece in llm.chat_stream(
cast("list[dict[str, str]]", messages), tools=tools
):
if isinstance(piece, ToolCallPiece):
calls.append(piece)
yield piece
if not calls:
return # the answer was streamed
call = calls[0] # a stream can carry several calls; run the first
result, list_left, read_left = _execute_tool(
db, call, seed_docs, holder, list_left, read_left
)
logger.info(
"agent tool=%s args=%s budget list_left=%d read_left=%d",
call.name,
json.dumps(call.arguments, ensure_ascii=False)[:200],
list_left,
read_left,
)
messages.append(
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": call.id,
"type": "function",
"function": {
"name": call.name,
"arguments": json.dumps(call.arguments),
},
}
],
}
)
messages.append({"role": "tool", "tool_call_id": call.id, "content": result})
tools = None if (list_left == 0 and read_left == 0) else AGENT_TOOLS
rounds += 1
if rounds >= max_rounds:
logger.warning(
"agent round cap reached (rounds=%d) — forcing a final "
"no-tools answer",
rounds,
)
async for piece in llm.chat_stream(
cast("list[dict[str, str]]", messages), tools=None
):
yield piece
return
+139 -15
View File
@@ -4,9 +4,10 @@ Provides the embeddings surface (importer, retrieval), one-shot chat
completions (phase 30: the ``lite`` model summarizes non-markdown
documents at import time), and chat streaming (PLAN A15) for the RAG
pipeline. Chat streaming yields typed :class:`StreamPiece` values
(phase 17): aipi's ``turbo`` model streams its reasoning as
``delta.reasoning_content`` chunks (deepseek/litellm wire convention,
verified live 2026-08-23) **before** the answer's
(phase 17) and — when the caller passes a ``tools`` list —
:class:`ToolCallPiece` values (phase 37): aipi's ``turbo`` model streams
its reasoning as ``delta.reasoning_content`` chunks (deepseek/litellm
wire convention, verified live 2026-08-23) **before** the answer's
``delta.content`` chunks, and reasoning counts against ``max_tokens``
(an answer can in principle be empty).
@@ -17,10 +18,11 @@ vectors that pgvector rejects.
"""
from __future__ import annotations
import json
import logging
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import Literal, cast
from typing import Any, Literal, cast
from openai import AsyncOpenAI
from openai.types.chat import ChatCompletionMessageParam
@@ -55,6 +57,78 @@ class StreamPiece:
text: str
@dataclass(frozen=True)
class ToolCallPiece:
"""One model-requested tool call accumulated from stream deltas (phase 37).
``id`` is the model's tool_call id (synthesized as ``call_<index>``
when the wire never carried one), ``name`` is the function name
(whatever the caller's ``tools`` list names — for the agent loop,
``list_documents`` / ``read_document``), and ``arguments`` is the
parsed JSON object (``{}`` when the model sent none).
"""
id: str # the model's tool_call id; synthesized "call_<index>" when absent
name: str # "list_documents" | "read_document" (whatever AGENT_TOOLS names)
arguments: dict[str, Any]
@dataclass
class _ToolCallSlot:
"""Mutable accumulator for one streamed tool call (phase 37, private).
``id`` and ``function.name`` arrive on the first partial for an index;
``function.arguments`` arrives in fragments to concatenate (OpenAI wire
convention, verified live against aipi 2026-08-26).
"""
id: str | None = None
name: str = ""
arguments: str = ""
def _materialize_tool_calls(
slots: dict[int, _ToolCallSlot],
) -> list[ToolCallPiece]:
"""Turn accumulated slots into ordered :class:`ToolCallPiece` values.
Malformed ``arguments`` JSON raises :class:`LLMError` — a silently
dropped tool call would corrupt the agent loop (fail-loud house
style). Empty/``null`` arguments become ``{}`` (a no-parameter call
such as ``list_documents``).
"""
pieces: list[ToolCallPiece] = []
for index in sorted(slots):
slot = slots[index]
raw = slot.arguments.strip()
label = slot.name or f"index {index}"
if raw:
try:
parsed: Any = json.loads(raw)
except json.JSONDecodeError as e:
raise LLMError(
f"model sent malformed tool-call arguments for '{label}': "
f"{raw[:200]!r} ({e})"
) from e
else:
parsed = None
if parsed is None:
arguments: dict[str, Any] = {}
elif isinstance(parsed, dict):
arguments = cast("dict[str, Any]", parsed)
else:
raise LLMError(
f"model sent non-object tool-call arguments for '{label}': "
f"{raw[:200]!r}"
)
pieces.append(
ToolCallPiece(
id=slot.id or f"call_{index}", name=slot.name, arguments=arguments
)
)
return pieces
# aipi's local embedding model rejects requests over ~1024 input tokens
# ("input is too large to process"). Batch by estimated tokens, with a
# safety margin under that cap — code-dense text can tokenize at ~3
@@ -232,8 +306,10 @@ class LLMClient:
return content.strip()
async def chat_stream(
self, messages: list[dict[str, str]]
) -> AsyncIterator[StreamPiece]:
self,
messages: list[dict[str, str]],
tools: list[dict[str, Any]] | None = None,
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
"""Stream assistant pieces from the chat model (PLAN A5/A15, phase 17).
``stream=True`` against the OpenAI-compatible endpoint, yielding
@@ -252,24 +328,61 @@ class LLMClient:
32 768) output tokens — the old hard 700-token cap cut long
answers off mid-sentence (owner report 2026-08-22).
Tool calls (phase 37): when *tools* (an OpenAI ``tools`` list) is
not None it is passed through as ``tools=…``; when None the key is
**not** included, so the request is byte-identical to pre-phase-37
and no tool pieces can be produced. A tool-calling model replies
with ``delta.tool_calls`` partials — keyed by ``index``, with
``id`` and ``function.name`` on the first partial and
``function.arguments`` in fragments — which are accumulated into
one :class:`ToolCallPiece` per call, yielded in index order at
stream end (or immediately once a chunk carries
``finish_reason="tool_calls"``). Malformed ``arguments`` JSON
raises :class:`LLMError`. Wire convention verified live against
aipi's ``turbo`` on 2026-08-26 via
``uv run python -m scripts.llm_probe --tools`` (phase 37, task 01:
``probe: turbo tool_calls=supported 2026-08-26``).
Any failure (network, HTTP, malformed stream) surfaces as
:class:`LLMError` so the API layer can turn it into an SSE
``error`` event instead of a hung request.
"""
try:
kwargs: dict[str, Any] = {
# ``{role, content}`` dicts are exactly what the message params
# accept; the cast keeps pyright honest about the SDK's union.
stream = await self._client.chat.completions.create(
model=self.settings.llm_chat_model,
messages=cast("list[ChatCompletionMessageParam]", messages),
temperature=0.4,
max_tokens=self.settings.max_output_tokens,
stream=True,
)
"model": self.settings.llm_chat_model,
"messages": cast("list[ChatCompletionMessageParam]", messages),
"temperature": 0.4,
"max_tokens": self.settings.max_output_tokens,
"stream": True,
}
if tools is not None:
kwargs["tools"] = tools
try:
stream = await self._client.chat.completions.create(**kwargs)
calls: dict[int, _ToolCallSlot] = {}
emitted = False
async for chunk in stream:
if not chunk.choices:
continue
delta = chunk.choices[0].delta
choice = chunk.choices[0]
delta = choice.delta
# Tool-call partials (phase 37) accumulate across chunks,
# keyed by index; a missing index (not seen on aipi) falls
# back to the next synthetic slot.
for tc in getattr(delta, "tool_calls", None) or []:
idx = getattr(tc, "index", None)
key = idx if isinstance(idx, int) else (max(calls) + 1 if calls else 0)
slot = calls.setdefault(key, _ToolCallSlot())
tc_id = getattr(tc, "id", None)
if tc_id and slot.id is None:
slot.id = tc_id
fn = getattr(tc, "function", None)
if fn is not None:
if fn.name:
slot.name += fn.name
if fn.arguments:
slot.arguments += fn.arguments
reasoning = getattr(delta, "reasoning_content", None)
if not reasoning:
# Future-proofing: the same wire convention under a
@@ -280,6 +393,17 @@ class LLMClient:
content = delta.content
if content:
yield StreamPiece("content", content)
if (
calls
and not emitted
and getattr(choice, "finish_reason", None) == "tool_calls"
):
for piece in _materialize_tool_calls(calls):
yield piece
emitted = True
if calls and not emitted:
for piece in _materialize_tool_calls(calls):
yield piece
except LLMError:
raise
except Exception as e: # noqa: BLE001 — wrap transport-level failures
+32 -5
View File
@@ -23,6 +23,12 @@ the ``<tuning>`` section (order: ``<relevance>`` →
``<knowledge_base>`` → ``<tuning>`` → mode body) — the agent knows
roughly what the KB contains before retrieval. With an empty row the
prompt is byte-identical to the pre-phase text.
Agent tools (phase 37): the **HIGH** prompt only carries a ``<tools>``
section after the ``<documents>`` body — the grounded turn may call the
server-side ``list_documents`` / ``read_document`` tools (budgeted, see
:mod:`app.rag.agent`). The LOW/deflection prompt never carries it and
stays byte-identical to the pre-phase text.
"""
from __future__ import annotations
@@ -66,6 +72,25 @@ _KB_INTRO = (
"(generated at import time):\n"
)
#: The ``<tools>`` instructions section — **HIGH prompt only** (phase 37,
#: task 03): a grounded turn may extend its context through the two
#: server-side tools (budgets: ``BOR_AGENT_LIST_CALLS`` /
#: ``BOR_AGENT_READ_CALLS``, see :mod:`app.rag.agent`). Appended after
#: the mode body (``<documents>``), so the instructions are the last
#: thing the model reads. The LOW/deflection prompt never carries it —
#: a deflection has no grounded context to extend — and stays
#: byte-identical to the pre-phase text. The E2E mock keys off the
#: ``<tools>`` marker's *presence*, not this wording.
TOOLS_SECTION: str = (
"<tools>\n"
"If the documents in your context reference other files, or you need "
"content that is not included above, call `list_documents` to see what "
"is indexed, then `read_document` to pull in exactly one more document. "
"Answer as soon as you have what you need — do not read more than one "
"extra document.\n"
"</tools>"
)
def _base(relevance: str) -> str:
if relevance not in ("HIGH", "LOW"):
@@ -155,11 +180,13 @@ def build_high_prompt(
kb_overview: str | None = None,
) -> str:
"""Grounded turn: locked persona (+ steering, + KB overview) + full
texts of the top documents.
texts of the top documents + the ``<tools>`` instructions (phase 37).
Section order (phase 31): ``<relevance>`` → ``<knowledge_base>`` →
``<tuning>`` → ``<documents>``; empty steering/overview omit their
section, keeping the prompt byte-identical to the pre-phase text.
Section order: ``<relevance>`` → ``<knowledge_base>`` → ``<tuning>``
→ ``<documents>`` → ``<tools>``; empty steering/overview omit their
section. ``<tools>`` is always present in the HIGH prompt (the
budgets — not the prompt — decide whether the tools are actually
offered to the model, see :mod:`app.rag.agent`).
"""
blocks = [
f'<document source="{doc.source}" path="{doc.path}" title="{doc.title}">\n'
@@ -174,7 +201,7 @@ texts of the top documents.
for part in (build_kb_section(kb_overview or ""), build_steering_section(notes or [])):
if part:
prompt += "\n" + part
return prompt + "\n<documents>\n" + body + "\n</documents>"
return prompt + "\n<documents>\n" + body + "\n</documents>\n" + TOOLS_SECTION
def build_deflect_prompt(