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

Record what each evidence capability retrieved and cited

Browse source 
`CapabilityEvidenceRecord` holds the relationships a transcript cannot express:
which chunks a capability retrieved, which it cited, in which questions, and at
which point in the conversation. RAG and analysis each own one in their own state
namespace. Nothing is co-written: the host's state is JSON storage, so a shared
record would be overwritten by whichever capability synced last, and merging
happens in transient per-request views instead.

Both clocks are derived from the conversation rather than counted locally, so
every participant computes the same values without sharing a counter. Question
identity is the message count when the question arrived; epoch is the message
count at an outcome. Epochs are therefore globally comparable, which is what
lets `citation_status` require a declaration to follow the newest evidence of
every capability, and what makes equal epochs mean one request.

A declaration is written only after `resolve_citations` succeeds, so a call
naming only unresolvable ids is not a citation. Status is derived, never stored,
so refs and status cannot contradict.

Resuming a question requires the host to carry the capability state from the run
being resumed. Without it the identity of the question in progress is unknowable,
and adopting the current message count would relabel that question as a new one
and judge every declaration in it against the wrong identity.

Nothing reads the records yet and no wire behaviour changes.
This commit is contained in:
Yiorgis Gozadinos 2026-08-10 15:09:53 +03:00
parent fa0c4a4f50
commit 85594a1fd1
No known key found for this signature in database
7 changed files with 722 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
CHANGELOG.md
View file

@ -1,6 +1,14 @@
# Changelog
## [Unreleased]
### Added
- `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.
### Changed
- 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.
### Fixed
- `cross-encoder` reranking no longer ties the scores of strongly-relevant candidates, which left their order to the sort. Scores remain 0-1.

64
haiku_rag_slim/haiku/rag/capabilities/_base.py
View file

@ -23,6 +23,7 @@ from pydantic_ai.tools import ToolDefinition
from pydantic_ai.toolsets import AgentToolset
from haiku.rag.capabilities._tools import CodeExecutionEntry, search_corpus
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord, EvidenceRef
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
@ -178,6 +179,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
request_count: int = field(default=0, repr=False)
grace_requests_used: int = field(default=0, repr=False)
turn_start: int = field(default=0, repr=False)
epoch: int = field(default=0, repr=False)
async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]":
"""Start a run's own copy, and decide what counts as an earlier question.
@ -187,12 +189,31 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
messages and replace its evidence with the earlier-question notice, leaving
the model to answer with the evidence taken away. Failing this way costs a
larger request; failing the other way costs the answer.
A new question takes the message count as its identity, which every
participant derives identically from the same history. A resumption keeps
the identity already recorded: the question is the one in progress, and
adopting the current count would relabel it as a new one and judge its
declarations against the wrong question. A resumption with no recorded
identity is a state this design does not produce, so it is reported rather
than guessed at — unless there is no history at all, where an absent prompt
means an instructions-only first question and nothing is in progress.
"""
outer = getattr(ctx.deps, "state", None)
outer_state = outer if isinstance(outer, dict) else None
raw_state = outer_state.get(self.state_namespace) if outer_state else None
resuming = _is_resumption(ctx.prompt, ctx.messages)
if resuming and ctx.messages and not (raw_state or {}).get("evidence"):
raise RuntimeError(
f"The {self.state_namespace} capability is resuming a question with "
"no stored question identity. Capabilities cannot be added, removed "
"or migrated while a question is unfinished, and the run's state "
"must be carried between its runs."
)
state = self.state_type.model_validate(raw_state or {})
_clear_invocation_state(state)
if not resuming:
cast(Any, state).evidence.question = len(ctx.messages)
run_capability = replace(
self,
state=state,
@ -203,9 +224,8 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
search_count=0,
request_count=0,
grace_requests_used=0,
turn_start=(
0 if _is_resumption(ctx.prompt, ctx.messages) else len(ctx.messages)
),
epoch=0,
turn_start=0 if resuming else len(ctx.messages),
)
run_capability._sync_state()
return run_capability
@ -238,6 +258,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
async def before_model_request(
self, ctx: RunContext[Any], request_context: ModelRequestContext
) -> ModelRequestContext:
self.epoch = len(ctx.messages)
if instruction := self._budget_notice():
current_request = request_context.messages[-1]
if isinstance(current_request, ModelRequest):
@ -390,6 +411,41 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
finally:
self._sync_state()
def _evidence_record(self) -> CapabilityEvidenceRecord:
assert self.state is not None
return cast(CapabilityEvidenceRecord, cast(Any, self.state).evidence)
def _note_evidence(self) -> None:
"""Record an outcome the model can ground an answer on.
Includes an empty search result and a failed execution that still printed
output: negative evidence grounds a refusal. Excludes a spent budget, which
yields nothing to ground anything on.
"""
self._evidence_record().note_evidence(self.epoch)
def _declare(self, citations: list[Citation]) -> None:
"""Record what the model cited, once the ids have resolved.
Declaring earlier would let a call naming only unresolvable ids read as a
grounded answer.
"""
state = cast(Any, self.state)
retrieved = {
result.chunk_id
for results in state.searches.values()
for result in results
if result.chunk_id
}
self._evidence_record().declare(
[
EvidenceRef(capability=self.state_namespace, chunk_id=c.chunk_id)
for c in citations
],
epoch=self.epoch,
retrieved_now=retrieved,
)
async def _search(self, query: str, limit: int | None) -> str | ToolReturn:
assert self.state is not None
self.search_count += 1
@ -407,6 +463,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
)
state = cast(Any, self.state)
state.searches[query] = results
self._note_evidence()
if self.vision and (parts := build_image_content_from_results(results)):
return ToolReturn(return_value=formatted, content=parts)
return formatted
@ -454,6 +511,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
"Copy chunk_ids verbatim from search results."
)
self._register_citations(citations)
self._declare(citations)
resolved = {citation.chunk_id for citation in citations}
unresolved = [cid for cid in missing if cid not in resolved]
if unresolved:

3
haiku_rag_slim/haiku/rag/capabilities/analysis.py
View file

@ -13,6 +13,7 @@ from haiku.rag.capabilities._base import (
RAGCapabilityBase,
resolve_db_path,
)
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
from haiku.rag.config.models import AppConfig
from haiku.rag.sandbox import AnalysisContext, Sandbox
from haiku.rag.store.models.chunk import SearchResult
@ -29,6 +30,7 @@ class AnalysisState(BaseModel):
executions: list[CodeExecutionEntry] = Field(default_factory=list)
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[str] = Field(default_factory=list)
evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord)
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
@ -107,6 +109,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
)
sandbox = await self._ensure_sandbox()
result = await sandbox.execute(code)
self._note_evidence()
if sandbox._search_results:
existing = self.state.searches.get("_sandbox", [])
seen = {item.chunk_id for item in existing}

155
haiku_rag_slim/haiku/rag/capabilities/ledger.py Normal file
View file

@ -0,0 +1,155 @@
from collections.abc import Iterable
from typing import Literal
from pydantic import BaseModel, Field
CitationStatus = Literal["missing", "grounded", "ungrounded"]
class EvidenceRef(BaseModel):
"""One piece of evidence, identified by its owner as well as its chunk.
A chunk id alone is not an identity: the same id can be reported by more than
one capability, and ownership is what tells compaction whose output it may
touch.
"""
capability: str
chunk_id: str
class EvidenceOccurrence(BaseModel):
"""Which questions retrieved a piece of evidence, and which cited it."""
capability: str
chunk_id: str
retrieved_in_questions: list[int] = Field(default_factory=list)
cited_in_questions: list[int] = Field(default_factory=list)
class CitationDeclaration(BaseModel):
"""What a question declared as its grounding, and when.
Bound to a question *and* an epoch: the epoch outlives a question, so matching
it alone would let a question that gathered no evidence inherit the previous
declaration and read as compliant having declared nothing.
"""
question: int
epoch: int
refs: list[EvidenceRef] = Field(default_factory=list)
class CapabilityEvidenceRecord(BaseModel):
"""What one evidence capability wrote, in its own state namespace.
Holds no content: ``Citation`` in ``citation_index`` is the canonical record
and already persists content, document id and picture refs. This is the index
over it that compaction needs and the transcript cannot provide.
Single-writer by construction. A record shared between capabilities would be
overwritten by whichever of them synced its state last; merging happens in the
transient views built by ``citation_status`` and the optional capabilities.
``question`` is the number of messages that existed when the question arrived,
and ``epoch`` the number when an outcome occurred. Both are derived from the
conversation rather than counted locally, so every participant computes the
same values without sharing a counter.
"""
occurrences: dict[str, EvidenceOccurrence] = Field(default_factory=dict)
question: int = 0
latest_evidence_epoch: int = 0
declaration: CitationDeclaration | None = None
def note_evidence(self, epoch: int) -> None:
"""Record that the model has seen an evidence outcome.
Called for anything an answer could rest on, including a search that
returned nothing and a failed execution that still printed output — a
fruitless search grounds a refusal. Not called for a failure that yields
no evidence at all, such as an exhausted budget.
"""
self.latest_evidence_epoch = max(self.latest_evidence_epoch, epoch)
def declare(
self,
refs: list[EvidenceRef],
*,
epoch: int,
retrieved_now: set[str] | None = None,
) -> None:
"""Record validated citations for the current question.
Repeated calls at the same epoch merge, so citing again cannot narrow what
was already declared: an empty call after a grounded one leaves it
grounded. A call at a later epoch declares afresh, because evidence the
model saw in between may be what it is now citing.
"""
current = self.declaration
if current is not None and (current.question, current.epoch) == (
self.question,
epoch,
):
known = {(ref.capability, ref.chunk_id) for ref in current.refs}
current.refs.extend(
ref for ref in refs if (ref.capability, ref.chunk_id) not in known
)
else:
self.declaration = CitationDeclaration(
question=self.question, epoch=epoch, refs=list(refs)
)
for ref in refs:
occurrence = self.occurrences.setdefault(
ref.chunk_id,
EvidenceOccurrence(capability=ref.capability, chunk_id=ref.chunk_id),
)
if self.question not in occurrence.cited_in_questions:
occurrence.cited_in_questions.append(self.question)
if (
retrieved_now
and ref.chunk_id in retrieved_now
and self.question not in occurrence.retrieved_in_questions
):
occurrence.retrieved_in_questions.append(self.question)
def citation_status(
records: Iterable[CapabilityEvidenceRecord], *, question: int
) -> CitationStatus:
"""Derived, never stored, so refs and status cannot contradict.
A declaration is current only for the question it was made in, and only if it
followed the newest evidence outcome of *every* capability: a question where
one capability cited and another then searched without citing is not grounded.
Strictly later, since a citation made in the same request as an outcome cannot
have read it.
A grounding *violation* is not one of these: that is an enforcement outcome
recorded by the policy capability, not something the model declared.
"""
records = list(records)
horizon = max((record.latest_evidence_epoch for record in records), default=0)
current = [
record.declaration
for record in records
if record.declaration is not None
and record.declaration.question == question
and record.declaration.epoch > horizon
]
if not current:
return "missing"
return (
"grounded" if any(declaration.refs for declaration in current) else "ungrounded"
)
__all__ = [
"CapabilityEvidenceRecord",
"CitationDeclaration",
"CitationStatus",
"EvidenceOccurrence",
"EvidenceRef",
"citation_status",
]

2
haiku_rag_slim/haiku/rag/capabilities/rag.py
View file

@ -12,6 +12,7 @@ from haiku.rag.capabilities._base import (
RAGCapabilityBase,
resolve_db_path,
)
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import Citation
@ -33,6 +34,7 @@ _instructions_path = Path(__file__).parent / "instructions" / "rag.md"
class RAGState(BaseModel):
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[str] = Field(default_factory=list)
evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord)
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)

346
tests/capabilities/test_capabilities.py
View file

@ -26,6 +26,10 @@ from haiku.rag.capabilities._base import (
)
from haiku.rag.capabilities.analysis import AnalysisCapability, AnalysisState
from haiku.rag.capabilities.analysis import create_capability as create_analysis
from haiku.rag.capabilities.ledger import (
CapabilityEvidenceRecord,
citation_status,
)
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGCapability, RAGState
from haiku.rag.capabilities.rag import create_capability as create_rag
from haiku.rag.config.models import AppConfig, PromptsConfig
@ -463,6 +467,44 @@ async def test_analysis_execution_limit_fails_the_tool(temp_db_path):
await capability._execute_code("print('done')")
@pytest.mark.asyncio
async def test_a_spent_execution_budget_is_not_evidence(temp_db_path):
"""Nothing was produced to ground an answer on, so nothing is recorded."""
config = AppConfig()
config.analysis.max_executions = 0
capability = create_analysis(db_path=temp_db_path, config=config)
capability.state = AnalysisState()
capability.epoch = 5
with pytest.raises(ToolFailed):
await capability._execute_code("print('done')")
assert capability.state.evidence.latest_evidence_epoch == 0
@pytest.mark.parametrize("success", [True, False])
@pytest.mark.asyncio
async def test_a_code_execution_is_evidence_even_when_it_fails(temp_db_path, success):
"""Output the model can read grounds an answer, whether the code raised or not."""
capability = create_analysis(db_path=temp_db_path, config=AppConfig())
capability.state = AnalysisState()
capability.epoch = 5
sandbox = AsyncMock(spec=Sandbox)
sandbox._search_results = []
sandbox.execute.return_value = SandboxResult(
stdout="42", stderr="" if success else "boom", success=success
)
capability.sandbox = sandbox
if success:
await capability._execute_code("print(42)")
else:
with pytest.raises(ToolFailed):
await capability._execute_code("print(42)")
assert capability.state.evidence.latest_evidence_epoch == 5
@pytest.mark.asyncio
async def test_spent_search_budget_is_announced_but_keeps_the_tool(rag_db):
"""A spent budget is announced; the tool stays declared to avoid a dead run.
@ -965,6 +1007,22 @@ async def test_compaction_never_reaches_the_stored_message_history(temp_db_path)
assert PRIOR_TURN_NOTICE not in returns
def _resuming_deps() -> Deps:
"""State as a resumption always finds it: the question already identified.
A run that resumes has been through ``for_run`` before, so the identity of the
question in progress is stored. Fabricating the history without it is a state
the design does not produce, and is rejected rather than guessed at.
"""
return Deps(
state={
"rag": RAGState(evidence=CapabilityEvidenceRecord(question=0)).model_dump(
mode="json"
)
}
)
def _in_flight_history() -> list[Any]:
"""A question already asked and searched, still awaiting its answer."""
return [
@ -1033,7 +1091,9 @@ async def test_a_prompt_on_an_unanswered_tail_is_treated_as_a_continuation(
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability])
await agent.run(
"a different question", deps=Deps(), message_history=_in_flight_history()
"a different question",
deps=_resuming_deps(),
message_history=_in_flight_history(),
)
assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"]
@ -1066,7 +1126,7 @@ async def test_a_resume_carrying_a_prompt_keeps_the_active_evidence(temp_db_path
"carry on",
deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}),
message_history=history,
deps=Deps(),
deps=_resuming_deps(),
)
assert "EVIDENCE FOR THE LIVE TURN" in _wire_returns(wire[-1])
@ -1105,7 +1165,9 @@ async def test_a_resumed_run_keeps_the_active_questions_evidence(
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability])
await agent.run(deps=Deps(), message_history=_in_flight_history(), **resume_kwargs)
await agent.run(
deps=_resuming_deps(), message_history=_in_flight_history(), **resume_kwargs
)
assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"]
@ -1161,3 +1223,281 @@ async def test_prior_turn_search_is_compacted_but_its_cite_receipt_is_not(temp_d
}
assert prior["rag_search"] == PRIOR_TURN_NOTICE
assert prior["rag_cite"] != PRIOR_TURN_NOTICE
def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord:
return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"])
async def _stub_search(self, query: str, _limit: int | None) -> str:
"""Record a result the way the real search does, so citing resolves."""
cast(Any, self.state).searches[query] = [
SearchResult(content="evidence", score=1.0, chunk_id="chunk-1")
]
self._note_evidence()
return "EVIDENCE"
@pytest.mark.asyncio
async def test_a_question_takes_its_own_identity_and_both_capabilities_agree(
temp_db_path,
):
"""Identity is derived from the conversation, so no counter is shared."""
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
)
async def model(_messages, _info):
return ModelResponse(parts=[TextPart("answer")])
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, analysis])
deps = Deps()
first = await agent.run("first question", deps=deps)
first_identity = _record(deps, "rag").question
await agent.run("second question", deps=deps, message_history=first.all_messages())
assert first_identity == 0
assert _record(deps, "rag").question > first_identity
assert _record(deps, "analysis").question == _record(deps, "rag").question
@pytest.mark.asyncio
async def test_a_resumption_keeps_the_identity_of_the_question_in_progress(
temp_db_path,
):
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
async def model(_messages, _info):
return ModelResponse(parts=[TextPart("answer")])
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
deps = Deps(
state={
"rag": RAGState(evidence=CapabilityEvidenceRecord(question=7)).model_dump(
mode="json"
)
}
)
history = [
*_in_flight_history(),
ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]),
]
await agent.run(
"carry on",
deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}),
message_history=history,
deps=deps,
)
assert _record(deps, "rag").question == 7
@pytest.mark.asyncio
async def test_resuming_without_a_stored_identity_fails_instead_of_guessing(
temp_db_path,
):
"""Adopting the message count would relabel a question already in progress.
Every declaration and epoch comparison in it would then be judged against the
wrong question, silently. This state is not one the design produces, so it is
reported rather than repaired.
"""
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")])
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
with pytest.raises(RuntimeError, match="no stored question identity"):
await agent.run(
"carry on",
deferred_tool_results=DeferredToolResults(
calls={"call-2": "external result"}
),
message_history=[
*_in_flight_history(),
ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]),
],
deps=Deps(),
)
@pytest.mark.asyncio
async def test_citing_after_searching_grounds_the_question(temp_db_path):
"""The whole rule, end to end, with no compactor and no policy capability."""
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
calls = iter(
[
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
[TextPart("answer")],
]
)
async def model(_messages, _info):
return ModelResponse(parts=next(calls))
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
deps = Deps()
with patch.object(RAGCapability, "_search", _stub_search):
await agent.run("what does the supervisor do?", deps=deps)
record = _record(deps, "rag")
assert record.declaration is not None
assert [ref.chunk_id for ref in record.declaration.refs] == ["chunk-1"]
assert record.occurrences["chunk-1"].retrieved_in_questions == [record.question]
assert citation_status([record], question=record.question) == "grounded"
@pytest.mark.asyncio
async def test_searching_after_citing_leaves_the_question_uncited(temp_db_path):
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
calls = iter(
[
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
[ToolCallPart("rag_search", {"query": "again"}, "call-3")],
[TextPart("answer")],
]
)
async def model(_messages, _info):
return ModelResponse(parts=next(calls))
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
deps = Deps()
with patch.object(RAGCapability, "_search", _stub_search):
await agent.run("what does the supervisor do?", deps=deps)
record = _record(deps, "rag")
assert record.declaration is not None
assert citation_status([record], question=record.question) == "missing"
@pytest.mark.asyncio
async def test_a_citation_in_the_same_request_as_its_search_is_not_current(
temp_db_path,
):
"""Two calls in one response share an epoch, and citing must follow seeing."""
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
calls = iter(
[
[
ToolCallPart("rag_search", {"query": "supervisor"}, "call-1"),
ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2"),
],
[TextPart("answer")],
]
)
async def model(_messages, _info):
return ModelResponse(parts=next(calls))
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
deps = Deps()
with patch.object(RAGCapability, "_search", _stub_search):
await agent.run("what does the supervisor do?", deps=deps)
record = _record(deps, "rag")
assert record.declaration is not None
assert record.declaration.epoch == record.latest_evidence_epoch
assert citation_status([record], question=record.question) == "missing"
@pytest.mark.asyncio
async def test_evidence_cited_in_two_questions_keeps_both_in_the_record(temp_db_path):
"""Occurrences outlive the question that wrote them, through the state dict."""
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
calls = iter(
[
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
[TextPart("first answer")],
[ToolCallPart("rag_search", {"query": "supervisor again"}, "call-3")],
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-4")],
[TextPart("second answer")],
]
)
async def model(_messages, _info):
return ModelResponse(parts=next(calls))
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
deps = Deps()
with patch.object(RAGCapability, "_search", _stub_search):
first = await agent.run("who supervises?", deps=deps)
first_question = _record(deps, "rag").question
await agent.run(
"and who supervises them?", deps=deps, message_history=first.all_messages()
)
record = _record(deps, "rag")
assert record.occurrences["chunk-1"].cited_in_questions == [
first_question,
record.question,
]
@pytest.mark.asyncio
async def test_a_run_with_no_prompt_and_no_history_starts_a_question(temp_db_path):
"""An instructions-only run is a first question, not a resumption.
There is no question in progress to keep an identity for, so nothing is
missing and the run proceeds with a fresh one.
"""
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
async def model(_messages, _info):
return ModelResponse(parts=[TextPart("answer")])
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
deps = Deps()
await agent.run(deps=deps)
assert _record(deps, "rag").question == 0
@pytest.mark.asyncio
async def test_citing_without_searching_grounds_the_question(temp_db_path):
"""A direct chunk-id citation stands on its own, with no evidence outcome.
Epochs count messages and so start above zero, which is what lets a
declaration made in the first request still beat an empty evidence horizon.
"""
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
calls = iter(
[
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-1")],
[TextPart("answer")],
]
)
async def model(_messages, _info):
return ModelResponse(parts=next(calls))
client = AsyncMock()
client.get_chunk_by_id.return_value = Chunk(
id="chunk-1", document_id="doc-1", content="evidence"
)
client.get_document_by_id.return_value = None
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
deps = Deps()
with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=client)):
await agent.run("cite chunk-1", deps=deps)
record = _record(deps, "rag")
assert record.latest_evidence_epoch == 0
assert record.declaration is not None
assert record.declaration.epoch > 0
assert record.occurrences["chunk-1"].retrieved_in_questions == []
assert citation_status([record], question=record.question) == "grounded"

150
tests/capabilities/test_evidence_ledger.py Normal file
View file

@ -0,0 +1,150 @@
from haiku.rag.capabilities.ledger import (
CapabilityEvidenceRecord,
CitationDeclaration,
EvidenceRef,
citation_status,
)
def rag_ref(chunk_id: str = "c1") -> EvidenceRef:
return EvidenceRef(capability="rag", chunk_id=chunk_id)
def test_a_record_survives_the_state_round_trip():
"""Capability state is persisted as JSON, so the schema must survive it.
A dict keyed by ``(capability, chunk_id)`` does not: the key serialises to
``"rag,c1"`` and fails revalidation as a tuple.
"""
record = CapabilityEvidenceRecord(question=4)
record.note_evidence(5)
record.declare([rag_ref()], epoch=7, retrieved_now={"c1"})
restored = CapabilityEvidenceRecord.model_validate(record.model_dump(mode="json"))
assert restored == record
assert citation_status([restored], question=4) == "grounded"
assert restored.occurrences["c1"].cited_in_questions == [4]
assert restored.occurrences["c1"].retrieved_in_questions == [4]
def test_no_declaration_reads_as_missing():
assert citation_status([CapabilityEvidenceRecord()], question=0) == "missing"
assert citation_status([], question=0) == "missing"
def test_refs_make_it_grounded_and_no_refs_make_it_ungrounded():
grounded = CapabilityEvidenceRecord()
grounded.declare([rag_ref()], epoch=1)
ungrounded = CapabilityEvidenceRecord()
ungrounded.declare([], epoch=1)
assert citation_status([grounded], question=0) == "grounded"
assert citation_status([ungrounded], question=0) == "ungrounded"
def test_an_earlier_questions_declaration_is_never_current():
"""Epochs outlive a question, so the epoch alone would inherit it."""
record = CapabilityEvidenceRecord(question=2)
record.declare([rag_ref()], epoch=3)
assert citation_status([record], question=2) == "grounded"
record.question = 8
assert record.declaration is not None
assert citation_status([record], question=8) == "missing"
def test_a_citation_in_the_same_request_as_the_evidence_is_not_current():
"""Citing must follow seeing: equal epochs mean one request."""
record = CapabilityEvidenceRecord()
record.note_evidence(5)
record.declare([rag_ref()], epoch=5)
assert citation_status([record], question=0) == "missing"
record.declare([rag_ref()], epoch=7)
assert citation_status([record], question=0) == "grounded"
def test_evidence_from_another_capability_after_citing_makes_it_uncited():
"""Currency spans capabilities, which only works because epochs are global."""
cited = CapabilityEvidenceRecord()
cited.note_evidence(3)
cited.declare([rag_ref()], epoch=5)
searched_after = CapabilityEvidenceRecord()
searched_after.note_evidence(7)
assert citation_status([cited], question=0) == "grounded"
assert citation_status([cited, searched_after], question=0) == "missing"
def test_declarations_at_the_same_epoch_merge():
record = CapabilityEvidenceRecord()
record.declare([rag_ref("c1")], epoch=3)
record.declare([rag_ref("c2")], epoch=3)
assert record.declaration is not None
assert [ref.chunk_id for ref in record.declaration.refs] == ["c1", "c2"]
def test_repeating_a_ref_at_the_same_epoch_does_not_duplicate_it():
record = CapabilityEvidenceRecord()
record.declare([rag_ref()], epoch=3)
record.declare([rag_ref()], epoch=3)
assert record.declaration is not None
assert len(record.declaration.refs) == 1
def test_neither_cite_order_downgrades_a_grounded_declaration():
grounded_then_empty = CapabilityEvidenceRecord()
grounded_then_empty.declare([rag_ref()], epoch=3)
grounded_then_empty.declare([], epoch=3)
empty_then_grounded = CapabilityEvidenceRecord()
empty_then_grounded.declare([], epoch=3)
empty_then_grounded.declare([rag_ref()], epoch=3)
assert citation_status([grounded_then_empty], question=0) == "grounded"
assert citation_status([empty_then_grounded], question=0) == "grounded"
def test_the_same_chunk_id_under_two_capabilities_stays_separate():
rag = CapabilityEvidenceRecord()
rag.declare([EvidenceRef(capability="rag", chunk_id="shared")], epoch=3)
analysis = CapabilityEvidenceRecord()
analysis.declare([EvidenceRef(capability="analysis", chunk_id="shared")], epoch=3)
assert rag.occurrences["shared"].capability == "rag"
assert analysis.occurrences["shared"].capability == "analysis"
def test_an_evidence_epoch_never_moves_backwards():
record = CapabilityEvidenceRecord()
record.note_evidence(9)
record.note_evidence(4)
assert record.latest_evidence_epoch == 9
def test_citing_the_same_chunk_in_two_questions_records_both():
record = CapabilityEvidenceRecord(question=2)
record.declare([rag_ref()], epoch=3, retrieved_now={"c1"})
record.question = 8
record.declare([rag_ref()], epoch=9)
occurrence = record.occurrences["c1"]
assert occurrence.cited_in_questions == [2, 8]
assert occurrence.retrieved_in_questions == [2]
def test_a_declaration_records_the_question_and_epoch_it_was_made_at():
record = CapabilityEvidenceRecord(question=6)
record.declare([rag_ref()], epoch=11)
assert record.declaration == CitationDeclaration(
question=6, epoch=11, refs=[rag_ref()]
)