307 lines
12 KiB
Python
307 lines
12 KiB
Python
"""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. The model is offered the two OpenAI functions in :data:`AGENT_TOOLS`
|
|
for the whole turn — phase 45 removed the phase-37 per-tool budgets
|
|
(owner permission 2026-08-27, ``TODO.md`` L8: "allow the LLM to make
|
|
as many tool calls as it wants"): ``list_documents`` and
|
|
``read_document`` can each be called as many times as the model needs,
|
|
re-lists included. With ``settings.agent_max_rounds``
|
|
(``BOR_AGENT_MAX_ROUNDS``, default 10) at 0 the loop makes exactly one
|
|
request with ``tools=None`` — byte-identical to the pre-phase-37 chat
|
|
path (the kill switch).
|
|
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 get a one-line refusal and count in nothing
|
|
(``holder.tool_calls`` tracks executed calls only): 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 …"``.
|
|
A rejected call still consumes a *round* in the loop, so a
|
|
pathological stream that keeps emitting rejected calls is bounded by
|
|
the cap (point 4).
|
|
4. Every call the model emits is appended back to the message history as
|
|
the assistant tool-call message + the tool result (refusals included),
|
|
consumes one round, and the model is called again. At the round cap —
|
|
``max_rounds = settings.agent_max_rounds`` (``BOR_AGENT_MAX_ROUNDS``,
|
|
default 10) — the loop forces one final ``chat_stream(messages,
|
|
tools=None)`` and returns: the cap is the **only** forced exit
|
|
(besides "the stream carried no calls"), and it bounds pathological
|
|
rejected-call streams.
|
|
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 executed tool calls (re-lists included); 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`` for the whole grounded turn —
|
|
#: phase 45 removed the per-tool budgets; the round cap
|
|
#: (``BOR_AGENT_MAX_ROUNDS``) is the only bound.
|
|
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 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 count in nothing
|
|
#: (``holder.tool_calls`` tracks executed calls); the round cap bounds
|
|
#: their pathological repetition (phase 45).
|
|
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 calls executed (re-lists included);
|
|
rejected calls (unknown tool, unknown/missing arguments or document,
|
|
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,
|
|
) -> str:
|
|
"""Execute one tool call server-side (DB only).
|
|
|
|
Returns the tool result text. A successful call bumps
|
|
``holder.tool_calls`` (a successful read also appends the
|
|
:class:`Document` to ``holder.read_docs``); rejected calls return
|
|
their refusal line and count in nothing.
|
|
"""
|
|
if call.name == "list_documents":
|
|
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
|
|
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
|
|
known = {(doc.source, doc.path) for doc in (*seed_docs, *holder.read_docs)}
|
|
if (source, path) in known:
|
|
return ALREADY_IN_CONTEXT
|
|
doc = find_document(db, source, path)
|
|
if doc is None:
|
|
return (
|
|
f"No document at {source}/{path} — check the list_documents output."
|
|
)
|
|
holder.read_docs.append(doc)
|
|
holder.tool_calls += 1
|
|
return f"Document {source}/{path}:\n{doc.content}"
|
|
return UNKNOWN_TOOL
|
|
|
|
|
|
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 executed tool-call count (re-lists included).
|
|
|
|
``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." — the rejection counts
|
|
in nothing, but it still consumes a round.
|
|
"""
|
|
messages: list[dict[str, Any]] = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_message},
|
|
]
|
|
# Phase 45: no per-tool budgets — the tools stay offered for the
|
|
# whole turn, bounded by the round cap. ``0`` is the no-tools kill
|
|
# switch: exactly one request with ``tools=None`` (the pre-phase-37
|
|
# path).
|
|
max_rounds = settings.agent_max_rounds
|
|
tools: list[dict[str, Any]] | None = AGENT_TOOLS if max_rounds > 0 else None
|
|
rounds = 0
|
|
while True:
|
|
calls: list[ToolCallPiece] = []
|
|
# Phase 48: bind the round's stream so a consumer abandon
|
|
# (GeneratorExit into the yield below) tears down the in-flight
|
|
# model stream deterministically — not GC-dependent. Awaiting
|
|
# ``aclose()`` in the ``finally`` is safe because it does not
|
|
# yield; on a fully consumed round it is a quiet no-op.
|
|
stream = llm.chat_stream(cast("list[dict[str, str]]", messages), tools=tools)
|
|
try:
|
|
async for piece in stream:
|
|
if isinstance(piece, ToolCallPiece):
|
|
calls.append(piece)
|
|
yield piece
|
|
finally:
|
|
await stream.aclose()
|
|
if not calls:
|
|
return # the answer was streamed
|
|
call = calls[0] # a stream can carry several calls; run the first
|
|
result = _execute_tool(db, call, seed_docs, holder)
|
|
rounds += 1 # every call the model emits consumes a round
|
|
logger.info(
|
|
"agent tool=%s args=%s round=%d/%d",
|
|
call.name,
|
|
json.dumps(call.arguments, ensure_ascii=False)[:200],
|
|
rounds,
|
|
max_rounds,
|
|
)
|
|
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})
|
|
if rounds >= max_rounds:
|
|
logger.warning(
|
|
"agent round cap reached (rounds=%d) — forcing a final "
|
|
"no-tools answer",
|
|
rounds,
|
|
)
|
|
# Phase 48: the forced final answer gets the same explicit
|
|
# teardown as the loop rounds (consumer abandon mid-final
|
|
# answer must still close the model's stream).
|
|
final = llm.chat_stream(cast("list[dict[str, str]]", messages), tools=None)
|
|
try:
|
|
async for piece in final:
|
|
yield piece
|
|
finally:
|
|
await final.aclose()
|
|
return
|