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, EvidenceCompactionCapability, build_capsule, group_label, ) from haiku.rag.capabilities.compaction import create_capability as create_compaction from haiku.rag.capabilities.evidence import DiscoveredEvidence, discover_evidence 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, state_carried=True, 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"}), cite_available=True, ) def test_a_retained_picture_carries_the_source_it_came_from(): """A retained picture carries its database: compaction re-fetches cited pictures through it.""" found = discovered(cited={"c1": [2]}, pictures={"c1": ["#/pictures/0"]}) found = replace( found, citations={"c1": replace_citation(found.citations["c1"], source="beta")}, ) capsule = build_capsule([found]) [picture] = capsule.pictures 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_a_picture_in_two_collections_is_retained_from_each(): """A document copied into another collection keeps its id and picture refs.""" found = discovered( cited={"a": [2], "b": [2]}, pictures={"a": ["#/pictures/0"], "b": ["#/pictures/0"]}, ) found = replace( found, citations={ "a": replace_citation( found.citations["a"], document_id="shared", source="papers" ), "b": replace_citation( found.citations["b"], document_id="shared", source="wiki" ), }, ) capsule = build_capsule([found]) assert [(picture.source, picture.self_ref) for picture in capsule.pictures] == [ ("papers", "#/pictures/0"), ("wiki", "#/pictures/0"), ] # Nothing else tells the two apart once they are attached. assert "Collection: papers." in capsule.pictures[0].label assert "Collection: wiki." in capsule.pictures[1].label def test_one_chunk_id_cited_from_two_collections_labels_each_picture(): """Both capabilities can cite the same id from different collections, where the reference and the document are the same too.""" found = [ replace( discovered( capability=capability, cited={"c1": [2]}, pictures={"c1": ["#/pictures/0"]}, ), citations={ "c1": replace_citation( citation("c1", pictures=["#/pictures/0"]), document_id="shared", source=source, ) }, ) for capability, source in (("analysis", "papers"), ("rag", "wiki")) ] capsule = build_capsule(found) labels = [picture.label for picture in capsule.pictures] assert all("[c1] (#/pictures/0)" in label for label in labels) assert "Collection: papers." in labels[0] assert "Collection: wiki." in labels[1] def test_a_picture_from_one_collection_is_not_labelled_with_it(): found = discovered(cited={"c1": [2]}, pictures={"c1": ["#/pictures/0"]}) found = replace( found, citations={"c1": replace_citation(found.citations["c1"], source="papers")}, ) [picture] = build_capsule([found]).pictures assert "Collection" not in picture.label 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(discover_evidence(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 == ""