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:
2026-08-26 22:39:14 -04:00
parent 9efffcb428
commit 15c1272828
30 changed files with 3594 additions and 67 deletions
+139 -15
View File
@@ -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