danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Build one capsule of cited evidence from the records

Browse source 
`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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-11 12:59:02 +03:00
parent 42d923fe4b
commit 0aa6d79f88
No known key found for this signature in database
5 changed files with 641 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
haiku_rag_slim/haiku/rag/capabilities/_base.py
View file

@ -312,7 +312,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
)
if spent := self._spent_tool_names():
names = ", ".join(sorted(spent))
if remaining := sorted(self._evidence_tool_names() - spent):
if remaining := sorted(self.evidence_tool_names() - spent):
return (
f"The {self.state_namespace} capability has spent its budget "
f"for {names}; further calls to them fail. Gather any further "
@ -350,8 +350,13 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
if tool.capability_id != self.id or tool.name == self._cite_tool_name
]
def _evidence_tool_names(self) -> set[str]:
"""Tools that can bring new evidence into the run."""
def evidence_tool_names(self) -> set[str]:
"""Tools that can bring new evidence into the run.
Public because compaction needs to know whose output on the wire is
evidence: a cite acknowledgement is a receipt of the model's own action and
must survive, while a code execution that reached the corpus is evidence.
"""
return {f"{self.state_namespace}_search"}
def _spent_tool_names(self) -> set[str]:

4
haiku_rag_slim/haiku/rag/capabilities/analysis.py
View file

@ -87,11 +87,11 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
self.sandbox = None
await super()._close()
def _evidence_tool_names(self) -> set[str]:
def evidence_tool_names(self) -> set[str]:
# Searching from inside the sandbox does not count against
# `qa.max_searches`, so code execution outlives a spent search budget
# as a way to reach new evidence.
return super()._evidence_tool_names() | {"analysis_execute_code"}
return super().evidence_tool_names() | {"analysis_execute_code"}
def _spent_tool_names(self) -> set[str]:
spent = super()._spent_tool_names()

243
haiku_rag_slim/haiku/rag/capabilities/compaction.py Normal file
View file

@ -0,0 +1,243 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, cast
from pydantic_ai import RunContext
from pydantic_ai.capabilities import AbstractCapability
from haiku.rag.capabilities._base import RAGCapabilityBase
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
from haiku.rag.store.models.citation import Citation
CAPABILITY_ID = "haiku-rag-evidence-compaction"
CAPSULE_HEADER = (
"[Evidence cited earlier in this conversation, kept so later questions can "
"rely on it. Cite these chunk_ids directly when you use them.]"
)
RECEIPT = (
"[Evidence retrieved for an earlier question, no longer shown. It does not "
"count as cited for the current question.]"
)
ENTRY_SEPARATOR = "\n\n"
def group_label(position: int) -> str:
"""Name a group by its position among the groups, not by question number.
A question identity is a message count, so a header built from it would present
an index as a turn number, and an ordinal over the groups is not the
conversation's ordinal either whenever a question in between cited nothing. The
label claims only what it is: a grouping, newest first.
"""
return f"[Cited evidence group {position}]"
def picture_label(chunk_id: str, self_ref: str) -> str:
return (
f"Page image retrieved from the knowledge base for cited evidence "
f"[{chunk_id}] ({self_ref}). Not provided by the user."
)
@dataclass(frozen=True)
class DiscoveredEvidence:
"""One evidence capability's records, as the compactor found them.
Read-only and rebuilt per request: the compactor merges these into a view and
persists nothing about evidence itself.
"""
capability: str
record: CapabilityEvidenceRecord
citations: Mapping[str, Citation]
tool_names: frozenset[str]
@dataclass(frozen=True)
class RetainedPicture:
"""A picture to re-attach, with the label that must accompany it.
Addressed by owner, document and reference, because a reference such as
``#/pictures/0`` repeats across documents and capabilities. The label travels
with it so it can never be emitted without its image.
"""
capability: str
chunk_id: str
document_id: str
self_ref: str
label: str
@dataclass(frozen=True)
class Capsule:
"""Everything the compactor would insert, and nothing about where it goes."""
text: str = ""
pictures: tuple[RetainedPicture, ...] = ()
@dataclass(frozen=True)
class _Entry:
capability: str
chunk_id: str
question: int
citation: Citation
def render(self) -> str:
title = self.citation.document_title
uri = self.citation.document_uri
source = f'"{title}"' if title else uri
if title and uri and uri != title:
source = f"{source} ({uri})"
return f"[{self.chunk_id}] Source: {source}\n{self.citation.content}"
def _eligible_entries(evidence: Sequence[DiscoveredEvidence]) -> list[_Entry]:
"""Cited evidence with content, newest citing question first.
Evidence cited in several questions belongs to the most recent one, so it is
rendered once and grouped where the model last used it.
An occurrence and its canonical ``Citation`` are written by the same call, so a
cited chunk without one is not a state this design produces. Rendering the rest
regardless would quietly drop evidence an answer rested on, so it is reported.
"""
entries = []
for discovered in evidence:
for chunk_id, occurrence in discovered.record.occurrences.items():
if not occurrence.cited_in_questions:
continue
citation = discovered.citations.get(chunk_id)
if citation is None:
raise ValueError(
f"{discovered.capability} cited {chunk_id} in question(s) "
f"{occurrence.cited_in_questions} but has no citation record "
"for it, so its content cannot be retained."
)
entries.append(
_Entry(
capability=discovered.capability,
chunk_id=chunk_id,
question=max(occurrence.cited_in_questions),
citation=citation,
)
)
entries.sort(key=lambda entry: (-entry.question, entry.capability, entry.chunk_id))
return entries
def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule:
"""Render every cited piece of evidence, grouped by the question that cited it.
Everything cited is kept whole and everything else is dropped. There is no
character budget: what a model can hold is the model's business, and a knob for
it would only half-rescue models that fail on long conversations regardless.
A host that needs earlier evidence pruned can compact further on top, on the wire
only. Removing or reordering the stored history breaks the message counts that
question identities and epochs are derived from, and the next record written is
refused.
Pure: no I/O and no message history, so what goes on the wire stays separable
from what it should contain. Picture bytes are fetched by the caller, which is
why a picture travels with its label rather than beside it.
"""
entries = _eligible_entries(evidence)
if not entries:
return Capsule()
lines = [CAPSULE_HEADER]
pictures: list[RetainedPicture] = []
seen: set[tuple[str, str, str]] = set()
position = 0
current_question: int | None = None
for entry in entries:
if entry.question != current_question:
position += 1
current_question = entry.question
lines.append(group_label(position))
lines.append(entry.render())
for self_ref in entry.citation.picture_refs:
# Overlapping chunks cite one figure, and a provider counts it twice.
# Identity is owner plus document plus reference, so the same reference
# in another document stays a different picture.
identity = (entry.capability, entry.citation.document_id, self_ref)
if identity in seen:
continue
seen.add(identity)
pictures.append(
RetainedPicture(
capability=entry.capability,
chunk_id=entry.chunk_id,
document_id=entry.citation.document_id,
self_ref=self_ref,
label=picture_label(entry.chunk_id, self_ref),
)
)
return Capsule(text=ENTRY_SEPARATOR.join(lines), pictures=tuple(pictures))
@dataclass
class EvidenceCompactionCapability(AbstractCapability[Any]):
"""Rewrites the history from what the evidence capabilities recorded.
Registering it is what turns compaction on: a host that leaves it out gets an
untouched transcript, which is why it has no enable flag. It reads the evidence
capabilities through the run's registry and holds no reference to any of them,
so a host running one capability, both, or neither needs no wiring change.
Registering two is rejected by pydantic-ai before the run starts, since they
would share this capability's id.
"""
def discover(self, ctx: RunContext[Any]) -> list[DiscoveredEvidence]:
"""Read what each evidence capability recorded, without writing anything.
The registry holds the per-run instances, which are the ones carrying
state; the registered objects never do. That includes a deferred capability
the model has not loaded, whose record is simply empty.
"""
discovered = []
for capability in ctx.capabilities.values():
if not isinstance(capability, RAGCapabilityBase):
continue
state = capability.state
discovered.append(
DiscoveredEvidence(
capability=capability.state_namespace,
record=cast(CapabilityEvidenceRecord, cast(Any, state).evidence),
citations=cast(Any, state).citation_index,
tool_names=frozenset(capability.evidence_tool_names()),
)
)
return sorted(discovered, key=lambda evidence: evidence.capability)
def create_capability() -> EvidenceCompactionCapability:
"""Create the capability that compacts history from recorded evidence."""
return EvidenceCompactionCapability(
id=CAPABILITY_ID,
description=(
"Replaces earlier questions' evidence on the wire with a capsule of "
"what was cited."
),
)
__all__ = [
"CAPABILITY_ID",
"CAPSULE_HEADER",
"RECEIPT",
"Capsule",
"DiscoveredEvidence",
"EvidenceCompactionCapability",
"RetainedPicture",
"build_capsule",
"create_capability",
"group_label",
"picture_label",
]

4
tests/capabilities/test_capabilities.py
View file

@ -614,7 +614,7 @@ async def test_spent_search_notice_points_at_code_while_it_has_budget(temp_db_pa
notice = capability._budget_notice()
assert notice is not None
assert "analysis_execute_code" in notice
assert capability._evidence_tool_names() <= capability._spent_tool_names()
assert capability.evidence_tool_names() <= capability._spent_tool_names()
@pytest.mark.asyncio
@ -629,7 +629,7 @@ async def test_spent_search_notice_tells_rag_to_answer(temp_db_path):
assert notice is not None
assert "rag_search" in notice
assert capability._evidence_tool_names() == {"rag_search"}
assert capability.evidence_tool_names() == {"rag_search"}
@pytest.mark.asyncio

386
tests/capabilities/test_evidence_capsule.py Normal file
View file

@ -0,0 +1,386 @@
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 == ""