diff --git a/CHANGELOG.md b/CHANGELOG.md index f3175b97..6491603a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,8 @@ ### Added +- `CitationPolicyCapability` (`haiku.rag.capabilities.policy.create_capability`): registering it requires every answer to declare its grounding. A question that ends undeclared is sent back once to record what grounded the answer already given; when the cite tool is no longer available the question is recorded in `CitationPolicyState.violations` instead. A question that gathered no evidence is left alone. +- `haiku.rag.capabilities.evidence.discover_evidence()` and `DiscoveredEvidence`, moved out of `compaction` so both optional capabilities share them. `RAGCapabilityBase.cite_available`. - `EvidenceCompactionCapability` (`haiku.rag.capabilities.compaction.create_capability`): registering it replaces earlier questions' evidence on the model request with the evidence that was cited, grouped by the question that cited it, cited page images re-attached, other earlier evidence returns reduced to a receipt. Requests only; `all_messages()` is untouched. No configuration. - `RAGState.evidence` / `AnalysisState.evidence` (`CapabilityEvidenceRecord`): which evidence a capability retrieved and cited, per question, keyed by message-count question identities and epochs. `haiku.rag.capabilities.ledger.citation_status(records, question=...)` derives `missing` / `grounded` / `ungrounded` across capabilities. - `RAGCapabilityBase.evidence_tool_names()` and `get_picture_bytes()`. @@ -10,6 +12,7 @@ ### Changed +- `rag_cite` / `analysis_cite` accept an empty `chunk_ids`, recording the answer as ungrounded rather than failing the call, and the instructions no longer exempt a refusal or a corpus-level computation from citing. - `RAGCapability` and `AnalysisCapability` no longer rewrite the model request. Register `create_capability()` from `haiku.rag.capabilities.compaction` alongside them to keep earlier questions compacted. - Resuming a run (no prompt, deferred tool results, an unfinished history tail) raises `RuntimeError` unless the host carries the capability state from the run being resumed. diff --git a/docs/capabilities/index.md b/docs/capabilities/index.md index bff95f8e..1c6f9b7f 100644 --- a/docs/capabilities/index.md +++ b/docs/capabilities/index.md @@ -7,6 +7,7 @@ haiku.rag provides native [Pydantic AI capabilities](https://ai.pydantic.dev/cap | [`RAGCapability`](rag.md) | Grounded document search and citations. | | [`AnalysisCapability`](analysis.md) | Corpus computation and structural analysis with sandboxed Python. | | `EvidenceCompactionCapability` | Optional. Shrinking a conversation's history to the evidence that was cited. | +| `CitationPolicyCapability` | Optional. Requiring every answer to declare what grounds it. | The two evidence capabilities are deferred by default. An agent initially sees only their descriptions and the standard `load_capability` tool. Instructions and tools enter the model context only when the model loads a capability. @@ -67,6 +68,36 @@ the host to carry the capability state from the run being resumed, alongside the message history. Without it the identity of the question in progress is unknowable and the run fails rather than silently treating it as a new question. +## Requiring citations + +Citing is always available and always recorded, but nothing requires it. Register the +citation policy capability to make every answer declare its grounding: + +```python +from haiku.rag.capabilities.policy import create_capability as citation_policy +from haiku.rag.capabilities.rag import create_capability as rag + +agent = Agent( + "openai:gpt-5", + capabilities=[rag(db_path="my.lancedb"), citation_policy()], +) +``` + +An empty citation is a valid declaration: a model that finds nothing relevant calls +the cite tool with an empty list, which records the answer as *ungrounded* — distinct +from an answer that declared nothing at all. That distinction is what makes requiring +a declaration possible without forcing the model to invent grounding. + +When a question ends undeclared, the model is asked once to record what grounded the +answer it already gave. It is not asked to change the answer. If the cite tool is no +longer available by then, the question is recorded as a violation in +`CitationPolicyState` under `"citation_policy"` instead, since pointing a model at a +tool that is gone costs it retries. A question that gathered no evidence at all — a +greeting, an aside — is left alone. + +Exactly one policy capability makes the decision, however many evidence capabilities +are registered, so two of them cannot each demand a citation for one answer. + ## State Capabilities use a plain `state: dict[str, Any]` attribute on agent dependencies when one is available. RAG state lives under `"rag"`; analysis state lives under `"analysis"`. This keeps state independent of any transport or UI protocol. diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index e2c051f0..85e5170e 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -1,6 +1,6 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from dataclasses import dataclass, field, replace -from typing import Any, cast +from typing import Any from pydantic_ai import RunContext from pydantic_ai.capabilities import AbstractCapability, WrapModelRequestHandler @@ -15,7 +15,11 @@ from pydantic_ai.messages import ( from pydantic_ai.models import ModelRequestContext from haiku.rag.capabilities._base import RAGCapabilityBase -from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord +from haiku.rag.capabilities.evidence import ( + DiscoveredEvidence, + discover_evidence, + question_in_progress, +) from haiku.rag.store.models.citation import Citation from haiku.rag.tools.search import RETRIEVED_IMAGE_TAG, decode_picture @@ -52,20 +56,6 @@ def picture_label(chunk_id: str, self_ref: str) -> str: ) -@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. @@ -329,8 +319,8 @@ class EvidenceCompactionCapability(AbstractCapability[Any]): of what was retrieved and break the message counts that question identities and epochs are derived from. """ - evidence = self.discover(ctx) - boundary = max((found.record.question or 0 for found in evidence), default=0) + evidence = discover_evidence(ctx) + boundary = question_in_progress(evidence) if boundary > 0: await self._build_once(ctx, evidence) request_context.messages = compact_history( @@ -395,28 +385,6 @@ class EvidenceCompactionCapability(AbstractCapability[Any]): content.append(picture) return tuple(content) - 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.""" @@ -434,7 +402,6 @@ __all__ = [ "CAPSULE_HEADER", "RECEIPT", "Capsule", - "DiscoveredEvidence", "EvidenceCompactionCapability", "RetainedPicture", "build_capsule", diff --git a/haiku_rag_slim/haiku/rag/capabilities/evidence.py b/haiku_rag_slim/haiku/rag/capabilities/evidence.py new file mode 100644 index 00000000..75a61cc4 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/capabilities/evidence.py @@ -0,0 +1,63 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, cast + +from pydantic_ai import RunContext + +from haiku.rag.capabilities._base import RAGCapabilityBase +from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord +from haiku.rag.store.models.citation import Citation + + +@dataclass(frozen=True) +class DiscoveredEvidence: + """One evidence capability's records, as another capability found them. + + Read-only and rebuilt per request: whoever discovers these merges them into a + view and persists nothing about evidence itself. + """ + + capability: str + record: CapabilityEvidenceRecord + citations: Mapping[str, Citation] + tool_names: frozenset[str] + cite_available: bool + + +def discover_evidence(ctx: RunContext[Any]) -> list[DiscoveredEvidence]: + """Read what each evidence capability recorded, without writing anything. + + Discovery runs one way through the run's capability registry, so no capability + holds a reference to another, and a host running one, both, or neither needs no + wiring change. 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 = [ + DiscoveredEvidence( + capability=capability.state_namespace, + record=cast(CapabilityEvidenceRecord, cast(Any, capability.state).evidence), + citations=cast(Any, capability.state).citation_index, + tool_names=frozenset(capability.evidence_tool_names()), + cite_available=capability.cite_available, + ) + for capability in ctx.capabilities.values() + if isinstance(capability, RAGCapabilityBase) + ] + return sorted(discovered, key=lambda evidence: evidence.capability) + + +def question_in_progress(evidence: list[DiscoveredEvidence]) -> int: + """The identity every evidence capability agrees this question has. + + They all derive it from the same history, so they agree; taking the maximum + rather than a first entry keeps the result independent of ordering. + """ + return max((found.record.question or 0 for found in evidence), default=0) + + +__all__ = [ + "DiscoveredEvidence", + "discover_evidence", + "question_in_progress", +] diff --git a/haiku_rag_slim/haiku/rag/capabilities/policy.py b/haiku_rag_slim/haiku/rag/capabilities/policy.py new file mode 100644 index 00000000..662518f3 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/capabilities/policy.py @@ -0,0 +1,147 @@ +from dataclasses import dataclass, field, replace +from typing import Any + +from pydantic import BaseModel, Field +from pydantic_ai import RunContext +from pydantic_ai.capabilities import AbstractCapability +from pydantic_ai.messages import ModelResponse, ToolCallPart +from pydantic_ai.models import ModelRequestContext + +from haiku.rag.capabilities.evidence import ( + DiscoveredEvidence, + discover_evidence, + question_in_progress, +) +from haiku.rag.capabilities.ledger import citation_status + +CAPABILITY_ID = "haiku-rag-citation-policy" + +STATE_NAMESPACE = "citation_policy" + +REDIRECT_HINT = "record what grounded the answer you already gave" + +REDIRECT = ( + "You answered without registering citations. This asks you to " + f"{REDIRECT_HINT} — it is not a request to change that answer, and not a " + "signal that it was wrong. Call the cite tool with the chunk_ids that support " + "it. If nothing in the knowledge base supports it, or you said you could not " + "find the information, call it with an empty list. Then repeat your answer " + "exactly as you gave it." +) + + +class CitationPolicyState(BaseModel): + """What the policy decided, for hosts and evaluations to read. + + ``violations`` holds the identities of questions that ended undeclared while + the cite tool was already gone, so no redirect was possible. It is an + enforcement outcome, which is why it lives here rather than in an evidence + capability's record: nothing the model declared says it. + """ + + violations: list[int] = Field(default_factory=list) + + +@dataclass +class CitationPolicyCapability(AbstractCapability[Any]): + """Requires every answer to declare what grounds it, once per question. + + Registering it is the only switch. Without it citations are still recorded and + still validated, they are simply not required. + + Enforcement needs exactly one decision-maker. If each evidence capability + enforced its own citations, both could redirect the model within one question + and neither could see what the other had declared, so this capability + discovers them all and merges their records before deciding. + + Registering two is rejected by pydantic-ai before the run starts, since they + would share this capability's id. + """ + + redirected: set[int] = field(default_factory=set, repr=False) + + async def for_run(self, ctx: RunContext[Any]) -> "CitationPolicyCapability": + """Give the run its own record of what it has already asked for.""" + return replace(self, redirected=set()) + + async def after_model_request( + self, + ctx: RunContext[Any], + *, + request_context: ModelRequestContext, + response: ModelResponse, + ) -> ModelResponse: + """Decide once, at the last moment a question can still be redirected. + + A response carrying no tool calls ends the question, so there is no later + opportunity. Citing is unconditional, so an undeclared answer is a protocol + breach whether the model answered or refused, and this never has to guess + which it was. + """ + if any(isinstance(part, ToolCallPart) for part in response.parts): + return response + evidence = discover_evidence(ctx) + question = question_in_progress(evidence) + if question in self.redirected or not _gathered_evidence(evidence): + return response + records = [found.record for found in evidence] + if citation_status(records, question=question) != "missing": + return response + + self.redirected.add(question) + if any(found.cite_available for found in evidence): + ctx.enqueue(REDIRECT, priority="when_idle") + else: + self._record_violation(ctx, question) + return response + + def _record_violation(self, ctx: RunContext[Any], question: int) -> None: + """Note a question that could not be asked to cite, the tool being gone.""" + outer = getattr(ctx.deps, "state", None) + if not isinstance(outer, dict): + return + state = CitationPolicyState.model_validate(outer.get(STATE_NAMESPACE) or {}) + state.violations.append(question) + outer[STATE_NAMESPACE] = state.model_dump(mode="json") + + async def before_run(self, ctx: RunContext[Any]) -> None: + """Publish an empty outcome, so a host can tell "none" from "not running".""" + outer = getattr(ctx.deps, "state", None) + if isinstance(outer, dict): + outer.setdefault( + STATE_NAMESPACE, CitationPolicyState().model_dump(mode="json") + ) + + +def _gathered_evidence(evidence: list[DiscoveredEvidence]) -> bool: + """Whether this question produced anything an answer could be grounded on. + + A question with no evidence outcome has nothing to declare — a greeting, or a + conversational aside. Read from the ledger rather than from ``state.searches``, + which a new question clears, so an answer grounded on code execution or on a + document read counts as well. + """ + question = question_in_progress(evidence) + return any(found.record.latest_evidence_epoch > question for found in evidence) + + +def create_capability() -> CitationPolicyCapability: + """Create the capability that requires an answer to declare its grounding.""" + return CitationPolicyCapability( + id=CAPABILITY_ID, + description=( + "Requires every answer to register the evidence that grounds it, or to " + "declare that nothing does." + ), + ) + + +__all__ = [ + "CAPABILITY_ID", + "REDIRECT", + "REDIRECT_HINT", + "STATE_NAMESPACE", + "CitationPolicyCapability", + "CitationPolicyState", + "create_capability", +] diff --git a/tests/capabilities/test_citation_policy.py b/tests/capabilities/test_citation_policy.py new file mode 100644 index 00000000..0c43afa6 --- /dev/null +++ b/tests/capabilities/test_citation_policy.py @@ -0,0 +1,282 @@ +from dataclasses import dataclass, field +from typing import Any, cast +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, ToolCallPart +from pydantic_ai.models.function import FunctionModel + +from haiku.rag.capabilities.analysis import create_capability as create_analysis +from haiku.rag.capabilities.policy import ( + REDIRECT_HINT, + CitationPolicyState, +) +from haiku.rag.capabilities.policy import ( + create_capability as create_policy, +) +from haiku.rag.capabilities.rag import RAGCapability +from haiku.rag.capabilities.rag import create_capability as create_rag +from haiku.rag.config.models import AppConfig +from haiku.rag.store.models.chunk import SearchResult + + +@dataclass +class Deps: + state: dict[str, Any] = field(default_factory=dict) + + +async def stub_search(self, query: str, _limit: int | None) -> str: + cast(Any, self.state).searches[query] = [ + SearchResult(content="evidence", score=1.0, chunk_id="chunk-1") + ] + self._note_evidence() + return "EVIDENCE" + + +def prompts_of(messages) -> list[str]: + return [ + str(part.content) + for message in messages + for part in message.parts + if type(part).__name__ == "UserPromptPart" + ] + + +async def run_with_policy(temp_db_path, responses, *, policy=True, config=None): + """Answer one question with the given model responses, policy optional.""" + rag = create_rag( + db_path=temp_db_path, config=config or AppConfig(), defer_loading=False + ) + capabilities: list[Any] = [rag] + if policy: + capabilities.append(create_policy()) + turns = iter(responses) + sent: list[list[Any]] = [] + + async def model(messages, _info): + sent.append(list(messages)) + return ModelResponse(parts=next(turns)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=capabilities) + deps = Deps() + with patch.object(RAGCapability, "_search", stub_search): + result = await agent.run("what does the supervisor do?", deps=deps) + return result, deps, sent + + +@pytest.mark.asyncio +async def test_an_answer_without_a_citation_is_sent_back_once(temp_db_path): + """The last response of a question is the last moment to notice.""" + result, deps, sent = await run_with_policy( + temp_db_path, + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [TextPart("an answer with no citation")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")], + [TextPart("an answer with no citation")], + ], + ) + + redirects = [prompt for prompt in prompts_of(sent[-1]) if REDIRECT_HINT in prompt] + assert len(redirects) == 1 + assert deps.state["rag"]["citations"] == ["chunk-1"] + assert result.output == "an answer with no citation" + + +@pytest.mark.asyncio +async def test_a_grounded_answer_is_left_alone(temp_db_path): + _, _, sent = await run_with_policy( + temp_db_path, + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")], + [TextPart("a grounded answer")], + ], + ) + + assert not [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p] + + +@pytest.mark.asyncio +async def test_an_explicitly_ungrounded_answer_is_left_alone(temp_db_path): + """Citing nothing is a declaration, not an omission.""" + _, deps, sent = await run_with_policy( + temp_db_path, + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ToolCallPart("rag_cite", {"chunk_ids": []}, "call-2")], + [TextPart("I cannot find this in the knowledge base")], + ], + ) + + assert not [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p] + assert deps.state["citation_policy"]["violations"] == [] + + +@pytest.mark.asyncio +async def test_a_question_that_gathered_no_evidence_is_left_alone(temp_db_path): + """Nothing was retrieved, so there is no grounding to declare.""" + _, _, sent = await run_with_policy(temp_db_path, [[TextPart("hello back")]]) + + assert not [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p] + + +@pytest.mark.asyncio +async def test_a_violation_is_recorded_when_the_cite_tool_is_gone(temp_db_path): + """Asking for a withdrawn tool costs the agent's unknown-tool retries.""" + with patch( + "haiku.rag.capabilities._base.RAGCapabilityBase.cite_available", + new_callable=lambda: property(lambda self: False), + ): + _, deps, sent = await run_with_policy( + temp_db_path, + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [TextPart("an answer with no citation")], + ], + ) + + assert not [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p] + assert deps.state["citation_policy"]["violations"] == [0] + + +@pytest.mark.asyncio +async def test_without_the_policy_capability_nothing_is_enforced(temp_db_path): + """Omission is the switch, so there is no flag to test.""" + _, deps, sent = await run_with_policy( + temp_db_path, + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [TextPart("an answer with no citation")], + ], + policy=False, + ) + + assert not [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p] + assert "citation_policy" not in deps.state + + +@pytest.mark.asyncio +async def test_one_decision_is_made_with_both_evidence_capabilities(temp_db_path): + """Two capabilities must not each demand a citation for one answer.""" + 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 + ) + turns = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [TextPart("an answer with no citation")], + [TextPart("an answer with no citation")], + ] + ) + sent: list[list[Any]] = [] + + async def model(messages, _info): + sent.append(list(messages)) + return ModelResponse(parts=next(turns)) + + agent = Agent( + FunctionModel(model), + deps_type=Deps, + capabilities=[rag, analysis, create_policy()], + ) + with patch.object(RAGCapability, "_search", stub_search): + await agent.run("what does the supervisor do?", deps=Deps()) + + assert len([p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]) == 1 + + +def test_two_policy_capabilities_fail_fast(temp_db_path): + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, _info): # pragma: no cover - never reached + return ModelResponse(parts=[TextPart("answer")]) + + with pytest.raises(UserError, match="unique within a run"): + Agent( + FunctionModel(model), + deps_type=Deps, + capabilities=[rag, create_policy(), create_policy()], + ) + + +def test_the_policy_state_round_trips(): + state = CitationPolicyState(violations=[4, 12]) + + restored = CitationPolicyState.model_validate(state.model_dump(mode="json")) + + assert restored.violations == [4, 12] + + +@pytest.mark.asyncio +async def test_a_second_question_can_be_redirected_again(temp_db_path): + """The redirect fires once per question, not once per conversation.""" + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + turns = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [TextPart("first, uncited")], + [TextPart("first, uncited")], + [ToolCallPart("rag_search", {"query": "again"}, "call-2")], + [TextPart("second, uncited")], + [TextPart("second, uncited")], + ] + ) + sent: list[list[Any]] = [] + + async def model(messages, _info): + sent.append(list(messages)) + return ModelResponse(parts=next(turns)) + + agent = Agent( + FunctionModel(model), deps_type=Deps, capabilities=[rag, create_policy()] + ) + deps = Deps() + with patch.object(RAGCapability, "_search", stub_search): + first = await agent.run("what does the supervisor do?", deps=deps) + await agent.run( + "and who supervises them?", deps=deps, message_history=first.all_messages() + ) + + assert len([p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]) == 2 + + +@dataclass +class StatelessDeps: + """A host that keeps no capability state, which is allowed.""" + + +@pytest.mark.asyncio +async def test_a_violation_with_nowhere_to_record_it_does_not_fail_the_run( + temp_db_path, +): + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + turns = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [TextPart("an answer with no citation")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(turns)) + + agent = Agent( + FunctionModel(model), + deps_type=StatelessDeps, + capabilities=[rag, create_policy()], + ) + + with ( + patch( + "haiku.rag.capabilities._base.RAGCapabilityBase.cite_available", + new_callable=lambda: property(lambda self: False), + ), + patch.object(RAGCapability, "_search", stub_search), + ): + result = await agent.run("what does the supervisor do?", deps=StatelessDeps()) + + assert result.output == "an answer with no citation" diff --git a/tests/capabilities/test_evidence_capsule.py b/tests/capabilities/test_evidence_capsule.py index 6f1a51b5..981db4c3 100644 --- a/tests/capabilities/test_evidence_capsule.py +++ b/tests/capabilities/test_evidence_capsule.py @@ -11,12 +11,12 @@ 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.evidence import DiscoveredEvidence, discover_evidence from haiku.rag.capabilities.ledger import ( CapabilityEvidenceRecord, EvidenceOccurrence, @@ -77,6 +77,7 @@ def discovered( for chunk_id in cited }, tool_names=frozenset({f"{capability}_search"}), + cite_available=True, ) @@ -279,7 +280,7 @@ def _spy_discovery(found: list[list[DiscoveredEvidence]]): async def spy(self, ctx): await original(self, ctx) - found.append(self.discover(ctx)) + found.append(discover_evidence(ctx)) return patch.object(EvidenceCompactionCapability, "before_run", spy)