11b1bfbc94
A request can carry several of this capability's returns — a model can call search twice in one response — and the carrier was identified by message alone, so each of them received the whole capsule. It is selected by message and part now, so exactly one carries it however many share the request. The chat TUI passed its persisted state into the run, so tool synchronisation mutated it in place while the message history was promoted only on success. A cancelled or failed run therefore kept the evidence the tools had recorded and discarded the messages that justified it, and the next question derived its identity from the shorter history: behind the recorded epoch, refused as non-append-only, the conversation unusable until cleared. The run gets a copy, promoted with the messages or not at all. Five decorators had been left attached to a helper by an earlier extraction, which pytest does not collect, so the resume case they carried was silently untested. The wire test covers both resume shapes again, no prompt and deferred results.
729 lines
24 KiB
Python
729 lines
24 KiB
Python
import base64
|
|
from dataclasses import dataclass, field, replace
|
|
from typing import Any, cast
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
from pydantic_ai import Agent, DeferredToolResults, RunContext
|
|
from pydantic_ai.messages import (
|
|
BinaryContent,
|
|
ModelRequest,
|
|
ModelResponse,
|
|
TextPart,
|
|
ToolCallPart,
|
|
ToolReturnPart,
|
|
UserPromptPart,
|
|
)
|
|
from pydantic_ai.models import ModelRequestContext, ModelRequestParameters
|
|
from pydantic_ai.models.function import FunctionModel
|
|
from pydantic_ai.models.test import TestModel
|
|
from pydantic_ai.usage import RunUsage
|
|
|
|
from haiku.rag.capabilities.compaction import (
|
|
RECEIPT,
|
|
Capsule,
|
|
compact_history,
|
|
picture_label,
|
|
)
|
|
from haiku.rag.capabilities.compaction import create_capability as create_compaction
|
|
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
|
|
from haiku.rag.capabilities.rag import RAGCapability, RAGState
|
|
from haiku.rag.capabilities.rag import create_capability as create_rag
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.store.models.chunk import SearchResult
|
|
|
|
OWNED = frozenset({"rag_search"})
|
|
PNG = BinaryContent(data=b"fake-image-bytes", media_type="image/png")
|
|
|
|
|
|
def retrieved_image(chunk_id: str = "chunk-1", self_ref: str = "#/pictures/0"):
|
|
"""A page image on the wire, labelled the way a search result attaches it."""
|
|
return UserPromptPart(content=[picture_label(chunk_id, self_ref), PNG])
|
|
|
|
|
|
def answered_question(question: str, *, evidence: str, images: bool = False):
|
|
"""One settled question: prompt, search, result, answer."""
|
|
returned: list[Any] = [ToolReturnPart("rag_search", evidence, "call-1")]
|
|
if images:
|
|
returned.append(retrieved_image())
|
|
return [
|
|
ModelRequest(parts=[UserPromptPart(question)]),
|
|
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
|
ModelRequest(parts=returned),
|
|
ModelResponse(parts=[TextPart("an answer")]),
|
|
]
|
|
|
|
|
|
def returns_of(messages) -> list[str]:
|
|
return [
|
|
str(part.content)
|
|
for message in messages
|
|
if isinstance(message, ModelRequest)
|
|
for part in message.parts
|
|
if isinstance(part, ToolReturnPart)
|
|
]
|
|
|
|
|
|
def images_of(messages) -> list[BinaryContent]:
|
|
return [
|
|
item
|
|
for message in messages
|
|
if isinstance(message, ModelRequest)
|
|
for part in message.parts
|
|
if isinstance(part, UserPromptPart) and not isinstance(part.content, str)
|
|
for item in part.content
|
|
if isinstance(item, BinaryContent)
|
|
]
|
|
|
|
|
|
def texts_of(messages) -> list[str]:
|
|
return [
|
|
item
|
|
for message in messages
|
|
if isinstance(message, ModelRequest)
|
|
for part in message.parts
|
|
if isinstance(part, UserPromptPart) and not isinstance(part.content, str)
|
|
for item in part.content
|
|
if isinstance(item, str)
|
|
]
|
|
|
|
|
|
def test_nothing_before_the_first_question_is_compacted():
|
|
messages = answered_question("first", evidence="EVIDENCE")
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=0, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert compacted == messages
|
|
|
|
|
|
def test_the_newest_earlier_return_carries_the_capsule():
|
|
messages = [
|
|
*answered_question("first", evidence="OLD EVIDENCE"),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert returns_of(compacted) == ["CAPSULE"]
|
|
assert "OLD EVIDENCE" not in returns_of(compacted)
|
|
|
|
|
|
def test_older_returns_become_receipts_and_only_the_newest_carries_the_capsule():
|
|
messages = [
|
|
*answered_question("first", evidence="OLDEST"),
|
|
*answered_question("second", evidence="NEWER"),
|
|
ModelRequest(parts=[UserPromptPart("third")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=8, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert returns_of(compacted) == [RECEIPT, "CAPSULE"]
|
|
|
|
|
|
def test_the_current_question_keeps_its_own_evidence():
|
|
messages = [
|
|
*answered_question("first", evidence="OLD EVIDENCE"),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-2")]),
|
|
ModelRequest(parts=[ToolReturnPart("rag_search", "LIVE EVIDENCE", "call-2")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert returns_of(compacted) == ["CAPSULE", "LIVE EVIDENCE"]
|
|
|
|
|
|
def test_another_capabilitys_return_is_left_alone():
|
|
messages = [
|
|
ModelRequest(parts=[UserPromptPart("first")]),
|
|
ModelResponse(parts=[ToolCallPart("other_tool", {}, "call-1")]),
|
|
ModelRequest(parts=[ToolReturnPart("other_tool", "NOT OURS", "call-1")]),
|
|
ModelResponse(parts=[TextPart("an answer")]),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert returns_of(compacted) == ["NOT OURS"]
|
|
|
|
|
|
def test_a_cite_acknowledgement_survives():
|
|
"""A receipt of the model's own action, not evidence."""
|
|
messages = [
|
|
ModelRequest(parts=[UserPromptPart("first")]),
|
|
ModelResponse(
|
|
parts=[ToolCallPart("rag_cite", {"chunk_ids": ["c1"]}, "call-1")]
|
|
),
|
|
ModelRequest(
|
|
parts=[ToolReturnPart("rag_cite", "Registered 1 citation.", "c1")]
|
|
),
|
|
ModelResponse(parts=[TextPart("an answer")]),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert returns_of(compacted) == ["Registered 1 citation."]
|
|
|
|
|
|
def test_an_uncited_earlier_image_is_dropped_with_its_label():
|
|
messages = [
|
|
*answered_question("first", evidence="OLD", images=True),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert images_of(compacted) == []
|
|
assert texts_of(compacted) == []
|
|
|
|
|
|
def test_cited_pictures_are_attached_beside_the_capsule():
|
|
messages = [
|
|
*answered_question("first", evidence="OLD", images=True),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
fresh = BinaryContent(data=b"cited-bytes", media_type="image/png")
|
|
|
|
compacted = compact_history(
|
|
messages,
|
|
boundary=4,
|
|
owned_tools=OWNED,
|
|
capsule_text="CAPSULE",
|
|
capsule_images=[picture_label("cited-chunk", "#/pictures/3"), fresh],
|
|
)
|
|
|
|
assert images_of(compacted) == [fresh]
|
|
assert texts_of(compacted) == [picture_label("cited-chunk", "#/pictures/3")]
|
|
carrier = [
|
|
index
|
|
for index, message in enumerate(compacted)
|
|
if isinstance(message, ModelRequest)
|
|
and any(
|
|
isinstance(part, ToolReturnPart) and part.content == "CAPSULE"
|
|
for part in message.parts
|
|
)
|
|
]
|
|
attached = [
|
|
index
|
|
for index, message in enumerate(compacted)
|
|
if isinstance(message, ModelRequest)
|
|
and any(
|
|
isinstance(part, UserPromptPart) and not isinstance(part.content, str)
|
|
for part in message.parts
|
|
)
|
|
]
|
|
assert carrier == attached
|
|
|
|
|
|
def test_a_user_attached_image_is_never_dropped():
|
|
"""The user's own picture is not ours to remove, even in an earlier question."""
|
|
mine = UserPromptPart(content=["look at this", PNG])
|
|
messages = [
|
|
ModelRequest(parts=[mine]),
|
|
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
|
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
|
|
ModelResponse(parts=[TextPart("an answer")]),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert images_of(compacted) == [PNG]
|
|
assert texts_of(compacted) == ["look at this"]
|
|
|
|
|
|
def test_the_stored_messages_are_never_mutated():
|
|
messages = [
|
|
*answered_question("first", evidence="OLD EVIDENCE", images=True),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
before = [list(message.parts) for message in messages]
|
|
|
|
compact_history(messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE")
|
|
|
|
assert [list(message.parts) for message in messages] == before
|
|
assert "OLD EVIDENCE" in returns_of(messages)
|
|
|
|
|
|
def test_nothing_cited_leaves_only_receipts():
|
|
messages = [
|
|
*answered_question("first", evidence="OLD"),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=4, owned_tools=OWNED, capsule_text=""
|
|
)
|
|
|
|
assert returns_of(compacted) == [RECEIPT]
|
|
|
|
|
|
@dataclass
|
|
class Deps:
|
|
state: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
def rag_and_compactor(temp_db_path):
|
|
return (
|
|
create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False),
|
|
create_compaction(),
|
|
)
|
|
|
|
|
|
def in_flight_history() -> list[Any]:
|
|
"""A question already asked and searched, still awaiting its answer."""
|
|
return [
|
|
ModelRequest(parts=[UserPromptPart("what does the supervisor do?")]),
|
|
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "s"}, "call-1")]),
|
|
ModelRequest(
|
|
parts=[ToolReturnPart("rag_search", "EVIDENCE FOR THE LIVE TURN", "call-1")]
|
|
),
|
|
]
|
|
|
|
|
|
def resuming_deps(question: int = 0) -> Deps:
|
|
"""State as a resumption always finds it: the question already identified."""
|
|
return Deps(
|
|
state={
|
|
"rag": RAGState(
|
|
evidence=CapabilityEvidenceRecord(question=question)
|
|
).model_dump(mode="json")
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_without_the_compactor_the_history_is_untouched(temp_db_path):
|
|
"""Omission is the switch: there is no flag to test, only absence."""
|
|
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
|
wire: list[list[Any]] = []
|
|
|
|
async def model(messages, _info):
|
|
wire.append(list(messages))
|
|
return ModelResponse(parts=[TextPart("answer")])
|
|
|
|
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
|
settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])]
|
|
|
|
await agent.run("a different question", deps=Deps(), message_history=settled)
|
|
|
|
assert returns_of(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_with_the_compactor_a_new_question_compacts_the_previous_one(
|
|
temp_db_path,
|
|
):
|
|
rag, compactor = rag_and_compactor(temp_db_path)
|
|
wire: list[list[Any]] = []
|
|
|
|
async def model(messages, _info):
|
|
wire.append(list(messages))
|
|
return ModelResponse(parts=[TextPart("answer")])
|
|
|
|
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
|
|
settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])]
|
|
|
|
await agent.run("a different question", deps=Deps(), message_history=settled)
|
|
|
|
assert returns_of(wire[-1]) == [RECEIPT]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resume_kwargs",
|
|
[
|
|
pytest.param({}, id="no prompt"),
|
|
pytest.param(
|
|
{"deferred_tool_results": DeferredToolResults()}, id="deferred results"
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_a_resumed_question_keeps_the_evidence_it_is_answering_from(
|
|
temp_db_path, resume_kwargs
|
|
):
|
|
"""The boundary is the stored identity of the question in progress.
|
|
|
|
An earlier question below it is compacted; the evidence the model is still
|
|
answering from sits above it and survives. Deriving the boundary from message
|
|
shape instead would put the live evidence below it and answer with a receipt
|
|
where the search result should be.
|
|
"""
|
|
rag, compactor = rag_and_compactor(temp_db_path)
|
|
wire: list[list[Any]] = []
|
|
|
|
async def model(messages, _info):
|
|
wire.append(list(messages))
|
|
return ModelResponse(parts=[TextPart("answer")])
|
|
|
|
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
|
|
history = [
|
|
*answered_question("an earlier question", evidence="EVIDENCE FOR THE OLD TURN"),
|
|
*in_flight_history(),
|
|
]
|
|
|
|
await agent.run(
|
|
deps=resuming_deps(question=4), message_history=history, **resume_kwargs
|
|
)
|
|
|
|
assert returns_of(wire[-1]) == [RECEIPT, "EVIDENCE FOR THE LIVE TURN"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compaction_never_reaches_the_stored_message_history(temp_db_path):
|
|
"""Rewriting is for the wire; hosts keep the evidence they gathered."""
|
|
rag, compactor = rag_and_compactor(temp_db_path)
|
|
turns = iter(
|
|
[
|
|
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
|
ModelResponse(parts=[TextPart("first answer")]),
|
|
ModelResponse(parts=[TextPart("second answer")]),
|
|
]
|
|
)
|
|
|
|
async def model(_messages, _info):
|
|
return next(turns)
|
|
|
|
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
|
|
deps = Deps()
|
|
|
|
with patch.object(
|
|
RAGCapability, "_search", AsyncMock(return_value="REAL EVIDENCE")
|
|
):
|
|
first = await agent.run("old question", deps=deps)
|
|
second = await agent.run(
|
|
"current question", deps=deps, message_history=first.all_messages()
|
|
)
|
|
|
|
assert "REAL EVIDENCE" in returns_of(second.all_messages())
|
|
assert RECEIPT not in returns_of(second.all_messages())
|
|
|
|
|
|
REAL_PNG = base64.b64decode(
|
|
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
|
|
)
|
|
|
|
|
|
async def _search_with_a_picture(self, query: str, _limit: int | None) -> str:
|
|
"""Record a result carrying a page image, the way a real search does."""
|
|
cast(Any, self.state).searches[query] = [
|
|
SearchResult(
|
|
content="evidence",
|
|
score=1.0,
|
|
chunk_id="chunk-1",
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/pictures/0"],
|
|
)
|
|
]
|
|
self._note_evidence()
|
|
return "EVIDENCE"
|
|
|
|
|
|
async def _cite_a_picture_chunk(temp_db_path, fetched: bytes | None):
|
|
"""Two questions: cite a picture chunk, then ask something else."""
|
|
rag, compactor = rag_and_compactor(temp_db_path)
|
|
calls = iter(
|
|
[
|
|
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
|
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
|
|
[TextPart("first answer")],
|
|
[TextPart("second answer")],
|
|
]
|
|
)
|
|
wire: list[list[Any]] = []
|
|
|
|
async def model(messages, _info):
|
|
wire.append(list(messages))
|
|
return ModelResponse(parts=next(calls))
|
|
|
|
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
|
|
deps = Deps()
|
|
|
|
with (
|
|
patch.object(RAGCapability, "_search", _search_with_a_picture),
|
|
patch.object(
|
|
RAGCapability, "get_picture_bytes", AsyncMock(return_value=fetched)
|
|
),
|
|
):
|
|
first = await agent.run("what does the supervisor do?", deps=deps)
|
|
await agent.run(
|
|
"and what else?", deps=deps, message_history=first.all_messages()
|
|
)
|
|
return wire
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_cited_picture_is_fetched_and_attached_with_its_label(temp_db_path):
|
|
wire = await _cite_a_picture_chunk(temp_db_path, REAL_PNG)
|
|
|
|
assert [picture.data for picture in images_of(wire[-1])] == [REAL_PNG]
|
|
assert texts_of(wire[-1]) == [picture_label("chunk-1", "#/pictures/0")]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_picture_that_cannot_be_fetched_emits_neither_image_nor_label(
|
|
temp_db_path,
|
|
):
|
|
"""A label without its picture tells the model a figure is there when it is not."""
|
|
wire = await _cite_a_picture_chunk(temp_db_path, None)
|
|
|
|
assert images_of(wire[-1]) == []
|
|
assert texts_of(wire[-1]) == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_picture_that_will_not_decode_emits_neither_image_nor_label(
|
|
temp_db_path,
|
|
):
|
|
"""One vision placeholder is rendered per attachment, so a corrupt one miscounts."""
|
|
wire = await _cite_a_picture_chunk(temp_db_path, b"not-an-image")
|
|
|
|
assert images_of(wire[-1]) == []
|
|
assert texts_of(wire[-1]) == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_the_capsule_is_built_once_per_request_and_again_for_the_next(
|
|
temp_db_path,
|
|
):
|
|
"""Two hook passes for one request must not rebuild; the next request must."""
|
|
rag, compactor = rag_and_compactor(temp_db_path)
|
|
deps = Deps()
|
|
ctx = RunContext(
|
|
deps=deps, model=TestModel(), usage=RunUsage(), run_id="run-1", run_step=1
|
|
)
|
|
run_rag = await rag.for_run(ctx)
|
|
run_compactor = await compactor.for_run(ctx)
|
|
cast(Any, run_rag.state).evidence.begin_question(4)
|
|
ctx = replace(ctx, capabilities={"rag": run_rag, "compaction": run_compactor})
|
|
builds = 0
|
|
|
|
def counting_build(evidence):
|
|
nonlocal builds
|
|
builds += 1
|
|
return Capsule(text="CAPSULE")
|
|
|
|
async def handler(_request_context):
|
|
return ModelResponse(parts=[TextPart("answer")])
|
|
|
|
request = ModelRequestContext(
|
|
messages=[*answered_question("first", evidence="OLD")],
|
|
model=TestModel(),
|
|
model_request_parameters=ModelRequestParameters(),
|
|
model_settings=None,
|
|
)
|
|
|
|
with patch("haiku.rag.capabilities.compaction.build_capsule", counting_build):
|
|
await run_compactor.wrap_model_request(
|
|
ctx, request_context=request, handler=handler
|
|
)
|
|
await run_compactor.wrap_model_request(
|
|
ctx, request_context=request, handler=handler
|
|
)
|
|
assert builds == 1
|
|
|
|
await run_compactor.wrap_model_request(
|
|
replace(ctx, run_step=2), request_context=request, handler=handler
|
|
)
|
|
|
|
assert builds == 2
|
|
|
|
|
|
def test_a_user_quoting_our_wording_keeps_their_picture_and_their_text():
|
|
"""Prose is not proof of ownership: a user can write any phrase.
|
|
|
|
Recognising our own pictures by a natural-language substring removed a user's
|
|
image, its text, and with it the whole message part.
|
|
"""
|
|
quoted = UserPromptPart(
|
|
content=[
|
|
"Here is a page image retrieved from the knowledge base for my report",
|
|
PNG,
|
|
]
|
|
)
|
|
messages = [
|
|
ModelRequest(parts=[quoted]),
|
|
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
|
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
|
|
ModelResponse(parts=[TextPart("an answer")]),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert images_of(compacted) == [PNG]
|
|
assert texts_of(compacted) == [
|
|
"Here is a page image retrieved from the knowledge base for my report"
|
|
]
|
|
assert all(message.parts for message in compacted)
|
|
|
|
|
|
def test_a_label_of_ours_with_no_picture_after_it_is_kept():
|
|
"""Only a genuine pair is ours to remove; a lone label is someone else's text."""
|
|
lonely = UserPromptPart(
|
|
content=[picture_label("chunk-1", "#/pictures/0"), "and more"]
|
|
)
|
|
messages = [
|
|
ModelRequest(parts=[lonely]),
|
|
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
|
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
|
|
ModelResponse(parts=[TextPart("an answer")]),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert texts_of(compacted) == [
|
|
picture_label("chunk-1", "#/pictures/0"),
|
|
"and more",
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_picture_whose_fetch_raises_costs_the_picture_not_the_answer(
|
|
temp_db_path,
|
|
):
|
|
rag, compactor = rag_and_compactor(temp_db_path)
|
|
calls = iter(
|
|
[
|
|
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
|
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
|
|
[TextPart("first answer")],
|
|
[TextPart("second answer")],
|
|
]
|
|
)
|
|
wire: list[list[Any]] = []
|
|
|
|
async def model(messages, _info):
|
|
wire.append(list(messages))
|
|
return ModelResponse(parts=next(calls))
|
|
|
|
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
|
|
deps = Deps()
|
|
|
|
with (
|
|
patch.object(RAGCapability, "_search", _search_with_a_picture),
|
|
patch.object(
|
|
RAGCapability,
|
|
"get_picture_bytes",
|
|
AsyncMock(side_effect=OSError("the read failed")),
|
|
),
|
|
):
|
|
first = await agent.run("what does the supervisor do?", deps=deps)
|
|
second = await agent.run(
|
|
"and what else?", deps=deps, message_history=first.all_messages()
|
|
)
|
|
|
|
assert second.output == "second answer"
|
|
assert images_of(wire[-1]) == []
|
|
assert texts_of(wire[-1]) == []
|
|
|
|
|
|
def test_a_request_is_never_left_with_no_parts():
|
|
"""Emptying a message would leave something that is not a message.
|
|
|
|
Our own pictures always travel with the tool return in their request, so this
|
|
shape does not come from us — but a rewritten history can hold it, and a
|
|
partless request is invalid whatever produced it.
|
|
"""
|
|
ours_alone = ModelRequest(
|
|
parts=[UserPromptPart(content=[picture_label("chunk-1", "#/pictures/0"), PNG])]
|
|
)
|
|
messages = [
|
|
ours_alone,
|
|
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
|
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
|
|
ModelResponse(parts=[TextPart("an answer")]),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert all(message.parts for message in compacted)
|
|
assert compacted[0] is ours_alone
|
|
|
|
|
|
def test_two_owned_returns_in_one_request_yield_one_capsule():
|
|
"""A model can search twice in one response, so a request can hold two returns.
|
|
|
|
Identifying the carrier by message alone gave every return in it the capsule,
|
|
which duplicates the whole thing — and it is unbounded.
|
|
"""
|
|
messages = [
|
|
ModelRequest(parts=[UserPromptPart("first")]),
|
|
ModelResponse(
|
|
parts=[
|
|
ToolCallPart("rag_search", {"query": "a"}, "call-1"),
|
|
ToolCallPart("rag_search", {"query": "b"}, "call-2"),
|
|
]
|
|
),
|
|
ModelRequest(
|
|
parts=[
|
|
ToolReturnPart("rag_search", "FIRST EVIDENCE", "call-1"),
|
|
ToolReturnPart("rag_search", "SECOND EVIDENCE", "call-2"),
|
|
]
|
|
),
|
|
ModelResponse(parts=[TextPart("an answer")]),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
|
|
compacted = compact_history(
|
|
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
|
)
|
|
|
|
assert returns_of(compacted) == [RECEIPT, "CAPSULE"]
|
|
|
|
|
|
def test_the_capsule_is_attached_beside_the_newest_return_of_that_request():
|
|
messages = [
|
|
ModelRequest(parts=[UserPromptPart("first")]),
|
|
ModelResponse(
|
|
parts=[
|
|
ToolCallPart("rag_search", {"query": "a"}, "call-1"),
|
|
ToolCallPart("rag_search", {"query": "b"}, "call-2"),
|
|
]
|
|
),
|
|
ModelRequest(
|
|
parts=[
|
|
ToolReturnPart("rag_search", "FIRST", "call-1"),
|
|
ToolReturnPart("rag_search", "SECOND", "call-2"),
|
|
]
|
|
),
|
|
ModelResponse(parts=[TextPart("an answer")]),
|
|
ModelRequest(parts=[UserPromptPart("second")]),
|
|
]
|
|
fresh = BinaryContent(data=b"cited-bytes", media_type="image/png")
|
|
|
|
compacted = compact_history(
|
|
messages,
|
|
boundary=4,
|
|
owned_tools=OWNED,
|
|
capsule_text="CAPSULE",
|
|
capsule_images=[picture_label("cited", "#/pictures/1"), fresh],
|
|
)
|
|
|
|
assert images_of(compacted) == [fresh]
|
|
assert returns_of(compacted) == [RECEIPT, "CAPSULE"]
|