0aa6d79f88
`EvidenceCompactionCapability` reads what the evidence capabilities recorded out of the run registry, and `build_capsule` renders it: every cited item, grouped by the question that last cited it, newest group first, each rendered once, with the pictures of cited evidence and the labels that must accompany them. Discovery runs one way and reads only, so no capability holds a reference to another and a host running one, both or neither needs no wiring change. Everything cited is kept whole and everything else is dropped. There is no character budget, no picture cap and nothing to configure: a cap would only half-rescue models that fail on long conversations regardless, and a host that needs earlier evidence pruned can compact its own requests further. A capability reports which of its tools produce evidence, so a cite acknowledgement is never mistaken for one. Pictures are identified by owner, document and reference, so one figure cited through overlapping chunks is attached once while the same reference in another document stays a different picture. The builder does no I/O and never sees the message history, so a picture travels with its label and the caller fetches the bytes. Nothing reaches the wire yet.
386 lines
12 KiB
Python
386 lines
12 KiB
Python
from dataclasses import dataclass, field, replace
|
|
from typing import Any
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.exceptions import UserError
|
|
from pydantic_ai.messages import ModelResponse, TextPart
|
|
from pydantic_ai.models.function import FunctionModel
|
|
|
|
from haiku.rag.capabilities.analysis import create_capability as create_analysis
|
|
from haiku.rag.capabilities.compaction import (
|
|
CAPSULE_HEADER,
|
|
DiscoveredEvidence,
|
|
EvidenceCompactionCapability,
|
|
build_capsule,
|
|
group_label,
|
|
)
|
|
from haiku.rag.capabilities.compaction import create_capability as create_compaction
|
|
from haiku.rag.capabilities.ledger import (
|
|
CapabilityEvidenceRecord,
|
|
EvidenceOccurrence,
|
|
)
|
|
from haiku.rag.capabilities.rag import create_capability as create_rag
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.store.models.citation import Citation
|
|
|
|
|
|
@dataclass
|
|
class Deps:
|
|
state: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
def citation(chunk_id: str, content: str = "evidence body", pictures=()) -> Citation:
|
|
return Citation(
|
|
document_id=f"doc-of-{chunk_id}",
|
|
chunk_id=chunk_id,
|
|
document_uri=f"test://{chunk_id}",
|
|
document_title=f"Title {chunk_id}",
|
|
content=content,
|
|
picture_refs=list(pictures),
|
|
)
|
|
|
|
|
|
def replace_citation(cited: Citation, **changes: Any) -> Citation:
|
|
return cited.model_copy(update=changes)
|
|
|
|
|
|
def discovered(
|
|
capability: str = "rag",
|
|
*,
|
|
cited: dict[str, list[int]] | None = None,
|
|
contents: dict[str, str] | None = None,
|
|
pictures: dict[str, list[str]] | None = None,
|
|
) -> DiscoveredEvidence:
|
|
"""One capability's records, as the compactor would find them."""
|
|
cited = cited or {}
|
|
contents = contents or {}
|
|
pictures = pictures or {}
|
|
record = CapabilityEvidenceRecord(question=max(max(cited.values(), default=[0])))
|
|
for chunk_id, questions in cited.items():
|
|
record.occurrences[chunk_id] = EvidenceOccurrence(
|
|
capability=capability,
|
|
chunk_id=chunk_id,
|
|
retrieved_in_questions=list(questions),
|
|
cited_in_questions=list(questions),
|
|
)
|
|
return DiscoveredEvidence(
|
|
capability=capability,
|
|
record=record,
|
|
citations={
|
|
chunk_id: citation(
|
|
chunk_id,
|
|
contents.get(chunk_id, "evidence body"),
|
|
pictures.get(chunk_id, ()),
|
|
)
|
|
for chunk_id in cited
|
|
},
|
|
tool_names=frozenset({f"{capability}_search"}),
|
|
)
|
|
|
|
|
|
def test_nothing_cited_produces_no_capsule():
|
|
capsule = build_capsule([discovered()])
|
|
|
|
assert capsule.text == ""
|
|
assert capsule.pictures == ()
|
|
|
|
|
|
def test_cited_evidence_is_grouped_newest_question_first():
|
|
capsule = build_capsule([discovered(cited={"old": [2], "new": [8]})])
|
|
|
|
assert capsule.text.index(group_label(1)) < capsule.text.index(group_label(2))
|
|
assert capsule.text.index("[new]") < capsule.text.index("[old]")
|
|
assert CAPSULE_HEADER in capsule.text
|
|
|
|
|
|
def test_an_entry_is_rendered_once_in_its_most_recent_citing_group():
|
|
capsule = build_capsule([discovered(cited={"reused": [2, 8], "only-old": [2]})])
|
|
|
|
assert capsule.text.count("[reused]") == 1
|
|
assert capsule.text.index("[reused]") < capsule.text.index("[only-old]")
|
|
|
|
|
|
def test_evidence_cited_in_one_question_forms_one_group():
|
|
capsule = build_capsule([discovered(cited={"a": [4], "b": [4]})])
|
|
|
|
assert group_label(1) in capsule.text
|
|
assert group_label(2) not in capsule.text
|
|
|
|
|
|
def test_every_cited_entry_is_kept_whole():
|
|
"""No budget: a long citation is retained in full rather than truncated."""
|
|
body = "L" * 20_000
|
|
capsule = build_capsule([discovered(cited={"long": [4]}, contents={"long": body})])
|
|
|
|
assert body in capsule.text
|
|
|
|
|
|
def test_both_capabilities_share_one_capsule():
|
|
capsule = build_capsule(
|
|
[
|
|
discovered("rag", cited={"from-rag": [4]}),
|
|
discovered("analysis", cited={"from-analysis": [6]}),
|
|
]
|
|
)
|
|
|
|
assert capsule.text.count(CAPSULE_HEADER) == 1
|
|
assert "[from-rag]" in capsule.text
|
|
assert "[from-analysis]" in capsule.text
|
|
|
|
|
|
def test_the_same_chunk_id_under_two_capabilities_is_kept_apart():
|
|
capsule = build_capsule(
|
|
[
|
|
discovered("rag", cited={"shared": [4]}, contents={"shared": "rag body"}),
|
|
discovered(
|
|
"analysis", cited={"shared": [4]}, contents={"shared": "analysis body"}
|
|
),
|
|
]
|
|
)
|
|
|
|
assert "rag body" in capsule.text
|
|
assert "analysis body" in capsule.text
|
|
|
|
|
|
def test_cited_evidence_with_no_canonical_citation_is_an_error():
|
|
"""Both are written by the same call, so divergence is not a valid state.
|
|
|
|
Rendering the rest would quietly drop evidence an answer rested on, against
|
|
the one guarantee this capsule makes.
|
|
"""
|
|
evidence = discovered(cited={"present": [4]})
|
|
evidence.record.occurrences["absent"] = EvidenceOccurrence(
|
|
capability="rag", chunk_id="absent", cited_in_questions=[4]
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="absent"):
|
|
build_capsule([evidence])
|
|
|
|
|
|
def test_retrieved_but_uncited_evidence_is_not_kept():
|
|
evidence = discovered(cited={"cited": [4]})
|
|
evidence.record.occurrences["seen-only"] = EvidenceOccurrence(
|
|
capability="rag", chunk_id="seen-only", retrieved_in_questions=[4]
|
|
)
|
|
|
|
capsule = build_capsule([evidence])
|
|
|
|
assert "[cited]" in capsule.text
|
|
assert "seen-only" not in capsule.text
|
|
|
|
|
|
def test_a_source_is_named_once_when_the_title_is_the_uri():
|
|
"""Real corpora set both to the document id, which reads as a stutter."""
|
|
evidence = discovered(cited={"a": [4]})
|
|
evidence.citations["a"].document_title = "2410.11843v5"
|
|
evidence.citations["a"].document_uri = "2410.11843v5"
|
|
|
|
capsule = build_capsule([evidence])
|
|
|
|
assert 'Source: "2410.11843v5"' in capsule.text
|
|
assert "2410.11843v5)" not in capsule.text
|
|
|
|
|
|
def test_pictures_of_cited_evidence_are_all_retained_newest_first():
|
|
capsule = build_capsule(
|
|
[
|
|
discovered(
|
|
cited={"a": [2], "c": [6]},
|
|
pictures={"a": ["#/pictures/0"], "c": ["#/pictures/1", "#/pictures/2"]},
|
|
)
|
|
]
|
|
)
|
|
|
|
assert [picture.self_ref for picture in capsule.pictures] == [
|
|
"#/pictures/1",
|
|
"#/pictures/2",
|
|
"#/pictures/0",
|
|
]
|
|
assert capsule.pictures[0].document_id == "doc-of-c"
|
|
assert capsule.pictures[0].capability == "rag"
|
|
|
|
|
|
def test_a_picture_of_uncited_evidence_is_not_retained():
|
|
found = discovered(cited={"cited": [4]}, pictures={"cited": ["#/pictures/0"]})
|
|
found.record.occurrences["seen-only"] = EvidenceOccurrence(
|
|
capability="rag", chunk_id="seen-only", retrieved_in_questions=[4]
|
|
)
|
|
evidence = replace(
|
|
found,
|
|
citations={
|
|
**found.citations,
|
|
"seen-only": citation("seen-only", pictures=["#/pictures/9"]),
|
|
},
|
|
)
|
|
|
|
capsule = build_capsule([evidence])
|
|
|
|
assert [picture.self_ref for picture in capsule.pictures] == ["#/pictures/0"]
|
|
|
|
|
|
def test_one_picture_cited_through_two_chunks_is_attached_once():
|
|
"""Overlapping chunks of one document share a figure, counted twice by a provider."""
|
|
found = discovered(
|
|
cited={"first": [4], "second": [4]},
|
|
pictures={"first": ["#/pictures/1"], "second": ["#/pictures/1"]},
|
|
)
|
|
shared = {
|
|
chunk_id: replace_citation(cited, document_id="doc-shared")
|
|
for chunk_id, cited in found.citations.items()
|
|
}
|
|
|
|
capsule = build_capsule([replace(found, citations=shared)])
|
|
|
|
assert len(capsule.pictures) == 1
|
|
assert capsule.pictures[0].chunk_id == "first"
|
|
|
|
|
|
def test_the_same_reference_in_two_documents_is_kept_twice():
|
|
"""``#/pictures/1`` means a different figure in a different document.
|
|
|
|
One capability throughout, so only the document differs: dropping the document
|
|
from the identity would have to fail this.
|
|
"""
|
|
capsule = build_capsule(
|
|
[
|
|
discovered(
|
|
"rag",
|
|
cited={"a": [4], "b": [4]},
|
|
pictures={"a": ["#/pictures/1"], "b": ["#/pictures/1"]},
|
|
)
|
|
]
|
|
)
|
|
|
|
assert len(capsule.pictures) == 2
|
|
assert {picture.capability for picture in capsule.pictures} == {"rag"}
|
|
assert {picture.document_id for picture in capsule.pictures} == {
|
|
"doc-of-a",
|
|
"doc-of-b",
|
|
}
|
|
|
|
|
|
def test_a_picture_label_names_the_chunk_it_belongs_to():
|
|
capsule = build_capsule(
|
|
[discovered(cited={"a": [4]}, pictures={"a": ["#/pictures/0"]})]
|
|
)
|
|
|
|
label = capsule.pictures[0].label
|
|
assert "[a]" in label
|
|
assert "#/pictures/0" in label
|
|
assert "knowledge base" in label
|
|
assert "Not provided by the user" in label
|
|
|
|
|
|
def _spy_discovery(found: list[list[DiscoveredEvidence]]):
|
|
"""Discover from ``before_run``, the earliest point the registry is reliable."""
|
|
original = EvidenceCompactionCapability.before_run
|
|
|
|
async def spy(self, ctx):
|
|
await original(self, ctx)
|
|
found.append(self.discover(ctx))
|
|
|
|
return patch.object(EvidenceCompactionCapability, "before_run", spy)
|
|
|
|
|
|
async def _answer(_messages, _info):
|
|
return ModelResponse(parts=[TextPart("answer")])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_the_compactor_discovers_both_evidence_capabilities(temp_db_path):
|
|
"""Discovery runs one way through the registry, so nothing needs wiring."""
|
|
compactor = create_compaction()
|
|
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
|
analysis = create_analysis(
|
|
db_path=temp_db_path, config=AppConfig(), defer_loading=False
|
|
)
|
|
found: list[list[DiscoveredEvidence]] = []
|
|
|
|
with _spy_discovery(found):
|
|
agent = Agent(
|
|
FunctionModel(_answer),
|
|
deps_type=Deps,
|
|
capabilities=[rag, analysis, compactor],
|
|
)
|
|
await agent.run("a question", deps=Deps())
|
|
|
|
assert {evidence.capability: set(evidence.tool_names) for evidence in found[0]} == {
|
|
"rag": {"rag_search"},
|
|
"analysis": {"analysis_search", "analysis_execute_code"},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_discovery_sees_the_run_instances_not_the_registered_ones(temp_db_path):
|
|
"""A registered capability holds no state; only its per-run copy does."""
|
|
compactor = create_compaction()
|
|
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
|
found: list[list[DiscoveredEvidence]] = []
|
|
|
|
with _spy_discovery(found):
|
|
agent = Agent(
|
|
FunctionModel(_answer), deps_type=Deps, capabilities=[rag, compactor]
|
|
)
|
|
await agent.run("a question", deps=Deps())
|
|
|
|
assert rag.state is None
|
|
assert found[0][0].record.question == 0
|
|
|
|
|
|
def test_two_compactors_fail_fast(temp_db_path):
|
|
"""Each would rewrite the same history and each would build its own capsule.
|
|
|
|
They share this capability's id, so pydantic-ai refuses at construction and
|
|
nothing here has to police it.
|
|
"""
|
|
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
|
|
|
with pytest.raises(UserError, match="unique within a run"):
|
|
Agent(
|
|
FunctionModel(_answer),
|
|
deps_type=Deps,
|
|
capabilities=[rag, create_compaction(), create_compaction()],
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_compactor_alone_discovers_nothing_and_still_runs():
|
|
found: list[list[DiscoveredEvidence]] = []
|
|
|
|
with _spy_discovery(found):
|
|
agent = Agent(
|
|
FunctionModel(_answer), deps_type=Deps, capabilities=[create_compaction()]
|
|
)
|
|
result = await agent.run("a question", deps=Deps())
|
|
|
|
assert found == [[]]
|
|
assert result.output == "answer"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_deferred_capability_the_model_never_loaded_has_an_empty_record(
|
|
temp_db_path,
|
|
):
|
|
"""It is still discovered, because every registered capability gets a run copy.
|
|
|
|
Nothing was retrieved under it, so its record contributes no entries and the
|
|
compactor needs no special case for it.
|
|
"""
|
|
deferred = create_rag(db_path=temp_db_path, config=AppConfig())
|
|
found: list[list[DiscoveredEvidence]] = []
|
|
|
|
with _spy_discovery(found):
|
|
agent = Agent(
|
|
FunctionModel(_answer),
|
|
deps_type=Deps,
|
|
capabilities=[deferred, create_compaction()],
|
|
)
|
|
await agent.run("a question", deps=Deps())
|
|
|
|
assert deferred.defer_loading is True
|
|
assert deferred.state is None
|
|
assert [evidence.capability for evidence in found[0]] == ["rag"]
|
|
assert found[0][0].record.occurrences == {}
|
|
assert build_capsule(found[0]).text == ""
|