501 lines
21 KiB
Python
501 lines
21 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 three 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``,
|
|
``read_document`` and ``search_documents`` can each be called as many
|
|
times as the model needs, re-lists and re-searches 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: X | path: Y | title: Z`` line per
|
|
document (phase 63: labeled fields — unambiguous for LLM parsing),
|
|
``GET /api/docs`` order (uncapped in v1; the UI never shows it, only
|
|
the model does) — ``read_document`` returns the document's **full**
|
|
content (A7-revised contract: never truncated) — and
|
|
``search_documents`` greps the indexed documents (or one named
|
|
document) for a case-insensitive fixed substring and returns up to 20
|
|
``source/path:line: text`` match lines (owner-locked A5, phase 68),
|
|
each line truncated to 200 chars. A search is a **locator**, not a
|
|
context-adder: it never appends to the answer context (only
|
|
``read_document`` does — ``holder.read_docs`` is untouched by a
|
|
search).
|
|
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 search
|
|
without a usable ``pattern`` (missing, blank or non-string) or with a
|
|
half-specified ``source``/``path`` target; a document already in
|
|
context (seed or previously read) → ``"Already in your
|
|
context."``; an unknown ``source/path`` (read or scoped search) →
|
|
``"No document at …"``. A search that ran but found nothing is NOT a
|
|
rejection — its ``"No matches for …"`` line is a (counted) result.
|
|
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 retried no-tools request
|
|
(``chat_stream_retried`` with ``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``).
|
|
7. Retries (phase 67, owner-locked A2): every model request — each tool
|
|
round and the forced final ``tools=None`` call — goes through
|
|
``chat_stream_retried``: a round that dies before its first piece is
|
|
restarted with the SAME messages (up to ``settings.llm_retries``
|
|
restarts, a flat ``settings.llm_retry_delay`` between attempts, each
|
|
preceded by a :class:`app.rag.llm.RetryPiece` the API layer turns into
|
|
an SSE ``retry`` frame); a round that already streamed a piece fails
|
|
the turn as before (no partial answer is ever redone). Retries are
|
|
invisible to the round cap: a round that needed a retry still consumes
|
|
exactly one round. With ``settings.llm_retries=0`` every request is a
|
|
single plain attempt (the pre-phase-67 path).
|
|
|
|
The DB accessors (:func:`list_catalog`, :func:`find_document`,
|
|
:func:`all_documents`) and the :func:`grep_document` line matcher 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,
|
|
RetryPiece,
|
|
StreamPiece,
|
|
ToolCallPiece,
|
|
chat_stream_retried,
|
|
)
|
|
|
|
logger = logging.getLogger("app.agent")
|
|
|
|
#: The three agent tools (phase 37; ``search_documents`` added in phase
|
|
#: 68): 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: X | path: Y | title: Z` 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, as shown after 'source: ' in the "
|
|
"list_documents output (e.g. 'Homelab' from "
|
|
"'source: Homelab | path: homelab/aws-route53.md')."
|
|
),
|
|
},
|
|
"path": {
|
|
"type": "string",
|
|
"description": (
|
|
"The document's path, as shown after 'path: ' in the "
|
|
"list_documents output (e.g. 'homelab/aws-route53.md' from "
|
|
"'source: Homelab | path: homelab/aws-route53.md')."
|
|
),
|
|
},
|
|
},
|
|
"required": ["source", "path"],
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "search_documents",
|
|
"description": (
|
|
"Search every indexed document for an exact string "
|
|
"(case-insensitive) and return up to 20 matching lines as "
|
|
"'source/path:line: text' — use this to locate content, "
|
|
"then read_document the winner. Optionally pass 'source' "
|
|
"and 'path' (as shown in list_documents) to search one "
|
|
"document only."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"pattern": {
|
|
"type": "string",
|
|
"description": (
|
|
"The exact text to search for (a plain "
|
|
"substring, not a regex)"
|
|
),
|
|
},
|
|
"source": {
|
|
"type": "string",
|
|
"description": (
|
|
"The document's source, as shown after 'source: ' in the "
|
|
"list_documents output (e.g. 'Homelab' from "
|
|
"'source: Homelab | path: homelab/aws-route53.md')."
|
|
),
|
|
},
|
|
"path": {
|
|
"type": "string",
|
|
"description": (
|
|
"The document's path, as shown after 'path: ' in the "
|
|
"list_documents output (e.g. 'homelab/aws-route53.md' from "
|
|
"'source: Homelab | path: homelab/aws-route53.md')."
|
|
),
|
|
},
|
|
},
|
|
"required": ["pattern"],
|
|
},
|
|
},
|
|
},
|
|
]
|
|
|
|
#: 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'."
|
|
MISSING_SEARCH_ARGS = "search_documents requires a string argument 'pattern'."
|
|
|
|
#: Search caps (owner-locked A5, phase 68): a global per-call match cap
|
|
#: (across documents, in catalog order) and a per-match-line char limit.
|
|
SEARCH_MAX_MATCHES = 20
|
|
SEARCH_LINE_LIMIT = 200
|
|
|
|
#: No-match result lines (templates — the pattern is truncated to 100
|
|
#: chars before formatting, to keep a long pattern from bloating the
|
|
#: tool result). A no-match line is a *result* of an executed search,
|
|
#: not a refusal (see the module docstring, point 3).
|
|
NO_MATCHES = "No matches for '{pattern}' in the knowledge base."
|
|
NO_MATCHES_SCOPED = "No matches for '{pattern}' in {source}/{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)
|
|
)
|
|
|
|
|
|
def all_documents(db: Session) -> list[Document]:
|
|
"""Every indexed document (full rows), ordered by ``(source, path)``
|
|
— catalog order.
|
|
|
|
The whole-KB ``search_documents`` path loads all contents in this one
|
|
bulk query (catalog order is the locked match order, owner-locked A5).
|
|
Module-level (not a method) so unit tests can monkeypatch it.
|
|
"""
|
|
return list(
|
|
db.execute(
|
|
select(Document).order_by(Document.source, Document.path)
|
|
).scalars()
|
|
)
|
|
|
|
|
|
def grep_document(content: str, pattern: str) -> list[tuple[int, str]]:
|
|
"""Every line of *content* that contains *pattern*, in file order.
|
|
|
|
Case-insensitive **fixed substring** (owner-locked A5: no regex — no
|
|
ReDoS surface, a simple contract for the model). Returns
|
|
``(1-based line number, line.rstrip())`` pairs; an empty *content*
|
|
never matches a non-empty pattern.
|
|
"""
|
|
needle = pattern.lower()
|
|
return [
|
|
(number, line.rstrip())
|
|
for number, line in enumerate(content.split("\n"), start=1)
|
|
if needle in line.lower()
|
|
]
|
|
|
|
|
|
@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``; a search never does — it
|
|
is a locator, locked A5); rejected calls return their refusal line
|
|
and count in nothing. A search that ran but found nothing is still a
|
|
successful (counted) call — its no-match line is a result, not a
|
|
refusal.
|
|
"""
|
|
if call.name == "list_documents":
|
|
rows = list_catalog(db)
|
|
listing = f"{len(rows)} documents:\n" + "\n".join(
|
|
f"source: {source} | path: {path} | title: {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}"
|
|
if call.name == "search_documents":
|
|
raw_pattern = call.arguments.get("pattern")
|
|
pattern = raw_pattern.strip() if isinstance(raw_pattern, str) else ""
|
|
if not pattern:
|
|
return MISSING_SEARCH_ARGS
|
|
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 (source == "") != (path == ""):
|
|
# A half-specified target is a model error — fail loud with
|
|
# the missing-args refusal instead of silently widening to a
|
|
# whole-KB search (house style).
|
|
return MISSING_SEARCH_ARGS
|
|
if source:
|
|
target = find_document(db, source, path)
|
|
if target is None:
|
|
return (
|
|
f"No document at {source}/{path} — check the list_documents output."
|
|
)
|
|
docs: list[Document] = [target]
|
|
else:
|
|
docs = all_documents(db)
|
|
matches: list[str] = []
|
|
for doc in docs:
|
|
for lineno, line in grep_document(doc.content, pattern):
|
|
matches.append(
|
|
f"{doc.source}/{doc.path}:{lineno}: {line[:SEARCH_LINE_LIMIT]}"
|
|
)
|
|
if len(matches) >= SEARCH_MAX_MATCHES:
|
|
break
|
|
if len(matches) >= SEARCH_MAX_MATCHES:
|
|
break # the global cap is hit — stop scanning
|
|
holder.tool_calls += 1 # the search executed (no-match counts too)
|
|
# Locked A5: a search never adds context — read_docs untouched.
|
|
if not matches:
|
|
shown = pattern[:100] # keep a long pattern short in the line
|
|
if source:
|
|
return NO_MATCHES_SCOPED.format(
|
|
pattern=shown, source=source, path=path
|
|
)
|
|
return NO_MATCHES.format(pattern=shown)
|
|
return "\n".join(matches)
|
|
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 | RetryPiece]:
|
|
"""Run the grounded-turn tool loop, yielding every stream piece.
|
|
|
|
Every piece (``thinking`` / ``content`` / tool calls /
|
|
:class:`RetryPiece`) is yielded as it arrives; the API layer (task 04)
|
|
turns tool-call pieces into SSE ``tool`` events and retry pieces into
|
|
SSE ``retry`` events. After the loop finishes, *holder* carries the
|
|
read documents and the executed tool-call count (re-lists included).
|
|
|
|
Retries (phase 67, owner-locked A2): every model request goes through
|
|
:func:`chat_stream_retried` — a failed round is retried **before** its
|
|
first piece (same messages, ``settings.llm_retries`` restarts, a flat
|
|
``settings.llm_retry_delay``); a round that already streamed pieces
|
|
fails the turn as before.
|
|
|
|
``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. Phase 67:
|
|
# the round goes through the retry primitive — a failure before
|
|
# the first piece restarts the request (locked A2) after a
|
|
# RetryPiece; closing the OUTER generator propagates GeneratorExit
|
|
# into ``chat_stream_retried``, whose own ``finally`` closes the
|
|
# in-flight inner ``chat_stream``, so teardown stays deterministic
|
|
# on consumer abandon. Awaiting ``aclose()`` in the ``finally`` is
|
|
# safe because it does not yield; on a fully consumed round it is
|
|
# a quiet no-op.
|
|
stream = chat_stream_retried(
|
|
llm,
|
|
cast("list[dict[str, str]]", messages),
|
|
tools=tools,
|
|
retries=settings.llm_retries,
|
|
delay=settings.llm_retry_delay,
|
|
)
|
|
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). Phase 67: the
|
|
# forced call retries under the same locked-A2 rule as the
|
|
# loop rounds.
|
|
final = chat_stream_retried(
|
|
llm,
|
|
cast("list[dict[str, str]]", messages),
|
|
tools=None,
|
|
retries=settings.llm_retries,
|
|
delay=settings.llm_retry_delay,
|
|
)
|
|
try:
|
|
async for piece in final:
|
|
yield piece
|
|
finally:
|
|
await final.aclose()
|
|
return
|