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

"""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 ``source`` argument containing a ``'/'``
   (the model passed the combined ``source/path`` form) is first
   self-corrected by splitting at the first slash (see
   :func:`_resolve_document` — source names are directory basenames and
   can never contain ``'/'``); if the split still matches nothing, the
   refusal teaches the split instead of repeating the combined form.
   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")

#: Parameter descriptions shared by ``read_document`` and
#: ``search_documents``. The model repeatedly conflated the two fields —
#: passing the combined ``source/path`` string (as printed in search
#: result lines, read-result headers and refusals) as ``source`` — so
#: the descriptions define the split explicitly: ``source`` is the part
#: BEFORE the first ``'/'``, ``path`` the part after it, with a worked
#: example in the ``read_document`` description itself.
_SOURCE_PARAM: dict[str, Any] = {
    "type": "string",
    "description": (
        "Top-level source name only (e.g. 'homelab') — the part BEFORE "
        "the first '/' of a combined 'source/path' string, exactly as "
        "shown after 'source: ' in the list_documents output. Must not "
        "contain '/' itself — do not pass the full source/path here."
    ),
}
_PATH_PARAM: dict[str, Any] = {
    "type": "string",
    "description": (
        "File path relative to the source directory (e.g. "
        "'active/container_caddy/caddy.md') — the part AFTER the first "
        "'/' of a combined 'source/path' string, exactly as shown after "
        "'path: ' in the list_documents output. Must not start with the "
        "source name."
    ),
}

#: 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. A document is identified by the "
                "(source, path) pair exactly as shown in the "
                "list_documents output: 'source' is the top-level "
                "source name only (e.g. 'homelab'), 'path' is the file "
                "path inside that source (e.g. "
                "'active/container_caddy/caddy.md'). If you only have a "
                "combined 'source/path' string (as in search_documents "
                "results), split it at the FIRST '/': the part before "
                "is the source, the part after is the path. Example: "
                "read_document(source='homelab', "
                "path='active/container_caddy/caddy.md')."
            ),
            "parameters": {
                "type": "object",
                "properties": {"source": _SOURCE_PARAM, "path": _PATH_PARAM},
                "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 (each result line's "
                "'source/path' splits at the first '/': the part before "
                "is the source, the part after is the path). 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": _SOURCE_PARAM,
                    "path": _PATH_PARAM,
                },
                "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 _resolve_document(
    db: Session, source: str, path: str
) -> tuple[Document | None, str, str]:
    """``(source, path)`` → document, with combined-form self-correction.

    The exact pair is tried first. If it misses and *source* contains a
    ``'/'``, the model passed the combined ``source/path`` form — search
    result lines, read-result headers and the generic refusal all print
    that form, so the model treats it as the document's identity. Source
    names are directory basenames (``app.rag.importer``: ``source =
    root.name``) and can never contain a ``'/'``, so the pair is retried
    at the FIRST slash: the part before is the source name, the part
    after is the path. A second candidate covers a split at a LATER
    slash (``source`` carried source + leading directories, ``path`` the
    remainder).

    Returns ``(doc, src, p)`` where ``(src, p)`` is the first-slash
    split when one was attempted (so a refusal can teach it), else the
    original pair.
    """
    doc = find_document(db, source, path)
    if doc is not None or "/" not in source:
        return doc, source, path
    split_source, _, split_path = source.partition("/")
    doc = find_document(db, split_source, split_path)
    if doc is None and path and path != split_path:
        doc = find_document(db, split_source, f"{split_path}/{path}")
    return doc, split_source, split_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. A combined-form ``source`` (containing a ``'/'``) is
    self-corrected through :func:`_resolve_document` before any 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, split_source, split_path = _resolve_document(db, source, path)
        if doc is None:
            if "/" in source:
                # Educational refusal: the combined form is the model's
                # mistake — teach the split instead of repeating it.
                return (
                    f"source must not contain '/': for '{source}' call "
                    f"read_document(source='{split_source}', "
                    f"path='{split_path}')."
                )
            return (
                f"No document at {source}/{path} — check the list_documents output."
            )
        if (doc.source, doc.path) in known:
            # A self-corrected combined form for a document already in
            # context (the raw pair above cannot have matched it).
            return ALREADY_IN_CONTEXT
        holder.read_docs.append(doc)
        holder.tool_calls += 1
        return f"Document {doc.source}/{doc.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, split_source, split_path = _resolve_document(db, source, path)
            if target is None:
                if "/" in source:
                    return (
                        f"source must not contain '/': for '{source}' use "
                        f"source='{split_source}', path='{split_path}'."
                    )
                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