feat(rag): agent document tools — list/read tools with env-tuned budgets, SSE tool events + "calling tool" UI
Grounded chat turns now run the agent loop (app/rag/agent.py) instead
of a bare chat_stream: while the per-turn budgets last
(BOR_AGENT_LIST_CALLS / BOR_AGENT_READ_CALLS, default 1 each) the model
gets list_documents (the indexed catalog, /api/docs order) and
read_document (full text, never truncated — A7-revised contract); once
both budgets are spent the tools key is dropped from the request and
the model must answer. Rejected calls (unknown tool, unknown/missing
path, document already in context, spent budget) consume no budget.
Budgets 0/0 make exactly one tools=None request — byte-identical to
the pre-phase path (budgets-as-kill-switch). Deflected turns keep the
direct chat_stream (A8 unchanged; the LOW prompt never carries the
<tools> section).
SSE contract gains {"type":"tool","name":...,"argument":
"source/path"|null} frames ahead of the answer deltas (PLAN §4
extension, owner permission 2026-08-26); done.sources, query_log.sources
and the per-turn log line (gains tool_calls=N) report the retrieval
docs + read docs, deduped. The UI shows a "calling tool"
button/label state and one visible .tool-call line per call above the
answer; the lines persist with the chat record and re-render on
reload. chat_stream passes tools through and accumulates streaming
tool_calls deltas into ToolCallPiece (tools=None stays byte-identical).
E2E: deterministic mock tool flow ("use your tools" + <tools> marker:
list -> read first catalog line -> quoted answer) plus the story suite
(marker flow, reload re-render, plain/deflected no-tool regressions).
Docs: .env.example + README (the two tools, the budgets, the SSE tool
frame, the "calling tool" UI state).
probe: turbo tool_calls=supported 2026-08-26 (uv run python -m
scripts.llm_probe --tools — non-streaming + streaming
finish_reason=tool_calls, indexed delta.tool_calls partials)
This commit is contained in:
+93
-5
@@ -46,6 +46,28 @@ prompts stay byte-identical to the pre-phase text (phase 15
|
||||
convention); ``TurnPlan.kb_chars`` records the length of the stored
|
||||
outline (0 when absent) and the per-turn log line records
|
||||
``kb_chars=N`` after ``tuning=N`` (PLAN §9 line extension).
|
||||
|
||||
Agent document tools (phase 37, PLAN §4 extension, owner permission
|
||||
2026-08-26): a **grounded** turn (``not plan.deflected``) no longer
|
||||
streams a bare ``chat_stream`` — it runs the agent loop
|
||||
(``app.rag.agent.run_agent``), which offers the model the two
|
||||
server-side tools ``list_documents`` / ``read_document`` while the
|
||||
per-turn budgets (``BOR_AGENT_LIST_CALLS`` / ``BOR_AGENT_READ_CALLS``,
|
||||
default 1 each) last; once both budgets are spent the ``tools`` key is
|
||||
dropped from the request and the model must answer. Each
|
||||
model-requested call streams as an SSE ``tool`` event —
|
||||
``{"type": "tool", "name": …, "argument": "source/path" | null}`` —
|
||||
ahead of the answer's ``delta`` frames. ``done.sources``,
|
||||
``query_log.sources`` and the per-turn log line all report the same
|
||||
combined source list (retrieval docs + the agent's read docs, deduped
|
||||
by ``(source, path)``, order preserved), and the log line records
|
||||
``tool_calls=N`` after ``thinking_chars=N`` (PLAN §9 line extension —
|
||||
``N`` counts budget-consuming executions; rejected calls do not
|
||||
count). **Deflected turns keep the direct ``chat_stream`` — byte-
|
||||
identical to the pre-phase path (A8):** the LOW prompt never carries
|
||||
tools, and with **both budgets at 0** ``run_agent`` makes exactly one
|
||||
``tools=None`` request, reproducing the pre-phase behavior (budgets-
|
||||
as-kill-switch).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -64,11 +86,13 @@ from app.api.steering import load_steering_notes
|
||||
from app.config import Settings, get_settings
|
||||
from app.db import db_available, get_db
|
||||
from app.models import Document, QueryLog
|
||||
from app.rag.agent import AgentHolder, run_agent
|
||||
from app.rag.llm import (
|
||||
EmbeddingError,
|
||||
LLMClient,
|
||||
LLMError,
|
||||
StreamPiece, # noqa: F401 (phase 17 typing: chat_stream yields StreamPiece)
|
||||
StreamPiece, # type of the answer pieces streamed by the agent loop
|
||||
ToolCallPiece, # phase 37: one model-requested tool call
|
||||
)
|
||||
from app.rag.overview import load_kb_overview
|
||||
from app.rag.prompts import build_deflect_prompt, build_high_prompt
|
||||
@@ -79,6 +103,7 @@ from app.schemas import (
|
||||
ChatErrorEvent,
|
||||
ChatRequest,
|
||||
ChatThinkingEvent,
|
||||
ChatToolEvent,
|
||||
SourceRef,
|
||||
)
|
||||
|
||||
@@ -260,7 +285,6 @@ async def chat(
|
||||
).model_dump()
|
||||
)
|
||||
return
|
||||
source_paths = [f"{d.source}/{d.path}" for d in plan.docs]
|
||||
messages = [
|
||||
{"role": "system", "content": plan.system_prompt},
|
||||
{"role": "user", "content": request.message},
|
||||
@@ -271,9 +295,44 @@ async def chat(
|
||||
# ahead of the ``delta`` events (PLAN §4 extension); the
|
||||
# kill-switch (``BOR_STREAM_THINKING=0``) suppresses the
|
||||
# frames, not the counting.
|
||||
# Phase 37: a grounded turn runs the agent loop instead of a
|
||||
# bare ``chat_stream`` — its ``ToolCallPiece``s stream as
|
||||
# ``tool`` events ahead of the answer. A deflected turn keeps
|
||||
# the direct ``chat_stream`` (byte-identical, A8): the LOW
|
||||
# prompt never carries tools, and with both budgets at 0
|
||||
# ``run_agent`` is a single ``tools=None`` request anyway.
|
||||
holder = AgentHolder()
|
||||
answer_stream: AsyncIterator[StreamPiece | ToolCallPiece]
|
||||
if plan.deflected:
|
||||
answer_stream = llm.chat_stream(messages)
|
||||
else:
|
||||
answer_stream = run_agent(
|
||||
llm,
|
||||
db,
|
||||
system_prompt=plan.system_prompt,
|
||||
user_message=request.message,
|
||||
seed_docs=plan.docs,
|
||||
settings=settings,
|
||||
holder=holder,
|
||||
)
|
||||
thinking_chars = 0
|
||||
try:
|
||||
async for piece in llm.chat_stream(messages): # StreamPiece (phase 17)
|
||||
async for piece in answer_stream: # StreamPiece | ToolCallPiece
|
||||
if isinstance(piece, ToolCallPiece):
|
||||
# Phase 37 (PLAN §4 extension): one SSE ``tool``
|
||||
# frame per model-requested call; ``argument`` is the
|
||||
# read_document "source/path" (null otherwise).
|
||||
yield sse_event(
|
||||
ChatToolEvent(
|
||||
name=piece.name,
|
||||
argument=(
|
||||
f"{piece.arguments.get('source')}/{piece.arguments.get('path')}"
|
||||
if piece.name == "read_document"
|
||||
else None
|
||||
),
|
||||
).model_dump()
|
||||
)
|
||||
continue
|
||||
if piece.kind == "thinking":
|
||||
thinking_chars += len(piece.text)
|
||||
if settings.stream_thinking:
|
||||
@@ -293,8 +352,35 @@ async def chat(
|
||||
).model_dump()
|
||||
)
|
||||
return
|
||||
except Exception: # noqa: BLE001 — a tool call hit the DB mid-stream
|
||||
# Phase 37: tool execution (list_catalog / find_document) runs
|
||||
# inside the stream now; a mid-turn DB failure gets the same
|
||||
# structured ``error`` event as the pre-stream retrieval path.
|
||||
logger.exception(
|
||||
"chat: tool execution failed question=%r total_ms=%d",
|
||||
request.message,
|
||||
int((time.monotonic() - started) * 1000),
|
||||
)
|
||||
yield sse_event(
|
||||
ChatErrorEvent(
|
||||
detail="The knowledge base went offline mid-question — is Postgres up?"
|
||||
).model_dump()
|
||||
)
|
||||
return
|
||||
|
||||
# 4. Durable record + required per-turn log line (PLAN §9).
|
||||
# Phase 37: the agent's read documents join the retrieval's —
|
||||
# deduped by (source, path), order preserved — and the same
|
||||
# combined list feeds done.sources, query_log.sources and the
|
||||
# log line (empty on deflected turns: the agent never runs).
|
||||
cited_docs: list[Document] = []
|
||||
seen: set[tuple[str, str]] = set()
|
||||
for doc in [*plan.docs, *holder.read_docs]:
|
||||
key = (doc.source, doc.path)
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
cited_docs.append(doc)
|
||||
source_paths = [f"{d.source}/{d.path}" for d in cited_docs]
|
||||
total_ms = int((time.monotonic() - started) * 1000)
|
||||
try:
|
||||
db.add(
|
||||
@@ -314,7 +400,8 @@ async def chat(
|
||||
|
||||
logger.info(
|
||||
"question=%r embed_ms=%d top_score=%.3f fts_hits=%d summary_hits=%d tuning=%d "
|
||||
"kb_chars=%d threshold=%.2f deflected=%s sources=%r thinking_chars=%d total_ms=%d",
|
||||
"kb_chars=%d threshold=%.2f deflected=%s sources=%r thinking_chars=%d "
|
||||
"tool_calls=%d total_ms=%d",
|
||||
request.message,
|
||||
embed_ms,
|
||||
plan.top_score,
|
||||
@@ -326,13 +413,14 @@ async def chat(
|
||||
plan.deflected,
|
||||
source_paths,
|
||||
thinking_chars,
|
||||
holder.tool_calls,
|
||||
total_ms,
|
||||
)
|
||||
yield sse_event(
|
||||
ChatDoneEvent(
|
||||
deflected=plan.deflected,
|
||||
sources=[
|
||||
SourceRef(source=d.source, path=d.path, title=d.title) for d in plan.docs
|
||||
SourceRef(source=d.source, path=d.path, title=d.title) for d in cited_docs
|
||||
],
|
||||
suggestions=plan.suggestions,
|
||||
).model_dump()
|
||||
|
||||
@@ -87,6 +87,13 @@ class Settings(BaseSettings):
|
||||
#: ``app.rag.overview``). Overflow is cut at the cap and the shared
|
||||
#: ``[…truncated…]`` marker is appended (summarizer convention).
|
||||
overview_input_max_chars: int = 40_000
|
||||
#: Per-turn opportunities to call the ``list_documents`` agent tool
|
||||
#: (phase 37, ``app.rag.agent``); 0 disables the tool entirely
|
||||
#: (pre-phase behavior with both budgets at 0).
|
||||
agent_list_calls: int = 1
|
||||
#: Per-turn opportunities to call the ``read_document`` agent tool
|
||||
#: (phase 37, ``app.rag.agent``); 0 disables the tool entirely.
|
||||
agent_read_calls: int = 1
|
||||
|
||||
# --- Hybrid retrieval (A7, revised 2026-08-21) ---
|
||||
# cosine top-N ∪ Postgres FTS top-N, fused with Reciprocal Rank Fusion
|
||||
|
||||
@@ -0,0 +1,300 @@
|
||||
"""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
|
||||
+139
-15
@@ -4,9 +4,10 @@ Provides the embeddings surface (importer, retrieval), one-shot chat
|
||||
completions (phase 30: the ``lite`` model summarizes non-markdown
|
||||
documents at import time), and chat streaming (PLAN A15) for the RAG
|
||||
pipeline. Chat streaming yields typed :class:`StreamPiece` values
|
||||
(phase 17): aipi's ``turbo`` model streams its reasoning as
|
||||
``delta.reasoning_content`` chunks (deepseek/litellm wire convention,
|
||||
verified live 2026-08-23) **before** the answer's
|
||||
(phase 17) and — when the caller passes a ``tools`` list —
|
||||
:class:`ToolCallPiece` values (phase 37): aipi's ``turbo`` model streams
|
||||
its reasoning as ``delta.reasoning_content`` chunks (deepseek/litellm
|
||||
wire convention, verified live 2026-08-23) **before** the answer's
|
||||
``delta.content`` chunks, and reasoning counts against ``max_tokens``
|
||||
(an answer can in principle be empty).
|
||||
|
||||
@@ -17,10 +18,11 @@ vectors that pgvector rejects.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal, cast
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
@@ -55,6 +57,78 @@ class StreamPiece:
|
||||
text: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ToolCallPiece:
|
||||
"""One model-requested tool call accumulated from stream deltas (phase 37).
|
||||
|
||||
``id`` is the model's tool_call id (synthesized as ``call_<index>``
|
||||
when the wire never carried one), ``name`` is the function name
|
||||
(whatever the caller's ``tools`` list names — for the agent loop,
|
||||
``list_documents`` / ``read_document``), and ``arguments`` is the
|
||||
parsed JSON object (``{}`` when the model sent none).
|
||||
"""
|
||||
|
||||
id: str # the model's tool_call id; synthesized "call_<index>" when absent
|
||||
name: str # "list_documents" | "read_document" (whatever AGENT_TOOLS names)
|
||||
arguments: dict[str, Any]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ToolCallSlot:
|
||||
"""Mutable accumulator for one streamed tool call (phase 37, private).
|
||||
|
||||
``id`` and ``function.name`` arrive on the first partial for an index;
|
||||
``function.arguments`` arrives in fragments to concatenate (OpenAI wire
|
||||
convention, verified live against aipi 2026-08-26).
|
||||
"""
|
||||
|
||||
id: str | None = None
|
||||
name: str = ""
|
||||
arguments: str = ""
|
||||
|
||||
|
||||
def _materialize_tool_calls(
|
||||
slots: dict[int, _ToolCallSlot],
|
||||
) -> list[ToolCallPiece]:
|
||||
"""Turn accumulated slots into ordered :class:`ToolCallPiece` values.
|
||||
|
||||
Malformed ``arguments`` JSON raises :class:`LLMError` — a silently
|
||||
dropped tool call would corrupt the agent loop (fail-loud house
|
||||
style). Empty/``null`` arguments become ``{}`` (a no-parameter call
|
||||
such as ``list_documents``).
|
||||
"""
|
||||
pieces: list[ToolCallPiece] = []
|
||||
for index in sorted(slots):
|
||||
slot = slots[index]
|
||||
raw = slot.arguments.strip()
|
||||
label = slot.name or f"index {index}"
|
||||
if raw:
|
||||
try:
|
||||
parsed: Any = json.loads(raw)
|
||||
except json.JSONDecodeError as e:
|
||||
raise LLMError(
|
||||
f"model sent malformed tool-call arguments for '{label}': "
|
||||
f"{raw[:200]!r} ({e})"
|
||||
) from e
|
||||
else:
|
||||
parsed = None
|
||||
if parsed is None:
|
||||
arguments: dict[str, Any] = {}
|
||||
elif isinstance(parsed, dict):
|
||||
arguments = cast("dict[str, Any]", parsed)
|
||||
else:
|
||||
raise LLMError(
|
||||
f"model sent non-object tool-call arguments for '{label}': "
|
||||
f"{raw[:200]!r}"
|
||||
)
|
||||
pieces.append(
|
||||
ToolCallPiece(
|
||||
id=slot.id or f"call_{index}", name=slot.name, arguments=arguments
|
||||
)
|
||||
)
|
||||
return pieces
|
||||
|
||||
|
||||
# aipi's local embedding model rejects requests over ~1024 input tokens
|
||||
# ("input is too large to process"). Batch by estimated tokens, with a
|
||||
# safety margin under that cap — code-dense text can tokenize at ~3
|
||||
@@ -232,8 +306,10 @@ class LLMClient:
|
||||
return content.strip()
|
||||
|
||||
async def chat_stream(
|
||||
self, messages: list[dict[str, str]]
|
||||
) -> AsyncIterator[StreamPiece]:
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
|
||||
"""Stream assistant pieces from the chat model (PLAN A5/A15, phase 17).
|
||||
|
||||
``stream=True`` against the OpenAI-compatible endpoint, yielding
|
||||
@@ -252,24 +328,61 @@ class LLMClient:
|
||||
32 768) output tokens — the old hard 700-token cap cut long
|
||||
answers off mid-sentence (owner report 2026-08-22).
|
||||
|
||||
Tool calls (phase 37): when *tools* (an OpenAI ``tools`` list) is
|
||||
not None it is passed through as ``tools=…``; when None the key is
|
||||
**not** included, so the request is byte-identical to pre-phase-37
|
||||
and no tool pieces can be produced. A tool-calling model replies
|
||||
with ``delta.tool_calls`` partials — keyed by ``index``, with
|
||||
``id`` and ``function.name`` on the first partial and
|
||||
``function.arguments`` in fragments — which are accumulated into
|
||||
one :class:`ToolCallPiece` per call, yielded in index order at
|
||||
stream end (or immediately once a chunk carries
|
||||
``finish_reason="tool_calls"``). Malformed ``arguments`` JSON
|
||||
raises :class:`LLMError`. Wire convention verified live against
|
||||
aipi's ``turbo`` on 2026-08-26 via
|
||||
``uv run python -m scripts.llm_probe --tools`` (phase 37, task 01:
|
||||
``probe: turbo tool_calls=supported 2026-08-26``).
|
||||
|
||||
Any failure (network, HTTP, malformed stream) surfaces as
|
||||
:class:`LLMError` so the API layer can turn it into an SSE
|
||||
``error`` event instead of a hung request.
|
||||
"""
|
||||
try:
|
||||
kwargs: dict[str, Any] = {
|
||||
# ``{role, content}`` dicts are exactly what the message params
|
||||
# accept; the cast keeps pyright honest about the SDK's union.
|
||||
stream = await self._client.chat.completions.create(
|
||||
model=self.settings.llm_chat_model,
|
||||
messages=cast("list[ChatCompletionMessageParam]", messages),
|
||||
temperature=0.4,
|
||||
max_tokens=self.settings.max_output_tokens,
|
||||
stream=True,
|
||||
)
|
||||
"model": self.settings.llm_chat_model,
|
||||
"messages": cast("list[ChatCompletionMessageParam]", messages),
|
||||
"temperature": 0.4,
|
||||
"max_tokens": self.settings.max_output_tokens,
|
||||
"stream": True,
|
||||
}
|
||||
if tools is not None:
|
||||
kwargs["tools"] = tools
|
||||
try:
|
||||
stream = await self._client.chat.completions.create(**kwargs)
|
||||
calls: dict[int, _ToolCallSlot] = {}
|
||||
emitted = False
|
||||
async for chunk in stream:
|
||||
if not chunk.choices:
|
||||
continue
|
||||
delta = chunk.choices[0].delta
|
||||
choice = chunk.choices[0]
|
||||
delta = choice.delta
|
||||
# Tool-call partials (phase 37) accumulate across chunks,
|
||||
# keyed by index; a missing index (not seen on aipi) falls
|
||||
# back to the next synthetic slot.
|
||||
for tc in getattr(delta, "tool_calls", None) or []:
|
||||
idx = getattr(tc, "index", None)
|
||||
key = idx if isinstance(idx, int) else (max(calls) + 1 if calls else 0)
|
||||
slot = calls.setdefault(key, _ToolCallSlot())
|
||||
tc_id = getattr(tc, "id", None)
|
||||
if tc_id and slot.id is None:
|
||||
slot.id = tc_id
|
||||
fn = getattr(tc, "function", None)
|
||||
if fn is not None:
|
||||
if fn.name:
|
||||
slot.name += fn.name
|
||||
if fn.arguments:
|
||||
slot.arguments += fn.arguments
|
||||
reasoning = getattr(delta, "reasoning_content", None)
|
||||
if not reasoning:
|
||||
# Future-proofing: the same wire convention under a
|
||||
@@ -280,6 +393,17 @@ class LLMClient:
|
||||
content = delta.content
|
||||
if content:
|
||||
yield StreamPiece("content", content)
|
||||
if (
|
||||
calls
|
||||
and not emitted
|
||||
and getattr(choice, "finish_reason", None) == "tool_calls"
|
||||
):
|
||||
for piece in _materialize_tool_calls(calls):
|
||||
yield piece
|
||||
emitted = True
|
||||
if calls and not emitted:
|
||||
for piece in _materialize_tool_calls(calls):
|
||||
yield piece
|
||||
except LLMError:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 — wrap transport-level failures
|
||||
|
||||
+32
-5
@@ -23,6 +23,12 @@ the ``<tuning>`` section (order: ``<relevance>`` →
|
||||
``<knowledge_base>`` → ``<tuning>`` → mode body) — the agent knows
|
||||
roughly what the KB contains before retrieval. With an empty row the
|
||||
prompt is byte-identical to the pre-phase text.
|
||||
|
||||
Agent tools (phase 37): the **HIGH** prompt only carries a ``<tools>``
|
||||
section after the ``<documents>`` body — the grounded turn may call the
|
||||
server-side ``list_documents`` / ``read_document`` tools (budgeted, see
|
||||
:mod:`app.rag.agent`). The LOW/deflection prompt never carries it and
|
||||
stays byte-identical to the pre-phase text.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -66,6 +72,25 @@ _KB_INTRO = (
|
||||
"(generated at import time):\n"
|
||||
)
|
||||
|
||||
#: The ``<tools>`` instructions section — **HIGH prompt only** (phase 37,
|
||||
#: task 03): a grounded turn may extend its context through the two
|
||||
#: server-side tools (budgets: ``BOR_AGENT_LIST_CALLS`` /
|
||||
#: ``BOR_AGENT_READ_CALLS``, see :mod:`app.rag.agent`). Appended after
|
||||
#: the mode body (``<documents>``), so the instructions are the last
|
||||
#: thing the model reads. The LOW/deflection prompt never carries it —
|
||||
#: a deflection has no grounded context to extend — and stays
|
||||
#: byte-identical to the pre-phase text. The E2E mock keys off the
|
||||
#: ``<tools>`` marker's *presence*, not this wording.
|
||||
TOOLS_SECTION: str = (
|
||||
"<tools>\n"
|
||||
"If the documents in your context reference other files, or you need "
|
||||
"content that is not included above, call `list_documents` to see what "
|
||||
"is indexed, then `read_document` to pull in exactly one more document. "
|
||||
"Answer as soon as you have what you need — do not read more than one "
|
||||
"extra document.\n"
|
||||
"</tools>"
|
||||
)
|
||||
|
||||
|
||||
def _base(relevance: str) -> str:
|
||||
if relevance not in ("HIGH", "LOW"):
|
||||
@@ -155,11 +180,13 @@ def build_high_prompt(
|
||||
kb_overview: str | None = None,
|
||||
) -> str:
|
||||
"""Grounded turn: locked persona (+ steering, + KB overview) + full
|
||||
texts of the top documents.
|
||||
texts of the top documents + the ``<tools>`` instructions (phase 37).
|
||||
|
||||
Section order (phase 31): ``<relevance>`` → ``<knowledge_base>`` →
|
||||
``<tuning>`` → ``<documents>``; empty steering/overview omit their
|
||||
section, keeping the prompt byte-identical to the pre-phase text.
|
||||
Section order: ``<relevance>`` → ``<knowledge_base>`` → ``<tuning>``
|
||||
→ ``<documents>`` → ``<tools>``; empty steering/overview omit their
|
||||
section. ``<tools>`` is always present in the HIGH prompt (the
|
||||
budgets — not the prompt — decide whether the tools are actually
|
||||
offered to the model, see :mod:`app.rag.agent`).
|
||||
"""
|
||||
blocks = [
|
||||
f'<document source="{doc.source}" path="{doc.path}" title="{doc.title}">\n'
|
||||
@@ -174,7 +201,7 @@ texts of the top documents.
|
||||
for part in (build_kb_section(kb_overview or ""), build_steering_section(notes or [])):
|
||||
if part:
|
||||
prompt += "\n" + part
|
||||
return prompt + "\n<documents>\n" + body + "\n</documents>"
|
||||
return prompt + "\n<documents>\n" + body + "\n</documents>\n" + TOOLS_SECTION
|
||||
|
||||
|
||||
def build_deflect_prompt(
|
||||
|
||||
@@ -3,6 +3,7 @@ from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
@@ -60,6 +61,25 @@ class ChatThinkingEvent(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
class ChatToolEvent(BaseModel):
|
||||
"""SSE frame for one agent tool call (phase 37, PLAN §4 extension).
|
||||
|
||||
A15 extension (owner permission 2026-08-26): a grounded turn may call
|
||||
the server-side document tools (``list_documents`` / ``read_document``,
|
||||
see :mod:`app.rag.agent`); each model-requested call streams as
|
||||
``{type: "tool", name: str, argument: str | null}`` ahead of the
|
||||
answer's ``delta`` frames. ``argument`` is the read document's
|
||||
``"source/path"`` for ``read_document`` and null otherwise. The client
|
||||
renders each frame as a "calling tool" line/state (phase 37 task 05);
|
||||
the ``delta`` / ``done`` shapes are unchanged — the read document is
|
||||
reflected in ``done.sources`` instead.
|
||||
"""
|
||||
|
||||
type: Literal["tool"] = "tool"
|
||||
name: str # "list_documents" | "read_document"
|
||||
argument: str | None = None # "source/path" for read_document
|
||||
|
||||
|
||||
class ChatDoneEvent(BaseModel):
|
||||
"""Final SSE event of a chat turn: metadata for the finished answer."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user