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:
@@ -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
|
||||
Reference in New Issue
Block a user