feat(rag): unbounded agent tool calls behind a round cap (owner revision)
Phase 45 (owner permission 2026-08-27, TODO.md L8: "allow the LLM
to make as many tool calls as it wants"): the phase-37 per-turn tool
budgets (BOR_AGENT_LIST_CALLS / BOR_AGENT_READ_CALLS, default 1 each)
and their exhaustion refusals are removed — a grounded turn now offers
list_documents / read_document for the whole turn (re-lists included),
bounded only by the round cap:
- app/config.py: agent_max_rounds (BOR_AGENT_MAX_ROUNDS, default 10,
negative rejected) replaces agent_list_calls / agent_read_calls;
.env.example + README document the single knob; app/rag/prompts.py
docstrings follow.
- app/rag/agent.py: the loop runs tools until the model answers or
rounds >= max_rounds, at which point it forces one final no-tools
answer (the cap is the only forced exit); 0 = no tools — exactly one
tools=None request, byte-identical to the pre-phase-37 path (the
kill switch). Rejected calls (unknown tool / missing args /
already-in-context / unknown path) still consume a round, so
pathological rejected-call streams are bounded by the cap. The
per-call log line is now tool/args/round=N/M; the per-turn
tool_calls=N field and the tool SSE event are unchanged.
- tests/e2e/mock_llm.py: MULTI_READ_TRIGGER ("read two documents") —
the deterministic list -> read #1 -> read #2 -> forced-answer flow
(byte-stable "I read <sp1> and <sp2>." line), classified by the
count of tool-role read results; the phase-37 single-read flow stays
byte-identical (unit-pinned in tests/unit/test_mock_tool_flow.py).
- tests/e2e/test_agent_unlimited_tools.py (new, story suite,
mock-only): three tool frames/lines in order (one list, two reads —
the second read is what the old read budget refused) + the
both-named non-deflected answer; done.sources + chips = retrieval
doc + both reads, deduped; no budget refusal rendered; the
single-read marker flow regression (exactly one read, single tool
pair).
- .agent/PLAN.md: the phase-45 SSE revision note (owner-locked, R2) —
the only PLAN edit this phase; the phase-37 note's budget clause is
marked removed.
Unit/integration rewrites (test_agent.py round-cap matrix incl. the
kill switch and rejected-call spam, test_config.py, test_chat_api.py
agent_max_rounds=0 fixtures) landed with the server core so every gate
stays green.
uv run pytest: 756 passed, app/ coverage 99%; ruff + pyright clean;
story E2E 4/4 in isolation (ran twice); regression E2E suites
(agent_document_tools unmodified, chat_rag, smoke) green in isolation.
Also records the 45_agent_unlimited_tools todo/ -> complete/ task-file
moves (00/01/02 pending in the working tree, task 03 moves on success).
This commit is contained in:
+158
-38
@@ -46,8 +46,9 @@ Implements just enough of the aipi surface:
|
||||
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:
|
||||
the deterministic SINGLE-READ tool flow, discriminated statelessly
|
||||
from the messages (the ``tools`` parameter gates the list/read
|
||||
steps — a no-tools request with no tool results is not the flow):
|
||||
* 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
|
||||
@@ -56,15 +57,37 @@ Implements just enough of the aipi surface:
|
||||
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.
|
||||
* request 3 (a ``tool``-role read result in the messages): 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.
|
||||
Reached regardless of the ``tools`` parameter (phase 45 keeps
|
||||
the tools offered until the round cap).
|
||||
The single-read flow stops at ONE read result; the MULTI-READ
|
||||
variant below reads two.
|
||||
- user message containing BOTH ``use your tools`` AND ``read two
|
||||
documents`` (``MULTI_READ_TRIGGER``, phase 45 task 02) **and** the
|
||||
system prompt carries the ``<tools>`` section -> the deterministic
|
||||
MULTI-READ flow (list → read #1 → read #2 → answer), classified by
|
||||
the COUNT of ``tool``-role read results (content starting with the
|
||||
agent's ``"Document <source/path>:"`` prefix):
|
||||
* 0 read results, no catalog yet: ``list_documents`` (id
|
||||
``call_0``);
|
||||
* 0 read results, catalog present: ``read_document`` on the FIRST
|
||||
catalog line (id ``call_1``);
|
||||
* 1 read result: ``read_document`` on the SECOND catalog line —
|
||||
the first listing line whose ``source/path`` differs from the
|
||||
one already read (id ``call_2``); a one-document catalog
|
||||
degenerates to the single-read answer (nothing second to read);
|
||||
* 2 read results: the forced answer, byte-stable: the single-read
|
||||
shape quoting the FIRST read result, plus the line ``I read
|
||||
<sp1> and <sp2>.`` naming both read paths in read order — so a
|
||||
suite can assert the model used BOTH documents.
|
||||
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).
|
||||
offered — e.g. ``agent_max_rounds=0``) behave exactly as today.
|
||||
``E2E_REAL_LLM=1`` ignores the mock entirely (the real model does
|
||||
what it does).
|
||||
- user message containing ``show me a table`` (phase 44, markdown
|
||||
tables, TODO.md L6) -> the fixed table answer (``TABLE_ANSWER``):
|
||||
a 3-column service table, an ``<img onerror>`` XSS probe line, and
|
||||
@@ -174,6 +197,15 @@ _DOCUMENTS_BLOCK_RE = re.compile(r"<documents>.*?</documents>", re.S)
|
||||
#: contain the phrase, so every other suite is unaffected.
|
||||
TOOLS_TRIGGER = "use your tools"
|
||||
|
||||
#: Phase 45 (agent-unlimited-tools story, task 02): a user message
|
||||
#: containing BOTH ``TOOLS_TRIGGER`` and this substring (case-insensitive
|
||||
#: — the check lowercases the user message) drives the deterministic
|
||||
#: MULTI-READ tool flow (list → read #1 → read #2 → the forced answer
|
||||
#: naming both read paths) — see the module docstring. The existing
|
||||
#: phase-37 E2E question carries ``TOOLS_TRIGGER`` but not this phrase,
|
||||
#: so the 3-step flow is untouched.
|
||||
MULTI_READ_TRIGGER = "read two documents"
|
||||
|
||||
#: Phase 44 (markdown-tables story, TODO.md L6): a user message
|
||||
#: containing this substring (case-insensitive) gets the fixed table
|
||||
#: answer (``TABLE_ANSWER`` below) — a 3-column table, an XSS probe
|
||||
@@ -211,46 +243,126 @@ TABLE_ANSWER = (
|
||||
_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).
|
||||
def _read_results(body: dict[str, Any]) -> list[tuple[str, str]]:
|
||||
"""The read results in the messages, in order: ``(source/path, content)``.
|
||||
|
||||
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).
|
||||
A read result is a ``tool``-role message whose content starts with
|
||||
the agent's read-result prefix (``app.rag.agent`` ``_execute_tool``):
|
||||
``"Document <source/path>:\n<content>"``. The header is stripped of
|
||||
the prefix AND the trailing colon so the path stays clean.
|
||||
"""
|
||||
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
|
||||
out: list[tuple[str, str]] = []
|
||||
for m in _messages(body):
|
||||
if m.get("role") != "tool":
|
||||
continue
|
||||
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
|
||||
out.append((sp, doc_content))
|
||||
return out
|
||||
|
||||
|
||||
def _catalog_docs(body: dict[str, Any]) -> list[tuple[str, str]]:
|
||||
"""Every ``source/path`` in the catalog tool result, in listing order.
|
||||
|
||||
Catalog lines are ``source/path — title`` (the agent's
|
||||
``list_documents`` output): split on ``" — "``, keep the head, and
|
||||
recover ``(source, path)`` with ``rsplit("/", 1)`` (``rpartition``)
|
||||
— the same convention the single-read flow's read step uses. The
|
||||
``"N documents:"`` header line carries no ``/`` and is skipped; read-
|
||||
result messages are full documents, not listings, and are skipped
|
||||
too.
|
||||
"""
|
||||
docs: list[tuple[str, str]] = []
|
||||
for m in _messages(body):
|
||||
if m.get("role") != "tool":
|
||||
continue
|
||||
content = str(m.get("content") or "")
|
||||
if content.startswith(_READ_RESULT_PREFIX):
|
||||
continue
|
||||
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)
|
||||
docs.append((source, path))
|
||||
return docs
|
||||
|
||||
|
||||
def _tool_flow(body: dict[str, Any]) -> tuple[str, ...] | None:
|
||||
"""Classify a marker request into one step of the tool flow.
|
||||
|
||||
Single-read (phase 37 — the user message carries ``TOOLS_TRIGGER``
|
||||
only):
|
||||
|
||||
* ``("list", "", "")`` — ``tools`` are offered and no tool results
|
||||
are in the messages yet: the model lists the catalog.
|
||||
* ``("read", source, path, "call_1")`` — 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. Reached
|
||||
regardless of the ``tools`` parameter (phase 45 keeps the tools
|
||||
offered until the round cap).
|
||||
|
||||
Multi-read (phase 45 task 02 — the user message carries BOTH
|
||||
``TOOLS_TRIGGER`` and ``MULTI_READ_TRIGGER``), classified by the
|
||||
count of ``tool``-role read results:
|
||||
|
||||
* 0 read results: ``("list", "", "")`` (no catalog yet) or
|
||||
``("read", source, path, "call_1")`` on the FIRST catalog doc.
|
||||
* 1 read result: ``("read", source, path, "call_2")`` on the SECOND
|
||||
catalog doc — the first listing line whose ``source/path``
|
||||
differs from the one already read. A one-document catalog
|
||||
degenerates to the single-read ``("answer", ...)`` shape (nothing
|
||||
second to read).
|
||||
* 2 read results: ``("multi_answer", "", text)`` — the forced
|
||||
answer, byte-stable: the single-read shape quoting the FIRST read
|
||||
result, plus ``I read <sp1> and <sp2>.`` (both read paths, read
|
||||
order). The second element is unused.
|
||||
|
||||
* ``None`` — not the marker flow: the request behaves exactly as
|
||||
before (marker absent, no ``<tools>`` section, or a no-tools
|
||||
request with no tool results — e.g. ``agent_max_rounds=0``).
|
||||
"""
|
||||
user = _user(body).lower()
|
||||
if TOOLS_TRIGGER not in user:
|
||||
return None
|
||||
if "<tools>" not in _system(body):
|
||||
return None
|
||||
reads = _read_results(body)
|
||||
if MULTI_READ_TRIGGER in user:
|
||||
if not reads:
|
||||
if not body.get("tools"):
|
||||
return None
|
||||
docs = _catalog_docs(body)
|
||||
if not docs:
|
||||
return ("list", "", "")
|
||||
return ("read", docs[0][0], docs[0][1], "call_1")
|
||||
if len(reads) == 1:
|
||||
skip = reads[0][0]
|
||||
second = next(
|
||||
(d for d in _catalog_docs(body) if f"{d[0]}/{d[1]}" != skip), None
|
||||
)
|
||||
if second is None:
|
||||
# One-document catalog: nothing second to read — the
|
||||
# single-read answer shape (deterministic degenerate).
|
||||
return ("answer", reads[0][0], reads[0][1])
|
||||
return ("read", second[0], second[1], "call_2")
|
||||
(sp1, c1), (sp2, _c2) = reads[0], reads[1]
|
||||
answer = f"Read {sp1}. {c1[:80]} I read {sp1} and {sp2}."
|
||||
return ("multi_answer", "", answer)
|
||||
# Phase-37 single-read flow — byte-identical to the original.
|
||||
if reads:
|
||||
return ("answer", reads[0][0], reads[0][1])
|
||||
if not body.get("tools"):
|
||||
return None
|
||||
docs = _catalog_docs(body)
|
||||
if docs:
|
||||
return ("read", docs[0][0], docs[0][1], "call_1")
|
||||
return ("list", "", "")
|
||||
|
||||
|
||||
@@ -653,11 +765,19 @@ def chat_completions(body: dict[str, Any]) -> Any:
|
||||
if flow[0] == "list":
|
||||
stream = _tool_call_stream("list_documents", {}, "call_0")
|
||||
elif flow[0] == "read":
|
||||
# flow[3] is the synthetic call id — "call_1" for the
|
||||
# single-read flow and the multi-read first read,
|
||||
# "call_2" for the multi-read second read (phase 45,
|
||||
# task 02).
|
||||
stream = _tool_call_stream(
|
||||
"read_document",
|
||||
{"source": flow[1], "path": flow[2]},
|
||||
"call_1",
|
||||
flow[3],
|
||||
)
|
||||
elif flow[0] == "multi_answer":
|
||||
# Phase 45 (task 02): the multi-read forced answer —
|
||||
# computed in _tool_flow, byte-stable.
|
||||
stream = _sse_stream(_apply_max_tokens(flow[2], body.get("max_tokens")), 0.0)
|
||||
else: # "answer" — quote the read document (first 80 chars)
|
||||
answer = _apply_max_tokens(
|
||||
f"Read {flow[1]}. {flow[2][:80]}", body.get("max_tokens")
|
||||
|
||||
@@ -0,0 +1,585 @@
|
||||
"""Phase 45 E2E (Playwright, mock-only): as many tool calls as the model wants.
|
||||
|
||||
Story: ``.agent/user_stories/agent-unlimited-tools.md``
|
||||
Run in isolation (DB must be up: ``podman compose up -d db``):
|
||||
|
||||
uv run pytest tests/e2e/test_agent_unlimited_tools.py -v --no-cov
|
||||
|
||||
MOCK-ONLY suite: ``E2E_REAL_LLM=1`` is not supported — the gate is the
|
||||
deterministic MULTI-READ marker flow in ``tests/e2e/mock_llm.py`` (user
|
||||
message contains BOTH ``use your tools`` (``TOOLS_TRIGGER``) and ``read
|
||||
two documents`` (``MULTI_READ_TRIGGER``) **and** the system prompt
|
||||
carries the ``<tools>`` section of the HIGH prompt):
|
||||
|
||||
1. request 1 (``tools`` offered, no tool results yet) → streams ONLY
|
||||
``tool_calls`` deltas calling ``list_documents`` (id ``call_0``);
|
||||
2. request 2 (the ``tool``-role catalog result) → ``read_document`` on
|
||||
the FIRST catalog line (id ``call_1``);
|
||||
3. request 3 (one ``tool``-role read result) → ``read_document`` on the
|
||||
SECOND catalog line (id ``call_2``) — the pre-phase-45 per-tool
|
||||
budgets would have refused exactly this second read (``No reading
|
||||
budget left — answer with what you have.``);
|
||||
4. request 4 (two read results) → the forced answer, byte-stable: the
|
||||
single-read shape quoting the FIRST read result, plus the line
|
||||
``I read <sp1> and <sp2>.`` naming both read paths in read order.
|
||||
|
||||
KB fixture (the ``test_agent_document_tools.py`` TRUNCATE-then-seed
|
||||
pattern, grown to three documents):
|
||||
|
||||
* ``Deployments/aaa-record-shape.json`` — read #1: indexed (in the
|
||||
agent's catalog, readable) but seeded WITHOUT chunks, so retrieval
|
||||
never puts it in context; sorts FIRST in the catalog;
|
||||
* ``Deployments/bbb-zone-sync.yaml`` — read #2: same shape; sorts
|
||||
SECOND;
|
||||
* ``Homelab/route53-notes.md`` — the ONLY retrievable document: one
|
||||
chunk whose embedding is the mock's own bag-of-words vector (genuine
|
||||
token overlap: the marker questions cosine ≈0.65/≈0.71 against it,
|
||||
well past the E2E 0.30 threshold, and they FTS-match too) → the
|
||||
grounded seed context.
|
||||
|
||||
Three documents (not two, as in phase 37) so BOTH reads land on
|
||||
documents outside the seed: with a two-document corpus the second read
|
||||
would be the already-in-context retrieval document and the agent would
|
||||
answer "Already in your context." — a rejection, not the multi-read
|
||||
flow this story proves.
|
||||
|
||||
Test → story mapping (Playwright Mapping Rule):
|
||||
1. ``test_multi_read_turn`` — the turn streams THREE ``tool`` frames /
|
||||
``.tool-call`` lines in order (one list — "is listing documents" —
|
||||
and two reads — "is reading <source/path>" — the #send-status
|
||||
transition recorded deterministically via MutationObserver), then a
|
||||
final non-deflected answer containing the mock's byte-stable
|
||||
``I read <sp1> and <sp2>.`` line; the round cap (default 10) bounds
|
||||
the turn, no budget refusal anywhere.
|
||||
2. ``test_done_sources_include_reads`` — the source chips under the
|
||||
answer list the retrieval doc PLUS both read documents, deduped
|
||||
(the phase-37 ``done.sources`` extension contract, now with 2
|
||||
reads); the same combined list lands in ``query_log.sources``.
|
||||
3. ``test_relist_allowed`` — the listing tool ran (its line rendered)
|
||||
and no pre-phase-45 budget refusal ("… budget left") appears
|
||||
anywhere in the message bubble or tool lines: the old
|
||||
``LIST_EXHAUSTED`` / ``READ_EXHAUSTED`` refusal strings are gone
|
||||
from the product (the source-level grep was task 01's job).
|
||||
4. ``test_single_tool_flow_regression`` (phase 37) — the original
|
||||
marker WITHOUT the multi-read trigger still answers after exactly
|
||||
ONE read with its single tool pair (list + one read). The full
|
||||
phase-37 suite runs unmodified in the regression pass.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from playwright.sync_api import Page, expect
|
||||
from sqlalchemy import select, text
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.db import SessionLocal
|
||||
from app.models import Chunk, Document, QueryLog
|
||||
from tests.e2e.mock_llm import embed_text
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Fixture documents (deterministic, token-controlled)
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
READ1_SOURCE = "Deployments"
|
||||
READ1_PATH = "aaa-record-shape.json"
|
||||
READ1_SP = f"{READ1_SOURCE}/{READ1_PATH}"
|
||||
|
||||
READ2_SOURCE = "Deployments"
|
||||
READ2_PATH = "bbb-zone-sync.yaml"
|
||||
READ2_SP = f"{READ2_SOURCE}/{READ2_PATH}"
|
||||
|
||||
SEED_SOURCE = "Homelab"
|
||||
SEED_PATH = "route53-notes.md"
|
||||
SEED_SP = f"{SEED_SOURCE}/{SEED_PATH}"
|
||||
|
||||
#: The retrievable document: references the record shape "for the exact
|
||||
#: JSON shape of reeselink.json" (the TODO failure, same story as the
|
||||
#: phase-37 fixture). The repeated record-file lines carry the marker
|
||||
#: questions' key tokens (aws, route53, hosted, zone, reeselink, json,
|
||||
#: exact, shape) — verified ≈0.65 (multi question) / ≈0.71 (single
|
||||
#: question) cosine against the mock's embeddings (E2E threshold 0.30)
|
||||
#: plus FTS hits, so both turns are solidly grounded.
|
||||
ROUTE53_CONTENT = (
|
||||
"# AWS Route 53 Notes\n\n"
|
||||
"## Record file\n\n"
|
||||
+ (
|
||||
"The aws route53 hosted zone for reeselink keeps every record in "
|
||||
"reseelink.json — the exact JSON shape of reeselink.json is "
|
||||
"documented in aaa-record-shape.json.\n"
|
||||
)
|
||||
* 10
|
||||
+ "\n## Sync job\n\n"
|
||||
"A cron job pushes reeselink.json to the aws route53 hosted zone "
|
||||
"every fifteen minutes; the diff is applied through the route53 api.\n"
|
||||
)
|
||||
|
||||
#: Read #1: the JSON shape. Its FIRST line is longer than 80 chars, so
|
||||
#: the mock's first-80-chars quote is newline-free (the rendered-text
|
||||
#: assertion matches it verbatim). Pinned by the assert below.
|
||||
RECORD_CONTENT = (
|
||||
'{"version": 4, "comment": "ReeseLink hosted zone records — the exact '
|
||||
'JSON shape of reeselink.json",\n'
|
||||
' "hosted_zone_id": "Z0RESEELINK45",\n'
|
||||
' "record_sets": [\n'
|
||||
' { "name": "www.reeselink.example", "type": "A", "ttl": 300 }\n'
|
||||
' ]\n'
|
||||
"}\n"
|
||||
)
|
||||
assert "\n" not in RECORD_CONTENT[:80] # the quote must stay one line
|
||||
|
||||
#: Read #2: the sync runbook.
|
||||
RUNBOOK_CONTENT = (
|
||||
"sync:\n"
|
||||
" schedule: every fifteen minutes\n"
|
||||
" target: reeselink.json\n"
|
||||
" engine: aws route53 api\n"
|
||||
" note: the diff is applied through the route53 api\n"
|
||||
)
|
||||
|
||||
#: Carries BOTH markers — ``use your tools`` (phase 37) and ``read two
|
||||
#: documents`` (phase 45 ``MULTI_READ_TRIGGER``).
|
||||
MULTI_QUESTION = (
|
||||
"Use your tools and read two documents: what is the exact JSON shape "
|
||||
"of reeselink.json for my aws route53 hosted zone?"
|
||||
)
|
||||
#: The phase-37 marker WITHOUT the multi-read trigger — the original
|
||||
#: 3-step single-read flow (regression test 4).
|
||||
SINGLE_QUESTION = (
|
||||
"Use your tools: what is the exact JSON shape of reeselink.json "
|
||||
"for my aws route53 hosted zone?"
|
||||
)
|
||||
assert "use your tools" in MULTI_QUESTION.lower()
|
||||
assert "read two documents" in MULTI_QUESTION.lower()
|
||||
assert "read two documents" not in SINGLE_QUESTION.lower()
|
||||
|
||||
#: The mock's byte-stable multi-read answer pieces (mock_llm
|
||||
#: ``_tool_flow``): the single-read shape quoting the FIRST read result,
|
||||
#: plus both read paths in read order.
|
||||
ANSWER_PREFIX = f"Read {READ1_SP}."
|
||||
ANSWER_QUOTE = RECORD_CONTENT[:80]
|
||||
BOTH_READS_LINE = f"I read {READ1_SP} and {READ2_SP}."
|
||||
|
||||
#: The pre-phase-45 budget refusals (phase 37 ``LIST_EXHAUSTED`` /
|
||||
#: ``READ_EXHAUSTED``) — gone from the app (task 01) and never rendered
|
||||
#: (test 3). The generic "budget left" fragment covers both exact
|
||||
#: strings.
|
||||
BUDGET_REFUSAL_FRAGMENTS = (
|
||||
"No listing budget left — answer with what you have.",
|
||||
"No reading budget left — answer with what you have.",
|
||||
"budget left",
|
||||
)
|
||||
|
||||
# The combined source list the app reports (app/api/chat.py): retrieval
|
||||
# docs first, then the agent's read docs, deduped by (source, path).
|
||||
EXPECTED_SOURCES = [
|
||||
(SEED_SOURCE, SEED_PATH),
|
||||
(READ1_SOURCE, READ1_PATH),
|
||||
(READ2_SOURCE, READ2_PATH),
|
||||
]
|
||||
EXPECTED_SOURCES_LINE = ", ".join(f"{s}/{p}" for s, p in EXPECTED_SOURCES)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# DB seeding (TRUNCATE-then-seed, cf. test_agent_document_tools.py)
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _doc(source: str, path: str, title: str, content: str) -> Document:
|
||||
return Document(
|
||||
source=source,
|
||||
path=path,
|
||||
full_path=f"/tmp/{path}",
|
||||
title=title,
|
||||
content=content,
|
||||
content_hash=hashlib.sha256(content.encode()).hexdigest(),
|
||||
indexed_at=datetime.now(UTC),
|
||||
)
|
||||
|
||||
|
||||
def _seed(db: Session) -> None:
|
||||
"""The three-document KB from the module docstring."""
|
||||
md = _doc(SEED_SOURCE, SEED_PATH, "AWS Route 53 Notes", ROUTE53_CONTENT)
|
||||
db.add(md)
|
||||
db.flush()
|
||||
# One chunk carrying the mock's own embedding → genuine token
|
||||
# overlap between the marker questions and this document (the only
|
||||
# retrievable document — the grounded seed context).
|
||||
db.add(
|
||||
Chunk(
|
||||
document_id=md.id,
|
||||
position=0,
|
||||
content=ROUTE53_CONTENT,
|
||||
embedding=embed_text(ROUTE53_CONTENT),
|
||||
)
|
||||
)
|
||||
# The two read documents: indexed, catalogued, readable — but NO
|
||||
# chunks, so retrieval never puts them in context.
|
||||
db.add(_doc(READ1_SOURCE, READ1_PATH, "Record Shape", RECORD_CONTENT))
|
||||
db.add(_doc(READ2_SOURCE, READ2_PATH, "Zone Sync Runbook", RUNBOOK_CONTENT))
|
||||
|
||||
|
||||
def _reset_db(seed: Callable[[Session], None] | None = None) -> None:
|
||||
"""Truncate the KB (plus the prompt-shaping tables), then re-seed.
|
||||
|
||||
``steering_notes`` / ``kb_overview`` are truncated too, so the HIGH
|
||||
prompt is exactly ``<relevance>`` + ``<documents>`` + ``<tools>``
|
||||
regardless of leftovers from other suites — byte-stable prompts,
|
||||
byte-stable answers.
|
||||
"""
|
||||
with SessionLocal() as db:
|
||||
db.execute(
|
||||
text("TRUNCATE chunks, documents, query_log, steering_notes, kb_overview")
|
||||
)
|
||||
db.commit()
|
||||
if seed is not None:
|
||||
seed(db)
|
||||
db.commit()
|
||||
|
||||
|
||||
def _last_query_log() -> QueryLog:
|
||||
with SessionLocal() as db:
|
||||
rows = db.scalars(select(QueryLog)).all()
|
||||
assert len(rows) == 1, f"expected exactly one query_log row, got {len(rows)}"
|
||||
return rows[0]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Page helpers
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
#: Records every value #send-label takes during the turn (a
|
||||
#: MutationObserver on the element), so the transient "Calling tool…"
|
||||
#: state is captured deterministically — no polling race (the phase-37
|
||||
#: flake fix, phase 44 task 03).
|
||||
LABEL_RECORDER = """
|
||||
() => {
|
||||
if (window.__labelsInstalled) return;
|
||||
window.__labelsInstalled = true;
|
||||
window.__labels = [];
|
||||
const el = document.querySelector('#send-label');
|
||||
if (!el) return;
|
||||
const rec = (v) => {
|
||||
const l = window.__labels;
|
||||
if (!l.length || l[l.length - 1] !== v) l.push(v);
|
||||
};
|
||||
rec(el.textContent);
|
||||
new MutationObserver(() => rec(el.textContent)).observe(el, {
|
||||
childList: true,
|
||||
subtree: true,
|
||||
});
|
||||
}
|
||||
"""
|
||||
|
||||
#: Records every value #send-status takes during the turn — the
|
||||
#: "… is listing documents" / "… is reading <source/path>" tool states
|
||||
#: are transient (the first delta switches the status to the streaming
|
||||
#: state), so the pre-submit observer is the deterministic source of
|
||||
#: truth for their order.
|
||||
STATUS_RECORDER = """
|
||||
() => {
|
||||
if (window.__statusesInstalled) return;
|
||||
window.__statusesInstalled = true;
|
||||
window.__statuses = [];
|
||||
const el = document.querySelector('#send-status');
|
||||
if (!el) return;
|
||||
const rec = (v) => {
|
||||
const l = window.__statuses;
|
||||
if (!l.length || l[l.length - 1] !== v) l.push(v);
|
||||
};
|
||||
rec(el.textContent);
|
||||
new MutationObserver(() => rec(el.textContent)).observe(el, {
|
||||
childList: true,
|
||||
subtree: true,
|
||||
});
|
||||
}
|
||||
"""
|
||||
|
||||
#: Captures the raw SSE ``data:`` payloads of the /api/chat stream
|
||||
#: (a response clone read in the background) — wire-level assertions for
|
||||
#: the ``tool`` frames, independent of the UI rendering.
|
||||
SSE_HOOK = """
|
||||
() => {
|
||||
if (window.__sseInstalled) return;
|
||||
window.__sseInstalled = true;
|
||||
window.__sseFrames = [];
|
||||
const origFetch = window.fetch;
|
||||
window.fetch = async function (...args) {
|
||||
const res = await origFetch.apply(this, args);
|
||||
try {
|
||||
const url = typeof args[0] === 'string' ? args[0] : args[0].url;
|
||||
if (url.includes('/api/chat')) {
|
||||
res.clone().text().then((bodyText) => {
|
||||
for (const block of bodyText.split('\\n\\n')) {
|
||||
const line = block.trim();
|
||||
if (line.startsWith('data: ')) {
|
||||
window.__sseFrames.push(line.slice(6));
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
} catch (e) { /* non-clonable responses: ignored */ }
|
||||
return res;
|
||||
};
|
||||
}
|
||||
"""
|
||||
|
||||
|
||||
def _install_page_hooks(page: Page) -> None:
|
||||
"""Install all hooks on the loaded page (post-goto, pre-submit)."""
|
||||
page.evaluate(SSE_HOOK)
|
||||
page.evaluate(LABEL_RECORDER)
|
||||
page.evaluate(STATUS_RECORDER)
|
||||
|
||||
|
||||
def _frames(page: Page) -> list[dict]:
|
||||
"""The captured SSE frames, once the hook's background read settles."""
|
||||
deadline = time.monotonic() + 10.0
|
||||
while True:
|
||||
raw = page.evaluate("() => window.__sseFrames || []")
|
||||
parsed = [json.loads(line) for line in raw if line]
|
||||
if any(f.get("type") == "done" for f in parsed):
|
||||
return parsed
|
||||
if time.monotonic() > deadline:
|
||||
raise AssertionError(
|
||||
f"SSE hook captured no `done` frame (frames so far: "
|
||||
f"{len(parsed)}) — hook install failed?"
|
||||
)
|
||||
time.sleep(0.05)
|
||||
|
||||
|
||||
def _tool_frames(frames: list[dict]) -> list[dict]:
|
||||
return [f for f in frames if f.get("type") == "tool"]
|
||||
|
||||
|
||||
def _submit(page: Page, question: str) -> None:
|
||||
page.fill("#message-input", question)
|
||||
page.click("#send-btn")
|
||||
# The user bubble lands synchronously with the submit handler.
|
||||
expect(page.locator(".msg.user .bubble").last).to_contain_text(question)
|
||||
|
||||
|
||||
def _wait_settled(page: Page) -> None:
|
||||
"""The turn is complete: answer text in the bubble, button recovered."""
|
||||
expect(page.locator(".msg.brain .bubble").last).not_to_have_text("", timeout=30_000)
|
||||
expect(page.locator("#send-btn")).to_be_enabled(timeout=30_000)
|
||||
expect(page.locator("#send-label")).to_have_text("Send")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 1. The multi-read turn: list → read #1 → read #2 → both-named answer
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_multi_read_turn(
|
||||
page: Page, app_url: str, mock_llm: int, db_ready: None
|
||||
) -> None:
|
||||
page.set_default_timeout(30_000)
|
||||
_reset_db(_seed)
|
||||
page.goto(app_url)
|
||||
_install_page_hooks(page)
|
||||
|
||||
_submit(page, MULTI_QUESTION)
|
||||
_wait_settled(page)
|
||||
|
||||
# Wire level: exactly THREE `tool` frames — list, read #1, read #2,
|
||||
# in order — and all ahead of the first `delta` frame. This third
|
||||
# frame is the one the pre-phase-45 read budget refused.
|
||||
frames = _frames(page)
|
||||
assert _tool_frames(frames) == [
|
||||
{"type": "tool", "name": "list_documents", "argument": None},
|
||||
{"type": "tool", "name": "read_document", "argument": READ1_SP},
|
||||
{"type": "tool", "name": "read_document", "argument": READ2_SP},
|
||||
]
|
||||
first_delta = next(i for i, f in enumerate(frames) if f.get("type") == "delta")
|
||||
assert all(
|
||||
i < first_delta for i, f in enumerate(frames) if f.get("type") == "tool"
|
||||
)
|
||||
done = next(f for f in frames if f.get("type") == "done")
|
||||
assert done["deflected"] is False
|
||||
|
||||
# The transient "calling tool" states, recorded deterministically:
|
||||
# the label shows "Calling tool…" and #send-status walked through
|
||||
# "… is listing documents" then "… is reading <sp>" for BOTH reads,
|
||||
# in order.
|
||||
labels = page.evaluate("() => window.__labels")
|
||||
assert "Calling tool…" in labels, labels
|
||||
assert labels.index("Calling tool…") > labels.index("Thinking…")
|
||||
statuses = page.evaluate("() => window.__statuses")
|
||||
i_list = next(
|
||||
(i for i, s in enumerate(statuses) if "is listing documents" in s), None
|
||||
)
|
||||
i_read1 = next(
|
||||
(i for i, s in enumerate(statuses) if f"is reading {READ1_SP}" in s), None
|
||||
)
|
||||
i_read2 = next(
|
||||
(i for i, s in enumerate(statuses) if f"is reading {READ2_SP}" in s), None
|
||||
)
|
||||
assert (
|
||||
i_list is not None and i_read1 is not None and i_read2 is not None
|
||||
), statuses
|
||||
assert i_list < i_read1 < i_read2, statuses
|
||||
|
||||
# Three visible tool lines, in order, above the answer.
|
||||
lines = page.locator(".msg.brain .tool-call")
|
||||
expect(lines).to_have_count(3)
|
||||
expect(lines.nth(0)).to_contain_text("Listing documents")
|
||||
expect(lines.nth(1)).to_contain_text("Reading ")
|
||||
expect(lines.nth(1)).to_contain_text(READ1_SP)
|
||||
expect(lines.nth(2)).to_contain_text("Reading ")
|
||||
expect(lines.nth(2)).to_contain_text(READ2_SP)
|
||||
|
||||
# The final answer is non-deflected, quotes the FIRST read result,
|
||||
# and names BOTH read paths (the mock's byte-stable line).
|
||||
last = page.locator(".msg.brain").last
|
||||
expect(last).not_to_have_class(re.compile(r"is-deflected"))
|
||||
bubble = last.locator(".bubble")
|
||||
expect(bubble).to_contain_text(ANSWER_PREFIX)
|
||||
expect(bubble).to_contain_text(ANSWER_QUOTE)
|
||||
expect(bubble).to_contain_text(BOTH_READS_LINE)
|
||||
|
||||
# Durable record: grounded, combined sources (retrieval + both
|
||||
# reads).
|
||||
row = _last_query_log()
|
||||
assert row.question == MULTI_QUESTION
|
||||
assert row.deflected is False
|
||||
assert row.sources == EXPECTED_SOURCES_LINE
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 2. done.sources / source chips: retrieval doc + BOTH reads, deduped
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_done_sources_include_reads(
|
||||
page: Page, app_url: str, mock_llm: int, db_ready: None
|
||||
) -> None:
|
||||
page.set_default_timeout(30_000)
|
||||
_reset_db(_seed)
|
||||
page.goto(app_url)
|
||||
_install_page_hooks(page)
|
||||
|
||||
_submit(page, MULTI_QUESTION)
|
||||
_wait_settled(page)
|
||||
|
||||
# Wire level: done.sources is the retrieval doc FIRST, then both
|
||||
# read documents — deduped (the retrieval doc was never read, the
|
||||
# reads are each read once; nothing appears twice).
|
||||
frames = _frames(page)
|
||||
done = next(f for f in frames if f.get("type") == "done")
|
||||
assert [(s["source"], s["path"]) for s in done["sources"]] == EXPECTED_SOURCES
|
||||
pairs = [(s["source"], s["path"]) for s in done["sources"]]
|
||||
assert len(pairs) == len(set(pairs)), "done.sources must be deduped"
|
||||
|
||||
# UI: exactly three source chips under the answer, in the same
|
||||
# order, each a viewer link — no duplicated chip.
|
||||
chips = page.locator(".msg.brain .source-chip")
|
||||
expect(chips).to_have_count(3)
|
||||
expect(chips.nth(0)).to_contain_text(SEED_SP)
|
||||
expect(chips.nth(1)).to_contain_text(READ1_SP)
|
||||
expect(chips.nth(2)).to_contain_text(READ2_SP)
|
||||
for i, (source, path) in enumerate(EXPECTED_SOURCES):
|
||||
expect(chips.nth(i)).to_have_attribute(
|
||||
"href", f"/document.html?source={source}&path={path}&back=%2F"
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 3. No budget refusal: the listing ran, and the pre-phase-45 refusal
|
||||
# strings are nowhere in the rendered message
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_relist_allowed(
|
||||
page: Page, app_url: str, mock_llm: int, db_ready: None
|
||||
) -> None:
|
||||
page.set_default_timeout(30_000)
|
||||
_reset_db(_seed)
|
||||
page.goto(app_url)
|
||||
_install_page_hooks(page)
|
||||
|
||||
_submit(page, MULTI_QUESTION)
|
||||
_wait_settled(page)
|
||||
|
||||
# The listing tool actually ran (its line rendered, its wire frame
|
||||
# present) — and the turn completed past the point where the old
|
||||
# per-tool budgets would have refused (list budget 1, read budget
|
||||
# 1 — this turn makes one list and TWO reads).
|
||||
frames = _frames(page)
|
||||
assert {"type": "tool", "name": "list_documents", "argument": None} in _tool_frames(
|
||||
frames
|
||||
)
|
||||
line0 = page.locator(".msg.brain .tool-call").nth(0)
|
||||
expect(line0).to_contain_text("Listing documents")
|
||||
|
||||
# No pre-phase-45 budget refusal anywhere in the message — neither
|
||||
# the exact old strings nor the generic fragment — not in the
|
||||
# bubble, not in any tool line.
|
||||
msg_text = page.locator(".msg.brain").last.text_content() or ""
|
||||
for fragment in BUDGET_REFUSAL_FRAGMENTS:
|
||||
assert fragment not in msg_text, (
|
||||
f"budget refusal {fragment!r} rendered: {msg_text!r}"
|
||||
)
|
||||
|
||||
# And it answered (a refusal would have left the model stuck — the
|
||||
# turn settled with a non-deflected, both-named answer).
|
||||
bubble = page.locator(".msg.brain .bubble").last
|
||||
expect(bubble).to_contain_text(BOTH_READS_LINE)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# 4. Phase-37 regression: the single-read marker flow still answers
|
||||
# after exactly ONE read with its single tool pair
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_single_tool_flow_regression(
|
||||
page: Page, app_url: str, mock_llm: int, db_ready: None
|
||||
) -> None:
|
||||
page.set_default_timeout(30_000)
|
||||
_reset_db(_seed)
|
||||
page.goto(app_url)
|
||||
_install_page_hooks(page)
|
||||
|
||||
_submit(page, SINGLE_QUESTION)
|
||||
_wait_settled(page)
|
||||
|
||||
# Exactly TWO tool frames — list then ONE read of the first catalog
|
||||
# line — no second read (the marker carries no multi-read trigger).
|
||||
frames = _frames(page)
|
||||
assert _tool_frames(frames) == [
|
||||
{"type": "tool", "name": "list_documents", "argument": None},
|
||||
{"type": "tool", "name": "read_document", "argument": READ1_SP},
|
||||
]
|
||||
lines = page.locator(".msg.brain .tool-call")
|
||||
expect(lines).to_have_count(2)
|
||||
expect(lines.nth(0)).to_contain_text("Listing documents")
|
||||
expect(lines.nth(1)).to_contain_text("Reading ")
|
||||
expect(lines.nth(1)).to_contain_text(READ1_SP)
|
||||
|
||||
# The single-read answer shape: quotes the read document; it does
|
||||
# NOT carry the multi-read both-named line (READ2 was never read).
|
||||
bubble = page.locator(".msg.brain .bubble").last
|
||||
expect(bubble).to_contain_text(ANSWER_PREFIX)
|
||||
expect(bubble).to_contain_text(ANSWER_QUOTE)
|
||||
expect(bubble).not_to_contain_text(BOTH_READS_LINE)
|
||||
expect(bubble).not_to_contain_text(READ2_SP)
|
||||
|
||||
# done: non-deflected; sources = retrieval doc + the single read
|
||||
# (READ2 absent — it was never read).
|
||||
done = next(f for f in frames if f.get("type") == "done")
|
||||
assert done["deflected"] is False
|
||||
assert [(s["source"], s["path"]) for s in done["sources"]] == [
|
||||
(SEED_SOURCE, SEED_PATH),
|
||||
(READ1_SOURCE, READ1_PATH),
|
||||
]
|
||||
|
||||
row = _last_query_log()
|
||||
assert row.question == SINGLE_QUESTION
|
||||
assert row.deflected is False
|
||||
assert row.sources == f"{SEED_SP}, {READ1_SP}"
|
||||
Reference in New Issue
Block a user