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

feat(rag): unbounded agent tool calls behind a round cap (owner revision)

Browse Source 
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:
ducoterra
2026-08-28 04:50:56 -04:00
parent bc70ce36e0
commit b855d0aef9
16 changed files with 1311 additions and 278 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/e2e/mock_llm.py
+158 -38
View File
@@ -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")
tests/e2e/test_agent_unlimited_tools.py
+585
View File
@@ -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}"
tests/integration/test_chat_api.py
+20 -19
View File
@@ -78,10 +78,10 @@ class FakeRagLLM:
#: is yielded for the *i*-th request that carries a non-None
#: ``tools`` parameter (a request the agent loop is offering tools
#: on). A request without tools — the deflected direct path, the
#: post-budget answer request, or the 0/0 single-request path —
#: always yields the thinking + answer stream below, so a
#: deflected turn through this fake is byte-identical to the
#: plain fake's output.
#: cap-forced answer request, or the kill-switch
#: (``agent_max_rounds=0``) single-request path — always yields the
#: thinking + answer stream below, so a deflected turn through this
#: fake is byte-identical to the plain fake's output.
self.tool_script: list[list[StreamPiece | ToolCallPiece]] = list(tool_script or [])
async def embed(self, texts: list[str]) -> list[list[float]]:
@@ -503,8 +503,8 @@ def test_grounded_turn_streams_tool_frames_and_cites_read_doc(
``thinking?/tool/tool/delta…/done``; ``done.sources`` and the
``query_log`` row include the read document (deduped, order
preserved); the per-turn log line carries ``tool_calls=2``.
The agent loop offers tools while budgets last and drops them
(``tools=None``) once both are spent."""
Phase 45: the agent loop keeps offering the tools for the whole
turn — the round cap (not per-tool budgets) is the bound."""
scripted = FakeRagLLM(
tool_script=[
[
@@ -518,8 +518,9 @@ def test_grounded_turn_streams_tool_frames_and_cites_read_doc(
arguments={"source": "docs", "path": "homelab/backups.md"},
)
],
# the post-budget answer request (tools=None) falls back to the
# fake's thinking + answer stream
# the answer request still carries the tools (2 rounds < the
# default cap of 10); the fake's tool_script is exhausted, so
# it falls back to the thinking + answer stream
]
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: scripted
@@ -557,12 +558,13 @@ def test_grounded_turn_streams_tool_frames_and_cites_read_doc(
assert len(sources) == len(set(sources)) # deduped by (source, path)
assert done["sources"][-1]["title"] == "Backup Strategy"
# The agent loop offered the tools while any budget remained and
# dropped them once both were spent (single post-budget request).
# Phase 45: the tools stay offered on every request — the round cap
# (not spent budgets) bounds the loop, and the model answered while
# still being offered the tools (2 rounds < default cap 10).
assert len(scripted.seen_messages) == 3
assert scripted.seen_tools[0] == AGENT_TOOLS
assert scripted.seen_tools[1] == AGENT_TOOLS # the read budget was still open
assert scripted.seen_tools[2] is None
assert scripted.seen_tools[1] == AGENT_TOOLS
assert scripted.seen_tools[2] == AGENT_TOOLS
# The query_log row carries the same combined source list.
(row,) = db.scalars(select(QueryLog)).all()
@@ -628,17 +630,17 @@ def test_deflected_turn_stays_byte_identical_without_tools(
assert "backups.md" not in row.sources
def test_zero_agent_budgets_reproduce_pre_phase_single_request(
def test_zero_max_rounds_reproduce_pre_phase_single_request(
client,
db,
seeded_kb: FakeRagLLM,
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None:
"""(c) ``BOR_AGENT_LIST_CALLS=0 BOR_AGENT_READ_CALLS=0``: no ``tool``
frames, exactly one request **without** a ``tools`` key (the
pre-phase request shape), ``done.sources`` unchanged, and
``tool_calls=0`` in the log line — budgets-as-kill-switch."""
"""(c) ``BOR_AGENT_MAX_ROUNDS=0``: no ``tool`` frames, exactly one
request **without** a ``tools`` key (the pre-phase request shape),
``done.sources`` unchanged, and ``tool_calls=0`` in the log line —
the kill switch survives the phase-45 budget removal."""
scripted = FakeRagLLM(
tool_script=[
[ToolCallPiece(id="call_1", name="list_documents", arguments={})],
@@ -658,8 +660,7 @@ def test_zero_agent_budgets_reproduce_pre_phase_single_request(
lambda: Settings(
_env_file=None, # pyright: ignore[reportCallIssue]
relevance_threshold=live.relevance_threshold,
agent_list_calls=0,
agent_read_calls=0,
agent_max_rounds=0,
),
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: scripted
tests/unit/test_agent.py
+155 -103
View File
@@ -3,10 +3,13 @@
A scripted fake LLM (canned stream sequences) + monkeypatched
``list_catalog`` / ``find_document`` — no database, no network. Covers
the loop mechanics: the list → read → answer happy path (event order,
holder state, the ``tools=None`` request after the budgets are spent,
the assistant/tool message history), the 0/0 single-call path, budget
exhaustion, dedupe, unknown tool / missing args / unknown path, the
round cap, and the ``<tools>`` prompt section (HIGH only).
holder state, the tools staying offered on every request — phase 45
removed the per-tool budgets, the assistant/tool message history), the
kill switch (``agent_max_rounds=0`` single-call path), the round cap
forcing a final no-tools answer (an always-calling stream and an
always-rejected stream), re-lists and multi-reads executing without
budgets, dedupe, unknown tool / missing args / unknown path, and the
``<tools>`` prompt section (HIGH only).
"""
from __future__ import annotations
@@ -102,6 +105,11 @@ def test_agent_tools_names_and_parameters() -> None:
read_params = by_name["read_document"]["function"]["parameters"]
assert read_params["required"] == ["source", "path"]
assert set(read_params["properties"]) == {"source", "path"}
# Phase 45: the per-tool budgets are gone — "exactly one more"
# dropped out of the read_document description.
assert by_name["read_document"]["function"]["description"] == (
"Add the full content of one more indexed document to your context"
)
# ---------- happy path: list → read → answer ----------
@@ -148,11 +156,12 @@ def test_list_then_read_then_answer(
assert holder.read_docs == [target]
assert holder.tool_calls == 2
# Default budgets (1/1): tools offered while any budget remains…
# Phase 45: no per-tool budgets — the tools stay offered on every
# request (the round cap, not spent budgets, bounds the loop), so
# the answer request still carries them (2 rounds < default cap 10).
assert llm.requests[0][1] == AGENT_TOOLS
assert llm.requests[1][1] == AGENT_TOOLS
# …and dropped (tools=None) once both are spent.
assert llm.requests[2][1] is None
assert llm.requests[2][1] == AGENT_TOOLS
assert len(llm.requests) == 3
# The follow-up request carries the assistant tool-call + tool result.
@@ -226,17 +235,53 @@ def test_content_and_tool_call_in_one_stream_keeps_both(
assert llm.requests[1][0][3]["content"] == "0 documents:\n"
# ---------- budgets ----------
# ---------- round cap (phase 45: replaces the per-tool budgets) ----------
def test_zero_budgets_is_one_request_without_tools() -> None:
"""BOR_AGENT_LIST_CALLS=0 BOR_AGENT_READ_CALLS=0 → byte-identical
single-call path: exactly one request, tools=None, no history growth."""
def test_always_list_bounded_by_round_cap(monkeypatch: pytest.MonkeyPatch) -> None:
"""A model that keeps calling ``list_documents`` gets exactly
``agent_max_rounds`` tool rounds, then one forced ``tools=None``
request streams the answer — the cap is the only forced exit."""
monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")])
listing = "1 documents:\nS/a.md — A"
holder = AgentHolder()
llm = ScriptedLLM(
[ToolCallPiece(id="call_1", name="list_documents", arguments={})],
[ToolCallPiece(id="call_2", name="list_documents", arguments={})],
[ToolCallPiece(id="call_3", name="list_documents", arguments={})],
[StreamPiece("content", "forced answer")],
)
pieces = asyncio.run(_run(llm, holder, _settings(agent_max_rounds=3)))
assert [type(p) for p in pieces] == [
ToolCallPiece,
ToolCallPiece,
ToolCallPiece,
StreamPiece,
]
assert len(llm.requests) == 4 # 3 tool rounds + the forced answer
# The three tool rounds were offered the tools…
assert llm.requests[0][1] == AGENT_TOOLS
assert llm.requests[1][1] == AGENT_TOOLS
assert llm.requests[2][1] == AGENT_TOOLS
# …and the forced final request carries no tools, whatever is left.
assert llm.requests[3][1] is None
# Every re-list executed and counted.
assert holder.tool_calls == 3
# The final request carries all three executed listings as history.
final_msgs = llm.requests[3][0]
assert len(final_msgs) == 8 # 2 + 3 rounds × (assistant + tool)
assert final_msgs[3]["content"] == listing
assert final_msgs[5]["content"] == listing
assert final_msgs[7]["content"] == listing
def test_zero_max_rounds_is_one_request_without_tools() -> None:
"""``agent_max_rounds=0`` — the kill switch: exactly one request,
``tools=None``, no tool lines, no history growth (byte-identical to
the pre-phase-37 path)."""
holder = AgentHolder()
llm = ScriptedLLM([StreamPiece("thinking", "t "), StreamPiece("content", "direct answer")])
pieces = asyncio.run(
_run(llm, holder, _settings(agent_list_calls=0, agent_read_calls=0))
)
pieces = asyncio.run(_run(llm, holder, _settings(agent_max_rounds=0)))
assert [type(p) for p in pieces] == [StreamPiece, StreamPiece]
assert len(llm.requests) == 1
assert llm.requests[0][1] is None
@@ -247,12 +292,92 @@ def test_zero_budgets_is_one_request_without_tools() -> None:
assert holder.read_docs == [] and holder.tool_calls == 0
def test_read_budget_exhausted_refuses_and_appends_nothing(
def test_rejected_read_spam_runs_to_round_cap(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Every call rejected (unknown path — "No document at …"): rejections
no longer end the loop early via budgets — the round cap bounds them
and forces the final no-tools answer."""
monkeypatch.setattr(agent, "find_document", lambda db, source, path: None)
holder = AgentHolder()
llm = ScriptedLLM(
[
ToolCallPiece(
id="call_1",
name="read_document",
arguments={"source": "S", "path": "ghost.md"},
)
],
[
ToolCallPiece(
id="call_2",
name="read_document",
arguments={"source": "S", "path": "ghost.md"},
)
],
[
ToolCallPiece(
id="call_3",
name="read_document",
arguments={"source": "S", "path": "ghost.md"},
)
],
[StreamPiece("content", "forced answer")],
)
asyncio.run(_run(llm, holder, _settings(agent_max_rounds=3)))
assert len(llm.requests) == 4 # 3 rejected rounds + the forced answer
assert llm.requests[0][1] == AGENT_TOOLS
assert llm.requests[1][1] == AGENT_TOOLS
assert llm.requests[2][1] == AGENT_TOOLS
assert llm.requests[3][1] is None # the forced final request: no tools
assert holder.read_docs == [] and holder.tool_calls == 0 # nothing executed
refusal = "No document at S/ghost.md — check the list_documents output."
assert llm.requests[1][0][3]["content"] == refusal
assert llm.requests[2][0][5]["content"] == refusal
assert llm.requests[3][0][7]["content"] == refusal
# ---------- unlimited calls: re-lists and multi-reads (phase 45) ----------
def test_relist_executes_and_counts(monkeypatch: pytest.MonkeyPatch) -> None:
"""Re-lists execute — a second ``list_documents`` in one turn returns
the catalog again and counts in ``tool_calls`` (no budget to
exhaust)."""
catalog = [
("Deployments", "backups.md", "Backup Strategy"),
("Homelab", "aws-route53.md", "AWS Route53 Records"),
]
monkeypatch.setattr(agent, "list_catalog", lambda db: catalog)
holder = AgentHolder()
llm = ScriptedLLM(
[ToolCallPiece(id="call_1", name="list_documents", arguments={})],
[ToolCallPiece(id="call_2", name="list_documents", arguments={})],
[StreamPiece("content", "ans")],
)
asyncio.run(_run(llm, holder, _settings()))
assert holder.tool_calls == 2 # both re-lists executed and counted
listing = (
"2 documents:\n"
"Deployments/backups.md — Backup Strategy\n"
"Homelab/aws-route53.md — AWS Route53 Records"
)
# The answer request carries the catalog a second time as a tool result.
assert llm.requests[2][0][3]["content"] == listing # first listing
assert llm.requests[2][0][5]["content"] == listing # the re-list
assert llm.requests[2][1] == AGENT_TOOLS # still offered (no budgets)
def test_multi_read_executes_without_budgets(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Reads are no longer budgeted either — two different documents can
be read in one turn (re-reading the same one is still deduped via
ALREADY_IN_CONTEXT — see the rejection tests)."""
a = _doc("S", "a.md", "A", "A-CONTENT")
b = _doc("S", "b.md", "B", "B-CONTENT")
monkeypatch.setattr(
agent, "find_document", lambda db, source, path: a if path == "a.md" else None
agent, "find_document", lambda db, source, path: {"a.md": a, "b.md": b}[path]
)
holder = AgentHolder()
llm = ScriptedLLM(
@@ -268,38 +393,15 @@ def test_read_budget_exhausted_refuses_and_appends_nothing(
],
[StreamPiece("content", "ans")],
)
asyncio.run(_run(llm, holder, _settings(agent_list_calls=1, agent_read_calls=1)))
assert holder.read_docs == [a] # the refused read appended nothing
assert holder.tool_calls == 1 # …and consumed no budget
refusal = llm.requests[2][0][5]
assert refusal == {
"role": "tool",
"tool_call_id": "call_2",
"content": agent.READ_EXHAUSTED,
}
# The list budget is still open, so tools stay offered after the refusal.
assert llm.requests[2][1] == AGENT_TOOLS
asyncio.run(_run(llm, holder, _settings()))
assert holder.read_docs == [a, b] # both reads appended, in order
assert holder.tool_calls == 2
assert llm.requests[1][0][3]["content"] == "Document S/a.md:\nA-CONTENT"
assert llm.requests[2][0][5]["content"] == "Document S/b.md:\nB-CONTENT"
assert llm.requests[2][1] == AGENT_TOOLS # the second read was still offered
def test_list_budget_exhausted_refuses_with_its_own_message(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(agent, "list_catalog", lambda db: [])
holder = AgentHolder()
llm = ScriptedLLM(
[ToolCallPiece(id="call_1", name="list_documents", arguments={})],
[ToolCallPiece(id="call_2", name="list_documents", arguments={})],
[StreamPiece("content", "ans")],
)
asyncio.run(_run(llm, holder, _settings(agent_list_calls=1, agent_read_calls=1)))
assert holder.tool_calls == 1
assert llm.requests[2][0][5]["content"] == agent.LIST_EXHAUSTED
# The read budget is still open, so tools stay offered after the refusal.
assert llm.requests[2][1] == AGENT_TOOLS
# ---------- rejections (no budget consumed) ----------
# ---------- rejections (non-budget; the round cap bounds their repetition) ----------
def test_reading_a_seed_doc_is_already_in_context(
@@ -326,7 +428,8 @@ def test_reading_a_seed_doc_is_already_in_context(
asyncio.run(_run(llm, holder, _settings(), seed_docs=seed))
assert holder.read_docs == [] and holder.tool_calls == 0
assert llm.requests[1][0][3]["content"] == agent.ALREADY_IN_CONTEXT
# No budget consumed → tools are still offered on the next request.
# Rejected → the tools are still offered on the next request (the
# round cap is the only bound).
assert llm.requests[1][1] == AGENT_TOOLS
@@ -349,15 +452,15 @@ def test_reading_an_already_read_doc_is_deduped(
],
[StreamPiece("content", "ans")],
)
asyncio.run(_run(llm, holder, _settings(agent_list_calls=1, agent_read_calls=1)))
asyncio.run(_run(llm, holder, _settings()))
assert holder.read_docs == [doc] # appended exactly once
assert holder.tool_calls == 1
assert llm.requests[2][0][5]["content"] == agent.ALREADY_IN_CONTEXT
# The read budget is intact after the deduped refusal…
# Rejected → the tools are still offered on the next request…
assert llm.requests[2][1] == AGENT_TOOLS
def test_unknown_path_refused_without_budget(
def test_unknown_path_refused(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(agent, "find_document", lambda db, source, path: None)
@@ -378,7 +481,7 @@ def test_unknown_path_refused_without_budget(
llm.requests[1][0][3]["content"]
== "No document at S/ghost.md — check the list_documents output."
)
assert llm.requests[1][1] == AGENT_TOOLS # budget intact
assert llm.requests[1][1] == AGENT_TOOLS # tools stay offered (cap bounds)
def test_unknown_tool_name_refused(
@@ -393,7 +496,7 @@ def test_unknown_tool_name_refused(
asyncio.run(_run(llm, holder, _settings()))
assert holder.read_docs == [] and holder.tool_calls == 0
assert llm.requests[1][0][3]["content"] == agent.UNKNOWN_TOOL
assert llm.requests[1][1] == AGENT_TOOLS # nothing was consumed
assert llm.requests[1][1] == AGENT_TOOLS # rejected → tools stay offered
@pytest.mark.parametrize(
@@ -426,57 +529,6 @@ def test_read_document_missing_arguments_refused(
assert llm.requests[1][1] == AGENT_TOOLS
# ---------- round cap (pathological stream) ----------
def test_round_cap_forces_a_final_no_tools_answer(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A model that keeps calling a budget-exhausted tool must be forced
to answer at ``max_rounds = 2 + list + read`` (= 4 for 1/1)."""
monkeypatch.setattr(agent, "list_catalog", lambda db: [])
holder = AgentHolder()
llm = ScriptedLLM(
[ToolCallPiece(id="call_1", name="list_documents", arguments={})],
[ToolCallPiece(id="call_2", name="list_documents", arguments={})],
[ToolCallPiece(id="call_3", name="list_documents", arguments={})],
[ToolCallPiece(id="call_4", name="list_documents", arguments={})],
[StreamPiece("content", "forced answer")],
)
pieces = asyncio.run(_run(llm, holder, _settings(agent_list_calls=1, agent_read_calls=1)))
assert [type(p) for p in pieces] == [
ToolCallPiece,
ToolCallPiece,
ToolCallPiece,
ToolCallPiece,
StreamPiece,
]
assert len(llm.requests) == 5
# The forced final request carries no tools, whatever is left.
assert llm.requests[4][1] is None
# Only the first call consumed budget; the three rejections did not.
assert holder.tool_calls == 1
# The 4th rejection sits at messages[2 + 4*2 - 1] of the final request.
assert llm.requests[4][0][9]["content"] == agent.LIST_EXHAUSTED
# ---------- settings ----------
def test_agent_budget_settings_default_to_one_each() -> None:
s = _settings()
assert s.agent_list_calls == 1
assert s.agent_read_calls == 1
def test_agent_budget_settings_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("BOR_AGENT_LIST_CALLS", "0")
monkeypatch.setenv("BOR_AGENT_READ_CALLS", "2")
s = _settings()
assert s.agent_list_calls == 0
assert s.agent_read_calls == 2
# ---------- prompts: <tools> section (HIGH only) ----------
tests/unit/test_chat_gate.py
+4 -3
View File
@@ -563,9 +563,10 @@ def test_endpoint_grounded_turn_runs_agent_loop_with_tools(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Phase 37: a grounded endpoint turn runs the agent loop — the
single no-tool-call request carries ``AGENT_TOOLS`` (default 1/1
budgets), no ``tool`` frames stream, and the ``done`` event is the
plain retrieval shape (the tool-free answer is byte-identical)."""
single no-tool-call request carries ``AGENT_TOOLS`` (the default
round cap keeps the tools offered), no ``tool`` frames stream, and
the ``done`` event is the plain retrieval shape (the tool-free
answer is byte-identical)."""
_session, llm = gate_env
doc = _doc("Kubernetes Homelab Cluster", "TALOS_DOC_SENT")
monkeypatch.setattr(chat_api, "retrieve", _fake_retriever([_chunk(doc, 0.90)]))
tests/unit/test_config.py
+20
View File
@@ -78,6 +78,26 @@ def test_max_output_tokens_env_override(monkeypatch) -> None:
assert s.max_output_tokens == 1234
def test_agent_max_rounds_default_and_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
"""Phase 45: the per-tool budgets are gone — ``BOR_AGENT_MAX_ROUNDS``
(default 10) is the single agent-loop knob; ``0`` is the no-tools
kill switch."""
monkeypatch.delenv("BOR_AGENT_MAX_ROUNDS", raising=False)
assert _settings().agent_max_rounds == 10
monkeypatch.setenv("BOR_AGENT_MAX_ROUNDS", "5")
assert _settings().agent_max_rounds == 5
monkeypatch.setenv("BOR_AGENT_MAX_ROUNDS", "0")
assert _settings().agent_max_rounds == 0
def test_agent_max_rounds_rejects_negative(monkeypatch: pytest.MonkeyPatch) -> None:
"""``0`` is the kill switch — a negative value is a typo, so the
validator fails loudly at startup."""
monkeypatch.setenv("BOR_AGENT_MAX_ROUNDS", "-1")
with pytest.raises(ValidationError, match="agent_max_rounds"):
_settings()
def test_stream_thinking_default_true_and_env_parse(monkeypatch: pytest.MonkeyPatch) -> None:
"""Phase 17 kill-switch (``BOR_STREAM_THINKING``): on by default,
``0``/``false`` turn the ``thinking`` SSE frames off."""
tests/unit/test_mock_tool_flow.py
+241
View File
@@ -0,0 +1,241 @@
"""Unit tests for the E2E mock's tool-flow classifier (phase 45, task 02).
The mock (``tests/e2e/mock_llm.py``) classifies marker requests
statelessly into one step of the agent tool flow. This file pins the
classification at unit speed — no Playwright, no LLM process:
* the phase-37 SINGLE-READ flow (``TOOLS_TRIGGER`` only) stays
byte-identical: list → read (first catalog line, ``call_1``) → answer;
* the phase-45 MULTI-READ flow (``TOOLS_TRIGGER`` + ``MULTI_READ_TRIGGER``)
classifies by the count of ``tool``-role read results: list → read #1
(``call_1``) → read #2 (second catalog line, ``call_2``) → the
byte-stable ``multi_answer`` naming both read paths.
"""
from __future__ import annotations
from typing import Any
from tests.e2e.mock_llm import (
MULTI_READ_TRIGGER,
TOOLS_TRIGGER,
_tool_flow,
)
# --------------------------------------------------------------------------
# Wire fixtures — byte-identical to what app/rag/agent.py produces
# --------------------------------------------------------------------------
#: The ``<tools>`` section marks the HIGH prompt (app/rag/prompts.py).
SYSTEM_HIGH = "<relevance>HIGH</relevance>\n<documents>\n</documents>\n<tools>\n…\n</tools>"
SYSTEM_LOW = "<relevance>LOW</relevance>\n"
#: A minimal truthy ``tools`` parameter (the mock only checks presence).
TOOLS = [{"type": "function", "function": {"name": "list_documents"}}]
#: The agent's ``list_documents`` output for a two-document KB
# (``app/rag/agent.py`` ``_execute_tool``): one ``source/path — title``
#: line per document, ``(source, path)`` order.
CATALOG_2 = (
"2 documents:\n"
"Deployments/example-record-file.json — Example Record File\n"
"Homelab/aws-route53.md — AWS Route 53 Notes"
)
CATALOG_1 = "1 documents:\nDeployments/example-record-file.json — Example Record File"
CATALOG_3 = (
"3 documents:\n"
"Deployments/aaa.md — AAA\n"
"Deployments/bbb.md — BBB\n"
"Homelab/ccc.md — CCC"
)
DOC1_SP = "Deployments/example-record-file.json"
DOC1_CONTENT = (
"The record file keeps every hosted zone record — first line is longer "
"than eighty characters so the quote truncation below is observable.\n"
"second line of the document content"
)
assert len(DOC1_CONTENT) > 80
DOC2_SP = "Homelab/aws-route53.md"
DOC2_CONTENT = "Route 53 notes — the second read, short on purpose."
SINGLE_USER = "Use your tools: what is the exact shape of the record file?"
#: Carries BOTH markers — ``use your tools`` then ``read two documents``.
MULTI_USER = "Use your tools and read two documents: compare the zone notes with the record file."
#: The multi marker alone — no ``use your tools``.
MULTI_ONLY_USER = "Please read two documents and compare them."
PLAIN_USER = "How does the sync job push records to the zone?"
assert TOOLS_TRIGGER in SINGLE_USER.lower() and MULTI_READ_TRIGGER not in SINGLE_USER.lower()
assert TOOLS_TRIGGER in MULTI_USER.lower() and MULTI_READ_TRIGGER in MULTI_USER.lower()
def _read_result(sp: str, content: str) -> str:
"""The agent's read-result text (``_execute_tool`` prefix)."""
return f"Document {sp}:\n{content}"
def _body(
user: str,
tool_msgs: tuple[str, ...] = (),
tools: Any = TOOLS,
system: str = SYSTEM_HIGH,
) -> dict[str, Any]:
"""A chat-completion body: system + user + the tool results in order."""
messages: list[dict[str, Any]] = [
{"role": "system", "content": system},
{"role": "user", "content": user},
]
for i, content in enumerate(tool_msgs):
messages.append(
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": f"call_{i}",
"type": "function",
"function": {"name": "list_documents", "arguments": "{}"},
}
],
}
)
messages.append({"role": "tool", "tool_call_id": f"call_{i}", "content": content})
return {"messages": messages, "tools": tools}
# --------------------------------------------------------------------------
# Phase-37 single-read flow — must stay byte-identical
# --------------------------------------------------------------------------
def test_single_flow_list_step() -> None:
assert _tool_flow(_body(SINGLE_USER)) == ("list", "", "")
def test_single_flow_read_step_first_catalog_line() -> None:
flow = _tool_flow(_body(SINGLE_USER, (CATALOG_3,)))
# The FIRST listing line (Deployments/aaa.md), rsplit convention.
assert flow == ("read", "Deployments", "aaa.md", "call_1")
def test_single_flow_answer_step_with_tools_offered() -> None:
# Phase 45: the round cap keeps the tools offered until it is hit —
# the answer step fires regardless of the ``tools`` parameter.
flow = _tool_flow(
_body(SINGLE_USER, (CATALOG_2, _read_result(DOC1_SP, DOC1_CONTENT)))
)
assert flow == ("answer", DOC1_SP, DOC1_CONTENT)
def test_single_flow_answer_step_without_tools() -> None:
flow = _tool_flow(
_body(
SINGLE_USER,
(CATALOG_2, _read_result(DOC1_SP, DOC1_CONTENT)),
tools=None,
)
)
assert flow == ("answer", DOC1_SP, DOC1_CONTENT)
def test_single_flow_no_tools_no_results_is_not_the_flow() -> None:
# agent_max_rounds=0 path: marker + <tools> prompt, but the request
# carries no tools and no tool results — regular answer, not a flow.
assert _tool_flow(_body(SINGLE_USER, tools=None)) is None
def test_single_flow_marker_without_tools_section_is_none() -> None:
assert _tool_flow(_body(SINGLE_USER, system=SYSTEM_LOW)) is None
def test_single_flow_plain_question_is_none() -> None:
assert _tool_flow(_body(PLAIN_USER)) is None
# --------------------------------------------------------------------------
# Phase-45 multi-read flow (task 02)
# --------------------------------------------------------------------------
def test_multi_flow_list_step() -> None:
assert _tool_flow(_body(MULTI_USER)) == ("list", "", "")
def test_multi_flow_read_first_step() -> None:
flow = _tool_flow(_body(MULTI_USER, (CATALOG_2,)))
assert flow == ("read", DOC1_SP.split("/", 1)[0], DOC1_SP.rsplit("/", 1)[1], "call_1")
def test_multi_flow_read_second_step_skips_already_read() -> None:
flow = _tool_flow(_body(MULTI_USER, (CATALOG_2, _read_result(DOC1_SP, DOC1_CONTENT))))
# The second catalog line — the first line differing from DOC1.
assert flow == ("read", "Homelab", "aws-route53.md", "call_2")
def test_multi_flow_read_second_is_listing_order_not_last() -> None:
# Three-doc catalog, first doc read: read #2 is the SECOND line
# (Deployments/bbb.md), not the last one.
flow = _tool_flow(_body(MULTI_USER, (CATALOG_3, _read_result("Deployments/aaa.md", "x"))))
assert flow == ("read", "Deployments", "bbb.md", "call_2")
def test_multi_flow_answer_step_names_both_paths() -> None:
flow = _tool_flow(
_body(
MULTI_USER,
(
CATALOG_2,
_read_result(DOC1_SP, DOC1_CONTENT),
_read_result(DOC2_SP, DOC2_CONTENT),
),
)
)
assert flow is not None
assert flow[0] == "multi_answer"
# Byte-stable: the single-read shape quoting the FIRST read result
# (first 80 chars), plus both read paths in read order.
assert flow[2] == f"Read {DOC1_SP}. {DOC1_CONTENT[:80]} I read {DOC1_SP} and {DOC2_SP}."
def test_multi_flow_answer_step_without_tools_offered() -> None:
# The forced answer is content, not a tool call — it must not be
# gated on the ``tools`` parameter.
flow = _tool_flow(
_body(
MULTI_USER,
(
CATALOG_2,
_read_result(DOC1_SP, DOC1_CONTENT),
_read_result(DOC2_SP, DOC2_CONTENT),
),
tools=None,
)
)
assert flow is not None
assert flow[0] == "multi_answer"
def test_multi_flow_one_document_catalog_degenerates_to_single_answer() -> None:
# Nothing second to read — the single-read answer shape, quoting the
# only read result.
flow = _tool_flow(
_body(MULTI_USER, (CATALOG_1, _read_result(DOC1_SP, DOC1_CONTENT)))
)
assert flow == ("answer", DOC1_SP, DOC1_CONTENT)
def test_multi_flow_no_tools_no_results_is_not_the_flow() -> None:
assert _tool_flow(_body(MULTI_USER, tools=None)) is None
def test_multi_trigger_without_tools_trigger_is_none() -> None:
# The multi marker alone (no ``use your tools``) is not the flow.
assert MULTI_READ_TRIGGER in MULTI_ONLY_USER
assert TOOLS_TRIGGER not in MULTI_ONLY_USER.lower()
assert _tool_flow(_body(MULTI_ONLY_USER)) is None
def test_multi_flow_requires_tools_section() -> None:
assert _tool_flow(_body(MULTI_USER, system=SYSTEM_LOW)) is None