diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 9752d2b0..93f5ac0c 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -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]: diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index 90bb8147..e2de8868 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -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() diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py new file mode 100644 index 00000000..7179ba27 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -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", +] diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index c0ab7f3f..2e140f6f 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -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 diff --git a/tests/capabilities/test_evidence_capsule.py b/tests/capabilities/test_evidence_capsule.py new file mode 100644 index 00000000..6f1a51b5 --- /dev/null +++ b/tests/capabilities/test_evidence_capsule.py @@ -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 == ""