Experimental/brain-of-reese
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(rag): agent document tools — list/read tools with env-tuned budgets, SSE tool events + "calling tool" UI

Browse Source 
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:
ducoterra
2026-08-26 22:39:14 -04:00
parent 9efffcb428
commit 15c1272828
30 changed files with 3594 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/e2e/mock_llm.py
+183 -9
View File
@@ -38,9 +38,33 @@ Implements just enough of the aipi surface:
the same echo convention for the overview's prompt injection.
- user message containing ``show the end of your notes`` (phase 24,
whole-document context) -> the answer quotes the **last 160 chars of
the document context** — a tail echo, byte-stable across runs, so a
sentinel placed at the *end* of a document appears in the rendered
answer iff the whole document was in the prompt.
the ``<documents>`` block** — a tail echo, byte-stable across runs, so
a sentinel placed at the *end* of a document appears in the rendered
answer iff the whole document was in the prompt. (Phase 37: the HIGH
prompt now ends with a ``<tools>`` section after ``</documents>``, so
the echo targets the block itself; its tail still includes the
closing tag — same sentinel semantics.)
- user message containing ``use your tools`` (phase 37, agent document
tools) **and** the system prompt carries the ``<tools>`` section ->
the deterministic tool-calling flow, discriminated statelessly from
the messages + the ``tools`` parameter:
* request 1 (``tools`` offered, no tool results yet): stream ONLY
``tool_calls`` deltas — ``list_documents`` (synthetic id
``call_0``, no arguments), ``finish_reason: "tool_calls"``, no
content;
* request 2 (a ``tool``-role catalog result in the messages):
parse the FIRST catalog line (``source/path — title`` → split on
``" — "`` → ``rsplit("/", 1)``) and stream a ``tool_calls`` delta
calling ``read_document`` on it (id ``call_1``);
* request 3 (the read result in the messages, no ``tools``
parameter): a content answer, deterministic: ``Read
<source/path>. <first 80 chars of the read document's content>``
— so a suite can assert the read document reached the model and
landed in the answer.
All other requests (including the marker without a ``<tools>``
section, or with the tool conversation not yet started and no tools
offered — e.g. budgets 0/0) behave exactly as today. ``E2E_REAL_LLM=1``
ignores the mock entirely (the real model does what it does).
``max_tokens`` is honored deterministically (token ≈ whitespace word),
like a real endpoint: an answer longer than the cap is truncated. This
@@ -126,6 +150,67 @@ PRE_CONTENT_PAUSE_S = 4.0
#: other suite is unaffected.
END_OF_NOTES_TRIGGER = "show the end of your notes"
#: The ``<documents>`` block of the system prompt (phase 37: the HIGH
#: prompt ends with the ``<tools>`` section after ``</documents>``, so the
#: phase-24 tail echo targets the block, not the raw message tail).
_DOCUMENTS_BLOCK_RE = re.compile(r"<documents>.*?</documents>", re.S)
#: Phase 37 (agent-document-tools story): a user message containing this
#: substring (case-insensitive) — combined with the ``<tools>`` section
#: in the system prompt — drives the deterministic tool flow documented
#: in the module docstring (list_documents → read_document on the first
#: catalog line → the quoted answer). Existing E2E questions do not
#: contain the phrase, so every other suite is unaffected.
TOOLS_TRIGGER = "use your tools"
#: The agent's ``read_document`` tool-result prefix (app.rag.agent
#: ``_execute_tool``): ``"Document <source/path>:\n<content>"``.
_READ_RESULT_PREFIX = "Document "
def _tool_flow(body: dict[str, Any]) -> tuple[str, str, str] | None:
"""Classify a marker request into one step of the tool flow (phase 37).
Returns one of:
* ``("list", "", "")`` — ``tools`` are offered and no tool results
are in the messages yet: the model lists the catalog.
* ``("read", source, path)`` — a ``tool``-role catalog result is in
the messages: the model reads its FIRST ``source/path — title``
line (split on ``" — "``, then ``rsplit("/", 1)``).
* ``("answer", "source/path", content)`` — a ``tool``-role read
result (``"Document <source/path>:\n<content>"``) is in the
messages: the model answers, quoting the read document.
* ``None`` — not the marker flow: the request behaves exactly as
today (marker absent, no ``<tools>`` section, or a no-tools first
request — the budgets-0/0 path).
"""
if TOOLS_TRIGGER not in _user(body).lower():
return None
if "<tools>" not in _system(body):
return None
tool_msgs = [m for m in _messages(body) if m.get("role") == "tool"]
for m in tool_msgs: # a read result means the forced-answer request
content = str(m.get("content") or "")
if content.startswith(_READ_RESULT_PREFIX):
# The header is "Document <source/path>:" — drop the prefix
# AND the trailing colon so the answer quotes a clean path.
header, _, doc_content = content.partition("\n")
sp = header[len(_READ_RESULT_PREFIX):].strip().removesuffix(":")
return ("answer", sp, doc_content)
if not body.get("tools"):
return None
for m in tool_msgs: # a catalog result means the read request
content = str(m.get("content") or "")
for line in content.splitlines():
head = line.split(" — ", 1)[0].strip()
if "/" in head:
source, _, path = head.rpartition("/")
if source and path:
return ("read", source, path)
return ("list", "", "")
def long_answer() -> str:
"""~900-word deterministic walkthrough (phase 11): numbered steps plus
@@ -222,12 +307,17 @@ def compose_answer(body: dict[str, Any]) -> str:
)
elif END_OF_NOTES_TRIGGER in user.lower():
# Whole-document-context story (phase 24): echo the tail of the
# context. Byte-stable across runs — a sentinel on the document's
# last line appears in the answer iff the whole document was in
# the prompt. (The tail includes the closing </documents> —
# harmless for the E2E sentinel assertions.)
# document context. Byte-stable across runs — a sentinel on the
# document's last line appears in the answer iff the whole
# document was in the prompt. (The tail includes the closing
# </documents> — harmless for the E2E sentinel assertions.)
# Phase 37: the HIGH prompt now ends with the <tools> section
# after </documents>, so the echo targets the <documents> block
# itself — the sentinel semantics are unchanged.
block = _DOCUMENTS_BLOCK_RE.search(_system(body))
tail_source = block.group(0) if block else _context(body)
answer = (
f"…and the very end of my notes reads: “{_context(body)[-160:]}” "
f"…and the very end of my notes reads: “{tail_source[-160:]}” "
"(Deterministic mock answer for E2E.)"
)
else:
@@ -436,11 +526,95 @@ def _apply_max_tokens(answer: str, max_tokens: Any) -> str:
return " ".join(words[:max_tokens])
def _tool_call_stream(name: str, arguments: dict[str, Any], call_id: str) -> Any:
"""SSE frames for one tool-call-only chat completion (phase 37).
The OpenAI wire convention the app accumulates (``app/rag/llm.py``):
the first partial of index 0 carries ``id`` + ``type`` +
``function.name`` plus the first ``function.arguments`` fragment;
the remaining fragments (deterministic 16-char split — so the
multi-fragment accumulation path is exercised) arrive on later
chunks; the final chunk carries ``finish_reason: "tool_calls"``.
No ``content`` / ``reasoning_content`` frames — the turn asked for a
tool instead of answering.
Pacing: 0.1 s per frame — deliberately SLOWER than the content
stream's 0.02 s, so the UI's transient "calling tool" state (held
from the first ``tool`` frame until the first answer ``delta``) is a
comfortable observation window for the story E2E (~1 s across the
two tool requests).
"""
model = "turbo"
chunk_id = f"chatcmpl-{uuid.uuid4()}"
raw_args = json_dumps(arguments) if arguments else "{}"
frags = [raw_args[i : i + 16] for i in range(0, len(raw_args), 16)] or ["{}"]
for i, frag in enumerate(frags):
tc: dict[str, Any] = {"index": 0, "function": {"arguments": frag}}
delta: dict[str, Any] = {"tool_calls": [tc]}
if i == 0:
tc = {
"index": 0,
"id": call_id,
"type": "function",
"function": {"name": name, "arguments": frag},
}
delta = {"role": "assistant", "tool_calls": [tc]}
payload = {
"id": chunk_id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [{"index": 0, "delta": delta, "finish_reason": None}],
}
yield f"data: {json_dumps(payload)}\n\n"
time.sleep(0.1)
yield (
"data: "
+ json_dumps(
{
"id": chunk_id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}],
}
)
+ "\n\n"
)
yield "data: [DONE]\n\n"
@app.post("/v1/chat/completions")
def chat_completions(body: dict[str, Any]) -> Any:
user_lower = _user(body).lower()
# Phase 37 (agent document tools): the deterministic marker flow.
# The app's chat path is the only streaming consumer of this mock, so
# the flow handles streaming requests; a non-streaming marker request
# (never issued by the app) falls through to the regular answer.
if body.get("stream"):
flow = _tool_flow(body)
if flow is not None:
if flow[0] == "list":
stream = _tool_call_stream("list_documents", {}, "call_0")
elif flow[0] == "read":
stream = _tool_call_stream(
"read_document",
{"source": flow[1], "path": flow[2]},
"call_1",
)
else: # "answer" — quote the read document (first 80 chars)
answer = _apply_max_tokens(
f"Read {flow[1]}. {flow[2][:80]}", body.get("max_tokens")
)
stream = _sse_stream(answer, 0.0)
return StreamingResponse(
stream,
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
)
answer = _apply_max_tokens(compose_answer(body), body.get("max_tokens"))
delay = 3.0 if "pretend to think slowly" in _user(body) else 0.0
user_lower = _user(body).lower()
thinking = compose_thinking(body) if THINKING_TRIGGER in user_lower else ""
pre_content = (
PRE_CONTENT_PAUSE_S if SLOW_PRETOKEN_TRIGGER in user_lower else 0.0