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