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

Require an answer to declare what grounds it, once per question

Browse source 
`CitationPolicyCapability` makes the single enforcement decision, whatever mix of
evidence capabilities is registered: two of them must not each demand a citation for
one answer. It decides in `after_model_request`, when a response carries no tool calls
and the question can still be redirected. An explicitly ungrounded answer is a
declaration and is left alone; a question that gathered no evidence is left alone too,
read from the ledger rather than from a searches dict that a new question clears. When
the cite tool is already withdrawn the question is recorded in
`CitationPolicyState.violations` instead of pointing the model at a tool that is gone.

Registering it is the only switch. `DiscoveredEvidence` and discovery move to
`capabilities.evidence` so both optional capabilities share them, and `cite_available`
joins `evidence_tool_names` as public for the same reason.

Measured on Qwen3.6-35B over two arms of 37 questions, 29 of them unanswerable from
the corpus: explicit ungrounded declarations rose from 23 to 26, grounded answers to
unanswerable questions fell from 4 to 2, answerable questions stayed at 8 of 8, and
one redirect fired in the whole arm. Reading every answer found no invented grounding.
This commit is contained in:
Yiorgis Gozadinos 2026-08-12 09:32:53 +03:00
parent e1e7936d15
commit 0f38417c60
No known key found for this signature in database
7 changed files with 538 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
CHANGELOG.md
View file

@ -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.

31
docs/capabilities/index.md
View file

@ -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.

51
haiku_rag_slim/haiku/rag/capabilities/compaction.py
View file

@ -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",

63
haiku_rag_slim/haiku/rag/capabilities/evidence.py Normal file
View file

@ -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",
]

147
haiku_rag_slim/haiku/rag/capabilities/policy.py Normal file
View file

@ -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",
]

282
tests/capabilities/test_citation_policy.py Normal file
View file

@ -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"

5
tests/capabilities/test_evidence_capsule.py
View file

@ -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)