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. 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