Files
brain-of-reese/app/rag/agent.py
T

575 lines
25 KiB
Python

"""Agent loop: the grounded-turn document tools (phase 37, task 03; the
harness-aligned ``ls``/``read``/``grep`` surface, phase 70).
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"): ``ls``, ``read`` and ``grep`` can
each be called as many times as the model needs, re-lists and
re-greps 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). Phase 70 (owner permission
2026-09-03: "match existing harnesses as much as possible") renamed
and reshaped the tools to the harness-trained surface —
``ls(path?)`` / ``read(path)`` / ``grep(pattern, path?)``, the
pi.dev tool shapes the model was trained on: the combined
``source/path`` string is the canonical document identity in every
tool argument, refusal, and result header, and the old two-argument
split (with its self-correction and "teach the split" refusals) is
gone — the model's combined form is now simply correct. The
phase-68 A5 match/output contract rides along under the new name.
2. Each tool call the model emits is executed server-side against
Postgres only (no LLM, no network): ``ls`` 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) — optionally scoped to one source name (a ``path``
argument matching no source name is a refusal; a registered source
with no indexed documents lists as ``0 documents:`` and counts) —
``read`` takes the combined ``source/path`` string, splits it at the
FIRST ``'/'`` (source names are directory basenames — they can never
contain ``'/'``), and returns the document's **full** content
(A7-revised contract: never truncated) — and ``grep`` greps the
indexed documents (or the one document a combined ``source/path``
names) 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 grep is a **locator**, not a
context-adder: it never appends to the answer context (only
``read`` does — ``holder.read_docs`` is untouched by a grep).
3. Rejected calls get a one-line refusal and count in nothing
(``holder.tool_calls`` tracks executed calls only): unknown tool name
→ ``"Unknown tool."``; a ``read`` without a usable ``path`` (missing,
blank or non-string) → ``"read requires a string argument
'path'."``; a ``grep`` without a usable ``pattern`` (missing, blank
or non-string) → ``"grep requires a string argument
'pattern'."``; a scoped ``ls`` whose ``path`` matches no source name
→ ``"No source named '…' — check the ls output."``; a document
already in context (seed or previously read) → ``"Already in your
context."``; an unknown document (a ``read`` or scoped ``grep`` whose
combined ``source/path`` matches nothing — a bare source name, which
can never be a document, included) → ``"No document at '…' — check
the ls output."`` with the argument echoed as passed (the model sees
its own form). A grep 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:`list_source_names`,
: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.git_sources import effective_sources
from app.rag.llm import (
LLMClient,
RetryPiece,
StreamPiece,
ToolCallPiece,
chat_stream_retried,
)
from app.rag.source_removal import resolve_source_name
logger = logging.getLogger("app.agent")
#: The three agent tools (phase 70: the harness-aligned surface —
#: ``ls`` / ``read`` / ``grep``, the pi.dev tool shapes the model was
#: trained on, replacing the phase-37 list/read and phase-68 search
#: names): 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. The combined
#: ``source/path`` string is the canonical document identity in every
#: argument (phase 70, owner permission 2026-09-03).
AGENT_TOOLS: list[dict[str, Any]] = [
{
"type": "function",
"function": {
"name": "ls",
"description": (
"List the indexed documents as `source: X | path: Y | "
"title: Z` lines."
),
"parameters": {
"type": "object",
"properties": {
"path": {
"type": "string",
"description": (
"Source name to list one source's documents "
"(e.g. 'homelab'); omit to list every "
"document."
),
}
},
"required": [],
},
},
},
{
"type": "function",
"function": {
"name": "read",
"description": (
"Add the full content of one indexed document to your "
"context."
),
"parameters": {
"type": "object",
"properties": {
"path": {
"type": "string",
"description": (
"The document to add to your context, as the "
"combined `source/path` string exactly as "
"shown in the `ls` output (e.g. "
"'homelab/active/container_caddy/caddy.md')."
),
}
},
"required": ["path"],
},
},
},
{
"type": "function",
"function": {
"name": "grep",
"description": (
"Search the indexed documents for an exact string "
"(case-insensitive) and return up to 20 matching lines "
"as `source/path:line: text` — a locator, not a "
"context-adder: read the winner with `read`."
),
"parameters": {
"type": "object",
"properties": {
"pattern": {
"type": "string",
"description": (
"The exact text to search for (a plain "
"substring, not a regex)"
),
},
"path": {
"type": "string",
"description": (
"Limit the search to one document, as a "
"combined `source/path` string from the "
"`ls` output (omit to search every "
"document)."
),
},
},
"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 requires a string argument 'path'."
MISSING_SEARCH_ARGS = "grep 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 grep,
#: 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 list_source_names(db: Session) -> list[str]:
"""Every registered source name, deduped, in registry order.
The source registry (the ``git_sources`` rows — the
``BOR_GIT_SOURCES`` env fallback while the table is empty) is the
source of truth for *source* names independent of document count:
a registered source with no indexed documents still lists (as
``0 documents:`` — the scoped ``ls`` must not refuse it as unknown).
Names resolve exactly as the import pipeline indexes them
(:func:`app.rag.source_removal.resolve_source_name` — reuse, or a
scoped ``ls`` would judge the wrong names unknown, the phase-69
"RAG consistent with the registry" invariant); two rows resolving
to the same name (the phase-69 sibling case) share documents, so
the name is listed once. Module-level (not a method) so unit tests
can monkeypatch it.
"""
rows, _origin = effective_sources(db)
names: list[str] = []
for row in rows:
name = resolve_source_name(row)
if name not in names:
names.append(name)
return names
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 _resolve_path(db: Session, combined: str) -> tuple[Document | None, str, str]:
"""The combined ``source/path`` identity → document (phase 70).
The canonical document identity in every tool argument, refusal and
result header is the combined string exactly as printed in the
``ls`` output, the ``Document …`` result headers, and the grep
result lines. Source names are directory basenames (``app.rag.importer``:
``source = root.name``) and can never contain a ``'/'``, so the
split at the FIRST slash is exact: the part before is the source
name, the part after is the path. Returns ``(doc, source, path)``
with the split pair (so callers can echo the canonical form, e.g.
the scoped no-match line); no ``'/'`` in the argument →
``(None, combined, "")`` — a bare source name is never a document
(no DB lookup; the refusal echoes the argument as passed).
"""
if "/" not in combined:
return None, combined, ""
source, _, path = combined.partition("/")
return find_document(db, source, path), source, path
def all_documents(db: Session) -> list[Document]:
"""Every indexed document (full rows), ordered by ``(source, path)``
— catalog order.
The whole-KB ``grep`` 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`` 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 ``grep`` never does —
it is a locator, locked A5); rejected calls return their refusal
line and count in nothing. A grep that ran but found nothing is
still a successful (counted) call — its no-match line is a result,
not a refusal. Document targets are combined ``source/path``
strings, resolved by :func:`_resolve_path` (the canonical identity,
phase 70).
"""
if call.name == "ls":
raw_path = call.arguments.get("path")
scope = raw_path.strip() if isinstance(raw_path, str) else ""
rows = list_catalog(db)
if scope:
if scope not in list_source_names(db):
return f"No source named '{scope}' — check the ls output."
rows = [row for row in rows if row[0] == scope]
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":
raw_path = call.arguments.get("path")
arg = raw_path.strip() if isinstance(raw_path, str) else ""
if not arg:
return MISSING_READ_ARGS
known = {(doc.source, doc.path) for doc in (*seed_docs, *holder.read_docs)}
# The dedupe check needs no DB: the split pair of a combined
# identity that is in context is in `known` as-is (the resolve
# below would find the same document).
if "/" in arg:
src, _, p = arg.partition("/")
if (src, p) in known:
return ALREADY_IN_CONTEXT
doc, _source, _path = _resolve_path(db, arg)
if doc is None:
# Echo the argument as passed — the model sees its own form
# (a bare source name can never be a document, no DB lookup).
return f"No document at '{arg}' — check the ls output."
holder.read_docs.append(doc)
holder.tool_calls += 1
return f"Document {doc.source}/{doc.path}:\n{doc.content}"
if call.name == "grep":
raw_pattern = call.arguments.get("pattern")
pattern = raw_pattern.strip() if isinstance(raw_pattern, str) else ""
if not pattern:
return MISSING_SEARCH_ARGS
raw_path = call.arguments.get("path")
scope = raw_path.strip() if isinstance(raw_path, str) else ""
scoped_to: tuple[str, str] | None = None
if scope:
target, src, p = _resolve_path(db, scope)
if target is None:
return f"No document at '{scope}' — check the ls output."
docs: list[Document] = [target]
scoped_to = (src, p) # the resolved (canonical) identity
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 grep executed (no-match counts too)
# Locked A5: a grep never adds context — read_docs untouched.
if not matches:
shown = pattern[:100] # keep a long pattern short in the line
if scoped_to is not None:
# The scoped no-match line is keyed on the resolved
# source/path (== the argument, stripped).
return NO_MATCHES_SCOPED.format(
pattern=shown, source=scoped_to[0], path=scoped_to[1]
)
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