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
+266 -7
View File
@@ -11,7 +11,7 @@ from __future__ import annotations
import asyncio
import json
from types import SimpleNamespace
from typing import Any
from typing import Any, cast
import pytest
@@ -22,6 +22,7 @@ from app.rag.llm import (
LLMClient,
LLMError,
StreamPiece,
ToolCallPiece,
)
@@ -242,21 +243,47 @@ def test_single_oversized_text_fails_actionably() -> None:
# ---------- chat streaming (phase 03) ----------
def _tool_call(
index: int,
id: str | None = None,
name: str | None = None,
arguments: str | None = None,
):
"""One fake ``delta.tool_calls[]`` partial (openai SDK shape, phase 37).
``function`` is None when neither *name* nor *arguments* is given —
mirroring the real wire, where id-only fragments carry no function.
"""
fn = None
if name is not None or arguments is not None:
fn = SimpleNamespace(name=name, arguments=arguments)
return SimpleNamespace(index=index, id=id, function=fn)
def _chunk(
content: str | None = "text", empty: bool = False, reasoning: str | None = None
content: str | None = "text",
empty: bool = False,
reasoning: str | None = None,
tool_calls: list | None = None,
finish_reason: str | None = None,
):
"""One fake ChatCompletionChunk (``choices[].delta`` shape).
``reasoning_content`` is present on the delta only when *reasoning*
is not None — mirroring the real wire, where the field exists only
when the model sends it.
``reasoning_content``, ``tool_calls`` and ``finish_reason`` are
present only when provided — mirroring the real wire, where the
fields exist only when the model sends them.
"""
if empty:
return SimpleNamespace(choices=[])
delta: SimpleNamespace = SimpleNamespace(content=content)
if reasoning is not None:
delta.reasoning_content = reasoning
return SimpleNamespace(choices=[SimpleNamespace(delta=delta)])
if tool_calls is not None:
delta.tool_calls = tool_calls
choice = SimpleNamespace(delta=delta)
if finish_reason is not None:
choice.finish_reason = finish_reason
return SimpleNamespace(choices=[choice])
class _FakeChatStream:
@@ -326,7 +353,11 @@ def _make_stream_client(
async def _collect(llm: LLMClient, messages: list[dict[str, str]]) -> list[StreamPiece]:
return [p async for p in llm.chat_stream(messages)]
"""Collect pieces from a tools-less stream (phase 37 task 02, test (a):
without tools, no ToolCallPiece can appear)."""
pieces = [p async for p in llm.chat_stream(messages)]
assert all(isinstance(p, StreamPiece) for p in pieces)
return cast("list[StreamPiece]", pieces)
def test_chat_stream_yields_deltas_in_order() -> None:
@@ -355,6 +386,9 @@ def test_chat_stream_uses_locked_generation_params() -> None:
# BOR_MAX_OUTPUT_TOKENS (default 32 768) so they are not cut off.
assert completions.kwargs["max_tokens"] == 32_768
assert completions.kwargs["messages"] == messages
# Phase 37: no tools passed ⇒ no `tools` key at all (byte-identical
# request to pre-phase-37).
assert "tools" not in completions.kwargs
def test_chat_stream_max_tokens_comes_from_settings() -> None:
@@ -455,6 +489,231 @@ def test_chat_stream_llm_error_passes_through_unwrapped() -> None:
asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))
# ---------- tool-call streaming (phase 37, task 02) ----------
#: The agent's tool list (phase 37) — the exact wire shape AGENT_TOOLS will
#: pass through (the names are whatever the caller's tools list names).
_AGENT_TOOLS: list[dict[str, Any]] = [
{
"type": "function",
"function": {
"name": "list_documents",
"description": "List the indexed documents.",
"parameters": {"type": "object", "properties": {}},
},
},
{
"type": "function",
"function": {
"name": "read_document",
"description": "Add one indexed document's full text to the context.",
"parameters": {
"type": "object",
"properties": {
"source": {"type": "string"},
"path": {"type": "string"},
},
"required": ["source", "path"],
},
},
},
]
def _collect_with_tools(
llm: LLMClient, messages: list[dict[str, str]], tools: list[dict[str, Any]]
) -> list[StreamPiece | ToolCallPiece]:
async def run() -> list[StreamPiece | ToolCallPiece]:
return [p async for p in llm.chat_stream(messages, tools=tools)]
return asyncio.run(run())
def test_chat_stream_passes_tools_when_given() -> None:
"""(e) A non-None tools list is forwarded verbatim to create()."""
llm, completions = _make_stream_client([_chunk("ok")], llm_chat_model="turbo")
_collect_with_tools(llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS)
assert completions.kwargs is not None
assert completions.kwargs["tools"] == _AGENT_TOOLS
def test_chat_stream_accumulates_tool_call_across_chunk_partials() -> None:
"""(b) name on the first partial, arguments in fragments — merged into
one ToolCallPiece with the concatenated JSON, at finish_reason."""
llm, _ = _make_stream_client(
[
_chunk(
None,
tool_calls=[
_tool_call(
0,
id="call_abc",
name="read_document",
arguments='{"source": "Homelab", "pa',
)
],
),
_chunk(None, tool_calls=[_tool_call(0, arguments='th": "kubernetes.md"}')]),
_chunk(None, finish_reason="tool_calls"),
]
)
pieces = _collect_with_tools(
llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS
)
assert pieces == [
ToolCallPiece(
id="call_abc",
name="read_document",
arguments={"source": "Homelab", "path": "kubernetes.md"},
)
]
def test_chat_stream_two_tool_calls_yielded_in_index_order() -> None:
"""(c) Indices 0 and 1, interleaved partials (index 1 seen first) —
both calls, in index order, each merged from its own fragments."""
llm, _ = _make_stream_client(
[
_chunk(
None,
tool_calls=[
_tool_call(1, id="call_b", name="read_document", arguments='{"sou')
],
),
_chunk(
None,
tool_calls=[
_tool_call(0, id="call_a", name="list_documents"),
_tool_call(1, arguments='rce": "Homelab", "path": "a.md"}')
],
),
_chunk(None, finish_reason="tool_calls"),
]
)
pieces = _collect_with_tools(
llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS
)
assert pieces == [
ToolCallPiece(id="call_a", name="list_documents", arguments={}),
ToolCallPiece(
id="call_b",
name="read_document",
arguments={"source": "Homelab", "path": "a.md"},
),
]
def test_chat_stream_tool_calls_yielded_at_stream_end_without_finish_reason() -> None:
"""The spec's other emission point: stream ends without a
finish_reason="tool_calls" chunk — pieces still materialize."""
llm, _ = _make_stream_client(
[
_chunk(
None,
tool_calls=[_tool_call(0, id="call_z", name="list_documents")],
)
]
)
pieces = _collect_with_tools(
llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS
)
assert pieces == [ToolCallPiece(id="call_z", name="list_documents", arguments={})]
def test_chat_stream_synthesizes_call_id_when_absent() -> None:
"""Wire never carried the call id ⇒ synthesized "call_<index>"."""
llm, _ = _make_stream_client(
[
_chunk(None, tool_calls=[_tool_call(2, name="read_document", arguments="{}")]),
_chunk(None, finish_reason="tool_calls"),
]
)
pieces = _collect_with_tools(
llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS
)
assert pieces == [
ToolCallPiece(
id="call_2",
name="read_document",
arguments={},
)
]
def test_chat_stream_null_arguments_become_empty_dict() -> None:
"""JSON "null" (and, by the same branch, absent arguments) ⇒ {}."""
llm, _ = _make_stream_client(
[
_chunk(
None,
tool_calls=[
_tool_call(0, id="call_n", name="list_documents", arguments="null")
],
),
_chunk(None, finish_reason="tool_calls"),
]
)
pieces = _collect_with_tools(
llm, [{"role": "user", "content": "q"}], _AGENT_TOOLS
)
assert pieces == [ToolCallPiece(id="call_n", name="list_documents", arguments={})]
def test_chat_stream_malformed_tool_arguments_raise_llm_error() -> None:
"""(d) A silently dropped tool call would corrupt the loop — malformed
arguments JSON must fail loudly."""
llm, _ = _make_stream_client(
[
_chunk(
None,
tool_calls=[
_tool_call(
0,
id="call_x",
name="read_document",
arguments='{"source": "Homelab",',
)
],
),
_chunk(None, finish_reason="tool_calls"),
]
)
async def drain() -> None:
async for _ in llm.chat_stream(
[{"role": "user", "content": "q"}], tools=_AGENT_TOOLS
):
pass
with pytest.raises(LLMError, match="malformed tool-call arguments"):
asyncio.run(drain())
def test_chat_stream_non_object_tool_arguments_raise_llm_error() -> None:
"""The OpenAI contract says arguments is a JSON *object* — a bare array
is malformed too."""
llm, _ = _make_stream_client(
[
_chunk(
None,
tool_calls=[
_tool_call(0, id="call_y", name="read_document", arguments='[1, 2]')
],
),
_chunk(None, finish_reason="tool_calls"),
]
)
async def drain() -> None:
async for _ in llm.chat_stream(
[{"role": "user", "content": "q"}], tools=_AGENT_TOOLS
):
pass
with pytest.raises(LLMError, match="non-object tool-call arguments"):
asyncio.run(drain())
# ---------- one-shot chat: LLMClient.chat (phase 30, task 01) ----------