1201 lines
46 KiB
Python
1201 lines
46 KiB
Python
"""Unit: the grounded-turn agent loop (phase 37, ``app.rag.agent``).
|
||
|
||
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 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, the
|
||
``<tools>`` prompt section (HIGH only), and the phase-67 per-round
|
||
retries (a dead-then-recovered round restarts before its first piece
|
||
with a ``RetryPiece``; a mid-stream drop stays terminal — locked A2;
|
||
the forced final no-tools call retries too; ``llm_retries=0`` is one
|
||
plain attempt; retries are invisible to the round cap; consumer abandon
|
||
mid-retry-sleep leaks nothing).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import uuid
|
||
from collections.abc import AsyncGenerator, AsyncIterator
|
||
from copy import deepcopy
|
||
from typing import Any, cast
|
||
|
||
import pytest
|
||
from sqlalchemy.orm import Session
|
||
|
||
from app.config import Settings
|
||
from app.models import Document
|
||
from app.rag import agent
|
||
from app.rag.agent import (
|
||
AGENT_TOOLS,
|
||
AgentHolder,
|
||
run_agent,
|
||
)
|
||
from app.rag.llm import LLMClient, LLMError, RetryPiece, StreamPiece, ToolCallPiece
|
||
from app.rag.prompts import TOOLS_SECTION, _base, build_deflect_prompt, build_high_prompt
|
||
|
||
|
||
def _settings(**kwargs: Any) -> Settings:
|
||
kwargs.setdefault("_env_file", None)
|
||
return Settings(**kwargs) # pyright: ignore[reportCallIssue]
|
||
|
||
|
||
def _doc(source: str, path: str, title: str = "Title", content: str = "CONTENT") -> Document:
|
||
return Document(
|
||
id=uuid.uuid4(),
|
||
source=source,
|
||
path=path,
|
||
full_path=f"/tmp/{path}",
|
||
title=title,
|
||
content=content,
|
||
content_hash="0" * 64,
|
||
)
|
||
|
||
|
||
class ScriptedLLM:
|
||
"""Canned stream sequences; records every ``chat_stream`` request so
|
||
the tests can assert on the messages and the ``tools`` passthrough."""
|
||
|
||
def __init__(self, *streams: list[StreamPiece | ToolCallPiece]) -> None:
|
||
self.streams: list[list[StreamPiece | ToolCallPiece]] = list(streams)
|
||
self.requests: list[tuple[list[dict[str, Any]], list[dict[str, Any]] | None]] = []
|
||
|
||
async def chat_stream(
|
||
self,
|
||
messages: list[dict[str, str]],
|
||
tools: list[dict[str, Any]] | None = None,
|
||
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
|
||
self.requests.append((deepcopy(messages), tools))
|
||
if not self.streams:
|
||
raise AssertionError("ScriptedLLM ran out of canned streams")
|
||
for piece in self.streams.pop(0):
|
||
yield piece
|
||
|
||
|
||
async def _run(
|
||
llm: ScriptedLLM | FailingLLM,
|
||
holder: AgentHolder,
|
||
settings: Settings,
|
||
seed_docs: list[Document] | None = None,
|
||
) -> list[StreamPiece | ToolCallPiece | RetryPiece]:
|
||
out: list[StreamPiece | ToolCallPiece | RetryPiece] = []
|
||
async for piece in run_agent(
|
||
cast("LLMClient", llm),
|
||
cast("Session", None),
|
||
system_prompt="SYSTEM_PROMPT",
|
||
user_message="QUESTION",
|
||
seed_docs=seed_docs or [],
|
||
settings=settings,
|
||
holder=holder,
|
||
):
|
||
out.append(piece)
|
||
return out
|
||
|
||
|
||
# ---------- AGENT_TOOLS shape ----------
|
||
|
||
|
||
def test_agent_tools_names_and_parameters() -> None:
|
||
by_name = {t["function"]["name"]: t for t in AGENT_TOOLS}
|
||
assert len(AGENT_TOOLS) == 3 # list / read / search (phase 68)
|
||
assert set(by_name) == {"list_documents", "read_document", "search_documents"}
|
||
assert all(t["type"] == "function" for t in AGENT_TOOLS)
|
||
list_params = by_name["list_documents"]["function"]["parameters"]
|
||
assert list_params["type"] == "object"
|
||
assert list_params["properties"] == {} # no parameters
|
||
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"
|
||
)
|
||
# Phase 63 (A2): the parameter descriptions point the LLM at the
|
||
# labeled `source:` / `path:` fields of the list_documents output.
|
||
assert read_params["properties"]["source"]["description"] == (
|
||
"The document's source, as shown after 'source: ' in the "
|
||
"list_documents output (e.g. 'Homelab' from "
|
||
"'source: Homelab | path: homelab/aws-route53.md')."
|
||
)
|
||
assert read_params["properties"]["path"]["description"] == (
|
||
"The document's path, as shown after 'path: ' in the "
|
||
"list_documents output (e.g. 'homelab/aws-route53.md' from "
|
||
"'source: Homelab | path: homelab/aws-route53.md')."
|
||
)
|
||
# Phase 68: search_documents — the third tool, a locator (locked A5).
|
||
search = by_name["search_documents"]["function"]
|
||
assert search["description"] == (
|
||
"Search every indexed document for an exact string "
|
||
"(case-insensitive) and return up to 20 matching lines as "
|
||
"'source/path:line: text' — use this to locate content, "
|
||
"then read_document the winner. Optionally pass 'source' "
|
||
"and 'path' (as shown in list_documents) to search one "
|
||
"document only."
|
||
)
|
||
search_params = search["parameters"]
|
||
assert search_params["type"] == "object"
|
||
assert search_params["required"] == ["pattern"]
|
||
assert set(search_params["properties"]) == {"pattern", "source", "path"}
|
||
assert search_params["properties"]["pattern"]["description"] == (
|
||
"The exact text to search for (a plain substring, not a regex)"
|
||
)
|
||
# Phase 63 labeled-field wording, same as read_document's parameters.
|
||
assert search_params["properties"]["source"]["description"] == (
|
||
"The document's source, as shown after 'source: ' in the "
|
||
"list_documents output (e.g. 'Homelab' from "
|
||
"'source: Homelab | path: homelab/aws-route53.md')."
|
||
)
|
||
assert search_params["properties"]["path"]["description"] == (
|
||
"The document's path, as shown after 'path: ' in the "
|
||
"list_documents output (e.g. 'homelab/aws-route53.md' from "
|
||
"'source: Homelab | path: homelab/aws-route53.md')."
|
||
)
|
||
|
||
|
||
# ---------- happy path: list → read → answer ----------
|
||
|
||
|
||
def test_list_then_read_then_answer(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
catalog = [
|
||
("Deployments", "backups.md", "Backup Strategy"),
|
||
("Homelab", "aws-route53.md", "AWS Route53 Records"),
|
||
]
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: catalog)
|
||
target = _doc("Homelab", "aws-route53.md", "AWS Route53 Records", "R53-CONTENT")
|
||
monkeypatch.setattr(agent, "find_document", lambda db, source, path: target)
|
||
seed = [_doc("Homelab", "kubernetes.md", "Kubernetes", "K8S-CONTENT")]
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[ToolCallPiece(id="call_1", name="list_documents", arguments={})],
|
||
[
|
||
ToolCallPiece(
|
||
id="call_2",
|
||
name="read_document",
|
||
arguments={"source": "Homelab", "path": "aws-route53.md"},
|
||
)
|
||
],
|
||
[StreamPiece("thinking", "hmm "), StreamPiece("content", "Done! ")],
|
||
)
|
||
|
||
pieces = asyncio.run(_run(llm, holder, _settings(), seed_docs=seed))
|
||
|
||
# Event order: tool pieces before the answer content/thinking.
|
||
assert [type(p) for p in pieces] == [
|
||
ToolCallPiece,
|
||
ToolCallPiece,
|
||
StreamPiece,
|
||
StreamPiece,
|
||
]
|
||
assert pieces[0] == ToolCallPiece(id="call_1", name="list_documents", arguments={})
|
||
assert isinstance(pieces[1], ToolCallPiece)
|
||
assert pieces[1].name == "read_document"
|
||
assert pieces[3] == StreamPiece("content", "Done! ")
|
||
# The read document is recorded for done.sources / query_log (task 04).
|
||
assert holder.read_docs == [target]
|
||
assert holder.tool_calls == 2
|
||
|
||
# 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
|
||
assert llm.requests[2][1] == AGENT_TOOLS
|
||
assert len(llm.requests) == 3
|
||
|
||
# The follow-up request carries the assistant tool-call + tool result.
|
||
msgs = llm.requests[1][0]
|
||
assert msgs[0] == {"role": "system", "content": "SYSTEM_PROMPT"}
|
||
assert msgs[1] == {"role": "user", "content": "QUESTION"}
|
||
assert msgs[2] == {
|
||
"role": "assistant",
|
||
"content": None,
|
||
"tool_calls": [
|
||
{
|
||
"id": "call_1",
|
||
"type": "function",
|
||
"function": {"name": "list_documents", "arguments": "{}"},
|
||
}
|
||
],
|
||
}
|
||
assert msgs[3] == {
|
||
"role": "tool",
|
||
"tool_call_id": "call_1",
|
||
"content": (
|
||
"2 documents:\n"
|
||
"source: Deployments | path: backups.md | title: Backup Strategy\n"
|
||
"source: Homelab | path: aws-route53.md | title: AWS Route53 Records"
|
||
),
|
||
}
|
||
# The second follow-up request carries the read call + the FULL text.
|
||
msgs = llm.requests[2][0]
|
||
assert msgs[4]["role"] == "assistant"
|
||
assert msgs[4]["tool_calls"][0]["id"] == "call_2"
|
||
assert json.loads(msgs[4]["tool_calls"][0]["function"]["arguments"]) == {
|
||
"source": "Homelab",
|
||
"path": "aws-route53.md",
|
||
}
|
||
assert msgs[5] == {
|
||
"role": "tool",
|
||
"tool_call_id": "call_2",
|
||
"content": "Document Homelab/aws-route53.md:\nR53-CONTENT", # full text, no cap
|
||
}
|
||
|
||
|
||
def test_empty_catalog_listing_says_zero_documents(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: [])
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[ToolCallPiece(id="call_1", name="list_documents", arguments={})],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert llm.requests[1][0][3]["content"] == "0 documents:\n"
|
||
assert holder.tool_calls == 1
|
||
|
||
|
||
def test_content_and_tool_call_in_one_stream_keeps_both(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""Rare stream with content AND a tool call: the content stays (it was
|
||
already emitted) and the tool still runs."""
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: [])
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
StreamPiece("content", "Let me check "),
|
||
ToolCallPiece(id="call_1", name="list_documents", arguments={}),
|
||
],
|
||
[StreamPiece("content", "the answer")],
|
||
)
|
||
pieces = asyncio.run(_run(llm, holder, _settings()))
|
||
assert [type(p) for p in pieces] == [StreamPiece, ToolCallPiece, StreamPiece]
|
||
assert holder.tool_calls == 1 # the tool ran despite the content
|
||
assert llm.requests[1][0][3]["content"] == "0 documents:\n"
|
||
|
||
|
||
# ---------- round cap (phase 45: replaces the per-tool budgets) ----------
|
||
|
||
|
||
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:\nsource: S | path: a.md | title: 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_max_rounds=0)))
|
||
assert [type(p) for p in pieces] == [StreamPiece, StreamPiece]
|
||
assert len(llm.requests) == 1
|
||
assert llm.requests[0][1] is None
|
||
assert llm.requests[0][0] == [
|
||
{"role": "system", "content": "SYSTEM_PROMPT"},
|
||
{"role": "user", "content": "QUESTION"},
|
||
]
|
||
assert holder.read_docs == [] and holder.tool_calls == 0
|
||
|
||
|
||
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"
|
||
"source: Deployments | path: backups.md | title: Backup Strategy\n"
|
||
"source: Homelab | path: aws-route53.md | title: 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.md": a, "b.md": b}[path]
|
||
)
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1", name="read_document", arguments={"source": "S", "path": "a.md"}
|
||
)
|
||
],
|
||
[
|
||
ToolCallPiece(
|
||
id="call_2", name="read_document", arguments={"source": "S", "path": "b.md"}
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
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
|
||
|
||
|
||
# ---------- rejections (non-budget; the round cap bounds their repetition) ----------
|
||
|
||
|
||
def test_reading_a_seed_doc_is_already_in_context(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
seed = [_doc("Homelab", "kubernetes.md", "Kubernetes", "K8S-CONTENT")]
|
||
|
||
def _boom(*_a: Any, **_k: Any) -> None:
|
||
raise AssertionError("find_document must not be called for a seeded doc")
|
||
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: [])
|
||
monkeypatch.setattr(agent, "find_document", _boom)
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1",
|
||
name="read_document",
|
||
arguments={"source": "Homelab", "path": "kubernetes.md"},
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
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
|
||
# Rejected → the tools are still offered on the next request (the
|
||
# round cap is the only bound).
|
||
assert llm.requests[1][1] == AGENT_TOOLS
|
||
|
||
|
||
def test_reading_an_already_read_doc_is_deduped(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
doc = _doc("S", "a.md", "A", "A-CONTENT")
|
||
monkeypatch.setattr(agent, "find_document", lambda db, source, path: doc)
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1", name="read_document", arguments={"source": "S", "path": "a.md"}
|
||
)
|
||
],
|
||
[
|
||
ToolCallPiece(
|
||
id="call_2", name="read_document", arguments={"source": "S", "path": "a.md"}
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
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
|
||
# Rejected → the tools are still offered on the next request…
|
||
assert llm.requests[2][1] == AGENT_TOOLS
|
||
|
||
|
||
def test_unknown_path_refused(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
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"},
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert holder.read_docs == [] and holder.tool_calls == 0
|
||
assert (
|
||
llm.requests[1][0][3]["content"]
|
||
== "No document at S/ghost.md — check the list_documents output."
|
||
)
|
||
assert llm.requests[1][1] == AGENT_TOOLS # tools stay offered (cap bounds)
|
||
|
||
|
||
def test_unknown_tool_name_refused(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: [])
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[ToolCallPiece(id="call_1", name="delete_universe", arguments={"x": 1})],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
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 # rejected → tools stay offered
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("arguments", "label"),
|
||
[
|
||
({}, "no arguments"),
|
||
({"source": "S"}, "path missing"),
|
||
({"path": "p.md"}, "source missing"),
|
||
({"source": "", "path": "p.md"}, "empty source"),
|
||
({"source": "S", "path": " "}, "blank path"),
|
||
({"source": 7, "path": "p.md"}, "non-string source"),
|
||
],
|
||
)
|
||
def test_read_document_missing_arguments_refused(
|
||
monkeypatch: pytest.MonkeyPatch, arguments: dict[str, Any], label: str
|
||
) -> None:
|
||
def _boom(*_a: Any, **_k: Any) -> None:
|
||
raise AssertionError(f"find_document must not be called ({label})")
|
||
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: [])
|
||
monkeypatch.setattr(agent, "find_document", _boom)
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[ToolCallPiece(id="call_1", name="read_document", arguments=arguments)],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert holder.read_docs == [] and holder.tool_calls == 0
|
||
assert llm.requests[1][0][3]["content"] == agent.MISSING_READ_ARGS
|
||
assert llm.requests[1][1] == AGENT_TOOLS
|
||
|
||
|
||
# ---------- search_documents (phase 68, locked A5/A6) ----------
|
||
|
||
|
||
def test_grep_document_case_insensitive_line_numbers() -> None:
|
||
"""Case-insensitive fixed substring, 1-based line numbers, file order,
|
||
repeated matches within a line collapse to one match (grep semantics)."""
|
||
content = "The NEEDLE is here\nno hit\nneedle again\nNEEDLE NEEDLE\n"
|
||
assert agent.grep_document(content, "NEEDLE") == [
|
||
(1, "The NEEDLE is here"),
|
||
(3, "needle again"),
|
||
(4, "NEEDLE NEEDLE"),
|
||
]
|
||
|
||
|
||
def test_grep_document_rstrips_lines_and_empty_content() -> None:
|
||
assert agent.grep_document("hello \t\nworld ", "WORLD") == [(2, "world")]
|
||
assert agent.grep_document("", "x") == []
|
||
assert agent.grep_document("no newlines", "NO") == [(1, "no newlines")]
|
||
assert agent.grep_document("a\nb\n", "MISSING") == []
|
||
|
||
|
||
def test_search_whole_kb_grep_style_output(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""Whole-KB search: catalog order, `source/path:line: text` lines,
|
||
case-insensitive; the call counts in ``tool_calls`` and never touches
|
||
``read_docs``; the tools stay offered on the answer request."""
|
||
d1 = _doc("Alpha", "a/one.md", "One", "first\nNEEDLE in one\nlast")
|
||
d2 = _doc("Beta", "b/two.md", "Two", "no hit\nneedle in two\n")
|
||
monkeypatch.setattr(agent, "all_documents", lambda db: [d1, d2])
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1", name="search_documents", arguments={"pattern": "needle"}
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert llm.requests[1][0][3]["content"] == (
|
||
"Alpha/a/one.md:2: NEEDLE in one\n"
|
||
"Beta/b/two.md:2: needle in two"
|
||
)
|
||
assert holder.tool_calls == 1
|
||
assert holder.read_docs == [] # locked A5: a search adds no context
|
||
assert llm.requests[1][1] == AGENT_TOOLS # tools stay offered
|
||
|
||
|
||
def test_search_capped_at_20_matches_in_catalog_order(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""The 20-match cap is GLOBAL across documents in catalog order, and
|
||
the scan stops once it is hit (a 35-match corpus yields exactly 20)."""
|
||
d1 = _doc("S", "a.md", "A", "\n".join(f"hit-{i}" for i in range(15)))
|
||
d2 = _doc("S", "b.md", "B", "\n".join(f"hit-{i}" for i in range(20)))
|
||
monkeypatch.setattr(agent, "all_documents", lambda db: [d1, d2])
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1", name="search_documents", arguments={"pattern": "hit-"}
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
lines = llm.requests[1][0][3]["content"].split("\n")
|
||
assert len(lines) == agent.SEARCH_MAX_MATCHES
|
||
assert lines[0] == "S/a.md:1: hit-0"
|
||
assert lines[14] == "S/a.md:15: hit-14" # all of a.md
|
||
assert lines[15] == "S/b.md:1: hit-0" # then b.md, in order
|
||
assert lines[19] == "S/b.md:5: hit-4" # cut at the global cap
|
||
assert holder.tool_calls == 1
|
||
|
||
|
||
def test_search_truncates_match_lines_at_200_chars(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""A 300-char match line yields exactly 200 chars of it (no crash)."""
|
||
d1 = _doc("S", "a.md", "A", "top\n" + "x" * 300 + " NEEDLE tail")
|
||
monkeypatch.setattr(agent, "all_documents", lambda db: [d1])
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1", name="search_documents", arguments={"pattern": "needle"}
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert (
|
||
llm.requests[1][0][3]["content"] == f"S/a.md:2: {'x' * agent.SEARCH_LINE_LIMIT}"
|
||
)
|
||
assert holder.tool_calls == 1
|
||
|
||
|
||
def test_search_scoped_to_one_document(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
"""Scoped search: only the named document is loaded (find_document),
|
||
``all_documents`` never runs, and the match line carries its path."""
|
||
d1 = _doc("S", "a.md", "A", "needle here")
|
||
|
||
def _find(db: Any, source: str, path: str) -> Document | None:
|
||
if (source, path) == ("S", "a.md"):
|
||
return d1
|
||
raise AssertionError(
|
||
f"find_document({source}, {path}) — the scoped "
|
||
"search must not load any other document"
|
||
)
|
||
|
||
def _boom(*_a: Any, **_k: Any) -> None:
|
||
raise AssertionError("all_documents must not run for a scoped search")
|
||
|
||
monkeypatch.setattr(agent, "find_document", _find)
|
||
monkeypatch.setattr(agent, "all_documents", _boom)
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1",
|
||
name="search_documents",
|
||
arguments={"pattern": "needle", "source": "S", "path": "a.md"},
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert llm.requests[1][0][3]["content"] == "S/a.md:1: needle here"
|
||
assert holder.tool_calls == 1
|
||
assert holder.read_docs == [] # searched doc did not enter the context
|
||
|
||
|
||
def test_search_scoped_missing_document_refused(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
monkeypatch.setattr(agent, "find_document", lambda db, source, path: None)
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1",
|
||
name="search_documents",
|
||
arguments={"pattern": "x", "source": "S", "path": "ghost.md"},
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert (
|
||
llm.requests[1][0][3]["content"]
|
||
== "No document at S/ghost.md — check the list_documents output."
|
||
)
|
||
assert holder.tool_calls == 0 and holder.read_docs == [] # a refusal
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("arguments", "label"),
|
||
[
|
||
({}, "no arguments"),
|
||
({"pattern": ""}, "empty pattern"),
|
||
({"pattern": " "}, "whitespace pattern"),
|
||
({"pattern": 42}, "non-string pattern"),
|
||
({"pattern": None}, "null pattern"),
|
||
({"pattern": "x", "source": "S"}, "source without path"),
|
||
({"pattern": "x", "path": "a.md"}, "path without source"),
|
||
],
|
||
)
|
||
def test_search_missing_arguments_refused(
|
||
monkeypatch: pytest.MonkeyPatch, arguments: dict[str, Any], label: str
|
||
) -> None:
|
||
"""Unusable pattern OR a half-specified source/path pair → the
|
||
missing-args refusal, with no DB access at all."""
|
||
|
||
def _boom(*_a: Any, **_k: Any) -> None:
|
||
raise AssertionError(f"no DB access for a refused search ({label})")
|
||
|
||
monkeypatch.setattr(agent, "all_documents", _boom)
|
||
monkeypatch.setattr(agent, "find_document", _boom)
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[ToolCallPiece(id="call_1", name="search_documents", arguments=arguments)],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert llm.requests[1][0][3]["content"] == agent.MISSING_SEARCH_ARGS
|
||
assert holder.tool_calls == 0 and holder.read_docs == []
|
||
assert llm.requests[1][1] == AGENT_TOOLS # rejected → tools stay offered
|
||
|
||
|
||
def test_search_no_matches_whole_kb(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
"""Zero hits across the KB → the no-match line (pattern quoted); the
|
||
search still executed, so it counts — and never adds context."""
|
||
monkeypatch.setattr(
|
||
agent, "all_documents", lambda db: [_doc("S", "a.md", "A", "nothing here")]
|
||
)
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1", name="search_documents", arguments={"pattern": "zebra"}
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert llm.requests[1][0][3]["content"] == (
|
||
"No matches for 'zebra' in the knowledge base."
|
||
)
|
||
assert holder.tool_calls == 1
|
||
assert holder.read_docs == []
|
||
|
||
|
||
def test_search_no_matches_scoped(monkeypatch: pytest.MonkeyPatch) -> None:
|
||
doc = _doc("S", "a.md", "A", "nothing here")
|
||
monkeypatch.setattr(agent, "find_document", lambda db, source, path: doc)
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1",
|
||
name="search_documents",
|
||
arguments={"pattern": "zebra", "source": "S", "path": "a.md"},
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert llm.requests[1][0][3]["content"] == "No matches for 'zebra' in S/a.md."
|
||
assert holder.tool_calls == 1
|
||
assert holder.read_docs == []
|
||
|
||
|
||
def test_search_no_match_truncates_long_pattern(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""A pattern longer than 100 chars is truncated in the no-match line
|
||
(kept short); the search itself still runs on the full pattern."""
|
||
monkeypatch.setattr(agent, "all_documents", lambda db: [])
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1",
|
||
name="search_documents",
|
||
arguments={"pattern": "p" * 150},
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert llm.requests[1][0][3]["content"] == (
|
||
f"No matches for '{'p' * 100}' in the knowledge base."
|
||
)
|
||
assert holder.tool_calls == 1
|
||
|
||
|
||
def test_search_counts_but_never_adds_context(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""The locate-then-read workflow: a search finds the document but does
|
||
NOT add it — the subsequent read_document does (and is not rejected as
|
||
already-in-context, because the search touched nothing)."""
|
||
doc = _doc("S", "a.md", "A", "needle here")
|
||
monkeypatch.setattr(agent, "all_documents", lambda db: [doc])
|
||
monkeypatch.setattr(agent, "find_document", lambda db, source, path: doc)
|
||
holder = AgentHolder()
|
||
llm = ScriptedLLM(
|
||
[
|
||
ToolCallPiece(
|
||
id="call_1", name="search_documents", arguments={"pattern": "needle"}
|
||
)
|
||
],
|
||
[
|
||
ToolCallPiece(
|
||
id="call_2",
|
||
name="read_document",
|
||
arguments={"source": "S", "path": "a.md"},
|
||
)
|
||
],
|
||
[StreamPiece("content", "ans")],
|
||
)
|
||
asyncio.run(_run(llm, holder, _settings()))
|
||
assert holder.tool_calls == 2 # search + read, both executed
|
||
assert holder.read_docs == [doc] # only the read added context (A5)
|
||
assert llm.requests[2][0][5]["content"] == "Document S/a.md:\nneedle here"
|
||
|
||
|
||
# ---------- retries inside the agent loop (phase 67, locked A2) ----------
|
||
|
||
|
||
class FailingLLM:
|
||
"""A scripted fake whose Nth ``chat_stream`` call yields pieces and
|
||
then raises (phase 67): ``attempts`` is a list of ``(pieces, error)``
|
||
— an error after zero pieces = "the endpoint died before the first
|
||
token"; after some pieces = a mid-stream drop. Records every
|
||
request's messages/tools and the indices of the attempts whose stream
|
||
teardown ran (``closed``)."""
|
||
|
||
def __init__(
|
||
self,
|
||
attempts: list[tuple[list[StreamPiece | ToolCallPiece], Exception | None]],
|
||
) -> None:
|
||
self.attempts = list(attempts)
|
||
self.requests: list[tuple[list[dict[str, Any]], list[dict[str, Any]] | None]] = []
|
||
#: Indices of attempts whose stream teardown has run.
|
||
self.closed: list[int] = []
|
||
|
||
def chat_stream(
|
||
self,
|
||
messages: list[dict[str, str]],
|
||
tools: list[dict[str, Any]] | None = None,
|
||
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
|
||
index = len(self.requests)
|
||
pieces, error = (
|
||
self.attempts[index]
|
||
if index < len(self.attempts)
|
||
else ([], LLMError("script exhausted"))
|
||
)
|
||
self.requests.append(
|
||
(deepcopy(messages), deepcopy(tools) if tools is not None else None)
|
||
)
|
||
return self._attempt(index, pieces, error)
|
||
|
||
async def _attempt(
|
||
self,
|
||
index: int,
|
||
pieces: list[StreamPiece | ToolCallPiece],
|
||
error: Exception | None,
|
||
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
|
||
try:
|
||
for piece in pieces:
|
||
yield piece
|
||
if error is not None:
|
||
raise error
|
||
finally:
|
||
self.closed.append(index)
|
||
|
||
|
||
def _record_sleeps(monkeypatch: pytest.MonkeyPatch) -> list[float]:
|
||
"""Monkeypatch ``asyncio.sleep`` (what ``chat_stream_retried`` awaits
|
||
for the flat pre-retry delay) and record every awaited delay."""
|
||
sleeps: list[float] = []
|
||
|
||
async def fake_sleep(seconds: float) -> None:
|
||
sleeps.append(seconds)
|
||
|
||
monkeypatch.setattr(asyncio, "sleep", fake_sleep)
|
||
return sleeps
|
||
|
||
|
||
def test_round_retried_before_first_piece(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
caplog: pytest.LogCaptureFixture,
|
||
) -> None:
|
||
"""A tool round that dies before its first piece is restarted with the
|
||
same messages: the stream carries a RetryPiece BEFORE the tool call,
|
||
the tool executes, the final answer streams, and the per-call log line
|
||
is still emitted exactly once (retries are invisible to the loop)."""
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")])
|
||
holder = AgentHolder()
|
||
llm = FailingLLM(
|
||
[
|
||
([], LLMError("connection refused")),
|
||
([ToolCallPiece(id="call_1", name="list_documents", arguments={})], None),
|
||
([StreamPiece("content", "Done!")], None),
|
||
]
|
||
)
|
||
sleeps = _record_sleeps(monkeypatch)
|
||
with caplog.at_level(logging.INFO, logger="app.agent"):
|
||
pieces = asyncio.run(
|
||
_run(llm, holder, _settings(agent_max_rounds=2, llm_retry_delay=2.5))
|
||
)
|
||
assert pieces == [
|
||
RetryPiece(2, 4), # default llm_retries=3 → 4 attempts
|
||
ToolCallPiece(id="call_1", name="list_documents", arguments={}),
|
||
StreamPiece("content", "Done!"),
|
||
]
|
||
assert holder.tool_calls == 1
|
||
assert holder.read_docs == []
|
||
# The restart is byte-identical: same messages, same tools offered.
|
||
assert len(llm.requests) == 3
|
||
assert llm.requests[0] == llm.requests[1]
|
||
assert llm.requests[0][1] == AGENT_TOOLS
|
||
assert llm.requests[2][1] == AGENT_TOOLS # the answer round still offered
|
||
# The flat delay was awaited exactly once, before the retry.
|
||
assert sleeps == [2.5]
|
||
tool_logs = [r for r in caplog.records if r.getMessage().startswith("agent tool=")]
|
||
assert len(tool_logs) == 1 # the retry did not re-run the tool or log
|
||
assert tool_logs[0].getMessage() == "agent tool=list_documents args={} round=1/2"
|
||
|
||
|
||
def test_round_failure_after_first_piece_is_terminal(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""Locked A2: a round that already streamed a piece fails the turn —
|
||
the LLMError propagates out of ``run_agent``, no RetryPiece, no
|
||
sleep, no second request, and the holder is untouched (the tool
|
||
never ran)."""
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: [])
|
||
holder = AgentHolder()
|
||
llm = FailingLLM(
|
||
[([StreamPiece("content", "partial ")], LLMError("mid-stream drop"))]
|
||
)
|
||
sleeps = _record_sleeps(monkeypatch)
|
||
|
||
async def drain() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
|
||
out: list[StreamPiece | ToolCallPiece | RetryPiece] = []
|
||
with pytest.raises(LLMError, match="mid-stream drop"):
|
||
async for piece in run_agent(
|
||
cast("LLMClient", llm),
|
||
cast("Session", None),
|
||
system_prompt="SYSTEM_PROMPT",
|
||
user_message="QUESTION",
|
||
seed_docs=[],
|
||
settings=_settings(),
|
||
holder=holder,
|
||
):
|
||
out.append(piece)
|
||
return out
|
||
|
||
out = asyncio.run(drain())
|
||
assert out == [StreamPiece("content", "partial ")] # no RetryPiece
|
||
assert len(llm.requests) == 1 # no retry
|
||
assert sleeps == []
|
||
assert holder.read_docs == [] and holder.tool_calls == 0
|
||
|
||
|
||
def test_forced_final_no_tools_call_is_retried(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""The forced final request (round cap reached) goes through the same
|
||
retry rule: a failure before its first piece yields a RetryPiece and
|
||
restarts with ``tools=None``; the answer from the retry streams."""
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")])
|
||
holder = AgentHolder()
|
||
llm = FailingLLM(
|
||
[
|
||
([ToolCallPiece(id="call_1", name="list_documents", arguments={})], None),
|
||
([ToolCallPiece(id="call_2", name="list_documents", arguments={})], None),
|
||
([], LLMError("down at the cap")),
|
||
([StreamPiece("content", "forced answer")], None),
|
||
]
|
||
)
|
||
pieces = asyncio.run(_run(llm, holder, _settings(agent_max_rounds=2)))
|
||
assert [type(p) for p in pieces] == [
|
||
ToolCallPiece,
|
||
ToolCallPiece,
|
||
RetryPiece,
|
||
StreamPiece,
|
||
]
|
||
assert pieces[2] == RetryPiece(2, 4)
|
||
assert pieces[3] == StreamPiece("content", "forced answer")
|
||
assert len(llm.requests) == 4 # 2 tool rounds + the final + its retry
|
||
# The forced final (and its retry) carry no tools, whatever is left.
|
||
assert llm.requests[2][1] is None
|
||
assert llm.requests[3][1] is None
|
||
# …and the restart is byte-identical.
|
||
assert llm.requests[2][0] == llm.requests[3][0]
|
||
assert holder.tool_calls == 2
|
||
|
||
|
||
def test_zero_retries_is_one_plain_attempt(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""The kill-switch path (``llm_retries=0``): a dead round raises
|
||
immediately — one request, no RetryPiece, no sleep (pre-phase-67
|
||
behavior)."""
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: [])
|
||
holder = AgentHolder()
|
||
llm = FailingLLM([([], LLMError("connection refused"))])
|
||
sleeps = _record_sleeps(monkeypatch)
|
||
|
||
async def drain() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
|
||
out: list[StreamPiece | ToolCallPiece | RetryPiece] = []
|
||
with pytest.raises(LLMError, match="connection refused"):
|
||
async for piece in run_agent(
|
||
cast("LLMClient", llm),
|
||
cast("Session", None),
|
||
system_prompt="SYSTEM_PROMPT",
|
||
user_message="QUESTION",
|
||
seed_docs=[],
|
||
settings=_settings(llm_retries=0),
|
||
holder=holder,
|
||
):
|
||
out.append(piece)
|
||
return out
|
||
|
||
out = asyncio.run(drain())
|
||
assert out == [] # nothing streamed, no RetryPiece
|
||
assert len(llm.requests) == 1
|
||
assert sleeps == []
|
||
assert holder.read_docs == [] and holder.tool_calls == 0
|
||
|
||
|
||
def test_abandon_mid_retry_sleep_leaks_nothing(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""Consumer abandon while a retried round is parked in the pre-retry
|
||
sleep (client disconnect): the driving task is cancelled cleanly, the
|
||
production teardown ``aclose()`` on ``run_agent`` does not raise, the
|
||
inner attempt's stream was torn down, and the retry never starts."""
|
||
entered = asyncio.Event()
|
||
|
||
async def parking_sleep(seconds: float) -> None:
|
||
entered.set()
|
||
await asyncio.Event().wait() # park until the abandon arrives
|
||
|
||
monkeypatch.setattr(asyncio, "sleep", parking_sleep)
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: [])
|
||
holder = AgentHolder()
|
||
llm = FailingLLM(
|
||
[([], LLMError("endpoint down")), ([StreamPiece("content", "never")], None)]
|
||
)
|
||
|
||
async def run() -> None:
|
||
gen = run_agent(
|
||
cast("LLMClient", llm),
|
||
cast("Session", None),
|
||
system_prompt="SYSTEM_PROMPT",
|
||
user_message="QUESTION",
|
||
seed_docs=[],
|
||
settings=_settings(),
|
||
holder=holder,
|
||
)
|
||
|
||
async def consumer() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
|
||
return [p async for p in gen]
|
||
|
||
task = asyncio.ensure_future(consumer())
|
||
await entered.wait() # the round's retry is parked in the sleep
|
||
assert not task.done()
|
||
task.cancel() # client disconnect: the driving task is cancelled
|
||
with pytest.raises(asyncio.CancelledError):
|
||
await task
|
||
# Production teardown (phase 48 pattern): must not raise. ``run_agent``
|
||
# is an async generator despite its AsyncIterator annotation.
|
||
await cast(
|
||
"AsyncGenerator[StreamPiece | ToolCallPiece | RetryPiece, None]", gen
|
||
).aclose()
|
||
|
||
asyncio.run(run())
|
||
assert len(llm.requests) == 1 # the retry never started
|
||
assert llm.closed == [0] # attempt 1's inner stream was torn down
|
||
assert holder.read_docs == [] and holder.tool_calls == 0
|
||
|
||
|
||
def test_retries_are_invisible_to_the_round_cap(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
caplog: pytest.LogCaptureFixture,
|
||
) -> None:
|
||
"""A failing-then-succeeding round consumes ONE round: with a cap of
|
||
2, the retried first round and the second tool round fill the cap —
|
||
the forced final follows the SECOND call, and the log lines read
|
||
round=1/2 and round=2/2."""
|
||
monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")])
|
||
holder = AgentHolder()
|
||
llm = FailingLLM(
|
||
[
|
||
([], LLMError("down")),
|
||
([ToolCallPiece(id="call_1", name="list_documents", arguments={})], None),
|
||
([ToolCallPiece(id="call_2", name="list_documents", arguments={})], None),
|
||
([StreamPiece("content", "forced answer")], None),
|
||
]
|
||
)
|
||
with caplog.at_level(logging.INFO, logger="app.agent"):
|
||
pieces = asyncio.run(_run(llm, holder, _settings(agent_max_rounds=2)))
|
||
assert [type(p) for p in pieces] == [
|
||
RetryPiece,
|
||
ToolCallPiece,
|
||
ToolCallPiece,
|
||
StreamPiece,
|
||
]
|
||
assert len(llm.requests) == 4 # 2 (round 1 + its retry) + 1 + the forced final
|
||
assert llm.requests[3][1] is None # the forced final, after round 2
|
||
assert holder.tool_calls == 2
|
||
msgs = [r.getMessage() for r in caplog.records]
|
||
assert "agent tool=list_documents args={} round=1/2" in msgs
|
||
assert "agent tool=list_documents args={} round=2/2" in msgs
|
||
assert any("round cap reached (rounds=2)" in m for m in msgs)
|
||
|
||
|
||
# ---------- prompts: <tools> section (HIGH only) ----------
|
||
|
||
|
||
def test_high_prompt_carries_tools_section_after_documents() -> None:
|
||
prompt = build_high_prompt([_doc("S", "a.md", "A", "A-CONTENT")])
|
||
assert TOOLS_SECTION in prompt
|
||
assert "call `list_documents`" in prompt
|
||
assert "then `read_document` to pull in exactly one more document" in prompt
|
||
assert "do not read more than one extra document" in prompt
|
||
# After the mode body: <tools> follows </documents>.
|
||
assert prompt.index("</documents>") < prompt.index("<tools>")
|
||
assert prompt.rstrip().endswith("</tools>")
|
||
|
||
|
||
def test_high_prompt_tools_section_with_notes_and_kb() -> None:
|
||
prompt = build_high_prompt(
|
||
[_doc("S", "a.md", "A", "A-CONTENT")], notes=["be concise"], kb_overview="- KB"
|
||
)
|
||
assert prompt.index("<knowledge_base>") < prompt.index("<tuning>")
|
||
assert prompt.index("<tuning>") < prompt.index("<documents>")
|
||
assert prompt.index("<documents>") < prompt.index("<tools>")
|
||
|
||
|
||
def test_low_prompt_is_byte_identical_and_tool_free() -> None:
|
||
expected = (
|
||
_base("LOW")
|
||
+ "\nDEFLECT_MODE: retrieval was weak — the titles below are the closest "
|
||
"your notes come to the question. They are titles only; do not pretend "
|
||
"they answer it. Use them to propose 2-3 alternative questions.\n"
|
||
+ "- T1\n- T2"
|
||
)
|
||
assert build_deflect_prompt(["T1", "T2"]) == expected
|
||
for prompt in (
|
||
build_deflect_prompt(["T1"]),
|
||
build_deflect_prompt(["T1"], notes=["be concise"]),
|
||
build_deflect_prompt(["T1"], kb_overview="- KB"),
|
||
build_deflect_prompt(["T1"], notes=["be concise"], kb_overview="- KB"),
|
||
):
|
||
assert "<tools>" not in prompt
|
||
assert TOOLS_SECTION not in prompt
|