From 9dc79e7fda19a366773f2be7587641b3dbf75c16 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 28 Aug 2026 10:35:38 +0300 Subject: [PATCH] Name the collection each piece of cited evidence came from A capsule replaces earlier questions' evidence with what they cited, so the searches that carried a `Collection:` line are gone by the time the model reads it. Cited content survived; which collection it came from did not, and a follow-up question attributed it to whichever the model guessed. Named the same way a search result names it: the capsule decides, and only when it spans more than one. --- docs/capabilities/compaction.md | 5 ++- .../haiku/rag/capabilities/compaction.py | 20 +++++++--- tests/capabilities/test_evidence_capsule.py | 40 +++++++++++++++++++ 3 files changed, 58 insertions(+), 7 deletions(-) diff --git a/docs/capabilities/compaction.md b/docs/capabilities/compaction.md index e0baeb25..057342d5 100644 --- a/docs/capabilities/compaction.md +++ b/docs/capabilities/compaction.md @@ -31,8 +31,9 @@ refuses rather than replace evidence it cannot retain. See On each request, evidence from earlier questions is replaced by the evidence those questions actually cited. Cited text and cited page images are kept in full, grouped by -the question that cited them, and stay citable by the same chunk ids. Every other -earlier evidence return becomes a short receipt. The current question is untouched. +the question that cited them, and stay citable by the same chunk ids. Evidence spanning +more than one collection carries a `Collection:` line naming the one it came from. Every +other earlier evidence return becomes a short receipt. The current question is untouched. Compaction rewrites the request, never the stored history, so `all_messages()` still holds everything the run gathered. diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index bd8a935b..ed171eb0 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -88,13 +88,21 @@ class _Entry: question: int citation: Citation - def render(self) -> str: + def render(self, *, include_collection: bool = False) -> str: + """Render an entry, optionally naming its collection.""" title = self.citation.document_title uri = self.citation.document_uri - source = f'"{title}"' if title else uri + document = 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}" + document = f"{document} ({uri})" + if include_collection and self.citation.source: + header = ( + f"[{self.chunk_id}]\nCollection: {self.citation.source}\n" + f"Source: {document}" + ) + else: + header = f"[{self.chunk_id}] Source: {document}" + return f"{header}\n{self.citation.content}" def _eligible_entries(evidence: Sequence[DiscoveredEvidence]) -> list[_Entry]: @@ -154,6 +162,8 @@ def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule: lines = [CAPSULE_HEADER] pictures: list[RetainedPicture] = [] seen: set[tuple[str, str, str]] = set() + # A capsule may combine citations from different search scopes. + include_collection = len({entry.citation.source for entry in entries}) > 1 position = 0 current_question: int | None = None for entry in entries: @@ -161,7 +171,7 @@ def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule: position += 1 current_question = entry.question lines.append(group_label(position)) - lines.append(entry.render()) + lines.append(entry.render(include_collection=include_collection)) 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 diff --git a/tests/capabilities/test_evidence_capsule.py b/tests/capabilities/test_evidence_capsule.py index c88e8b55..03217592 100644 --- a/tests/capabilities/test_evidence_capsule.py +++ b/tests/capabilities/test_evidence_capsule.py @@ -97,6 +97,46 @@ def test_a_retained_picture_carries_the_source_it_came_from(): assert picture.source == "beta" +def test_the_capsule_names_the_collection_evidence_came_from(): + found = discovered(cited={"a": [2], "b": [2]}) + found = replace( + found, + citations={ + "a": replace_citation(found.citations["a"], source="papers"), + "b": replace_citation(found.citations["b"], source="wiki"), + }, + ) + + lines = build_capsule([found]).text.splitlines() + + def rendered(chunk_id: str) -> list[str]: + start = lines.index(f"[{chunk_id}]") + return lines[start : start + 3] + + assert rendered("a") == [ + "[a]", + "Collection: papers", + 'Source: "Title a" (test://a)', + ] + assert rendered("b") == ["[b]", "Collection: wiki", 'Source: "Title b" (test://b)'] + + +def test_evidence_from_one_collection_does_not_name_it(): + found = discovered(cited={"a": [2], "b": [2]}) + found = replace( + found, + citations={ + chunk_id: replace_citation(cited, source="papers") + for chunk_id, cited in found.citations.items() + }, + ) + + lines = build_capsule([found]).text.splitlines() + + assert not [line for line in lines if line.startswith("Collection:")] + assert '[a] Source: "Title a" (test://a)' in lines + + def test_nothing_cited_produces_no_capsule(): capsule = build_capsule([discovered()])