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:
+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
|
||||
|
||||
Reference in New Issue
Block a user