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

Merge pull request #538 from ggozad/feat/citation-policy

Browse source 
Citation policy: require a declaration, accept an empty one
This commit is contained in:
Yiorgis Gozadinos 2026-08-13 13:55:19 +03:00 committed by GitHub
commit 0a0ddbd3b6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 961 additions and 67 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, in any conversation that has something to declare — this question retrieved evidence, or something was cited earlier. A question that ends undeclared is sent back once to record what grounded the answer already given, and is recorded in `CitationPolicyState.violations` if it finishes undeclared regardless. A conversation with neither a current-question evidence outcome nor any earlier citation is not enforced.
- `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.

39
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,44 @@ 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, or the question finishes undeclared anyway, it is recorded as
a violation in `CitationPolicyState` under `"citation_policy"`, since pointing a model
at a tool that is gone costs it retries.
What gets enforced is every answer in a conversation that has something to declare:
either this question retrieved evidence, or the conversation has already cited
something, which stays available to later answers. So a follow-up about evidence cited
earlier is enforced even though it searched nothing — that case is the reason the
capability exists. It also means that once anything has been cited, later turns are
enforced too, a greeting included; the model satisfies the policy by citing an empty
list, at the cost of one extra request. A conversation with neither a current-question
evidence outcome nor any earlier citation is not enforced.
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.

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

@ -273,6 +273,16 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
if tool.capability_id != self.id or tool.name == self._cite_tool_name
]
@property
def cite_available(self) -> bool:
"""Whether this capability's cite tool is still declared to the model.
Public because the citation policy has to know whether asking for a
citation is even possible: past the grace window the tool is gone, and
pointing the model at it would cost the agent's unknown-tool retries.
"""
return not self._citation_grace_expired
def evidence_tool_names(self) -> set[str]:
"""Tools that can bring new evidence into the run.
@ -419,12 +429,17 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
return formatted
async def _cite(self, chunk_ids: list[str]) -> str:
"""Register the evidence behind this answer, or declare there is none.
An empty list is a valid answer to "what grounds this?", and the only way
the model can say "nothing" other than staying silent — which is
indistinguishable from forgetting to cite at all. It declares the question
ungrounded, which is not the same as leaving it undeclared.
"""
assert self.state is not None
if not chunk_ids:
raise ModelRetry(
"No citations registered: chunk_ids was empty. Pass the chunk_ids "
"you want to cite, copied verbatim from search results."
)
self._declare([])
return "Recorded: this answer cites no knowledge-base evidence."
all_results: list[SearchResult] = []
state = cast(Any, self.state)

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

10
haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md
View file

@ -3,9 +3,9 @@
You answer questions over a document knowledge base. Two common workflows:
- **`analysis_search → analysis_cite → answer`** when the answer is grounded on specific document content. Call `analysis_cite` with the supporting chunk_ids before writing the answer.
- **`analysis_execute_code → answer`** when the answer is a count, aggregation, listing, or structural computation over the corpus (e.g. "how many documents?", "average page count"). No `analysis_cite` is needed when no specific chunks support the answer.
- **`analysis_execute_code → answer`** when the answer is a count, aggregation, listing, or structural computation over the corpus (e.g. "how many documents?", "average page count"). Call `analysis_cite` with an empty list when no specific chunks support the answer.
You can mix the two. The rule: cite when grounded on retrieved evidence; don't fabricate citations for corpus-level computation.
You can mix the two. The rule: always call `analysis_cite` before answering — pass the grounding chunk_ids, or an empty list for a corpus-level computation. Never fabricate citations.
## Tools
@ -25,7 +25,7 @@ Search the knowledge base directly (outside code execution). Each result has a `
### analysis_cite
Register the chunk IDs that ground your answer. **You must call `analysis_cite` before writing any final answer that uses retrieved evidence — search results, items.jsonl rows, toc.json nodes, or content.txt content.** Skipping `analysis_cite` leaves the answer ungrounded and is treated as a failure.
`analysis_cite` is **not** required when your answer is a corpus-level computation that doesn't draw on specific chunks — counts, aggregations, listings, averages across documents. Don't fabricate citations for these.
When your answer is a corpus-level computation that doesn't draw on specific chunks — counts, aggregations, listings, averages across documents — call `analysis_cite` with an empty list. Don't fabricate citations for these.
Chunk IDs come from two places:
- The `chunk_id` field on `search` / `await search(...)` results
@ -104,7 +104,7 @@ The user may attach images to their question. An attached image is part of the q
4. For questions about a *known document's* structure ("which section contains X", "list the sections of doc Y", "summarise section Z"), read `/documents/{id}/toc.json` first. Each node carries `item_range` (a slice into `items.jsonl`) and `chunk_ids` (citable). Prefer this over `search()` for in-document navigation — `search()` ranks across the whole corpus and can return chunks from unrelated documents.
5. Before writing your final response, call `analysis_cite` with the chunk_ids that ground your answer.
You MUST call `analysis_cite` with at least one chunk ID before producing your final answer **when your answer is grounded on retrieved evidence**. Skip `analysis_cite` in two cases: (a) you are refusing for lack of information, or (b) your answer is a corpus-level computation (count, aggregation, listing) that doesn't draw on specific chunks. In those cases do **not** fabricate citations.
You MUST call `analysis_cite` before producing your final answer, every time, with no exceptions. Pass the chunk IDs that ground the answer, or an empty list when none do — because you are refusing for lack of information, or because the answer is a corpus-level computation. An answer not preceded by `analysis_cite` is a protocol violation.
## Important
@ -114,4 +114,4 @@ You MUST call `analysis_cite` with at least one chunk ID before producing your f
- Use `await` for all async functions inside `analysis_execute_code` (`search`, `list_documents`)
- Read files with `Path.read_text()` or `open()`/`with`. For lines use `.readlines()` or `.read().split("\n")`, never `for line in f`. The `collections` module is unavailable.
- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `analysis_cite` tool separately to register citations. `cite{...}` markdown-style inline references do nothing; only an actual `analysis_cite` tool call registers a citation.
- **Before you write your final answer, invoke the `analysis_cite` tool with the supporting chunk_ids.** This is the last tool call before answering whenever your answer draws on retrieved evidence.
- **Before you write your final answer, invoke the `analysis_cite` tool with the supporting chunk_ids, or with an empty list if there are none.** This is the last tool call before answering, every time.

6
haiku_rag_slim/haiku/rag/capabilities/instructions/rag.md
View file

@ -17,7 +17,7 @@ Each result includes:
When a result's Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text. Use the image directly to answer questions about figures, diagrams, charts, screenshots.
### rag_cite
Register the chunk IDs that ground your answer. Call this BEFORE writing your final answer, with the `chunk_id` values from search results that support each claim. Every answer that uses search results must be backed by `rag_cite`.
Register the chunk IDs that ground your answer. Call this BEFORE writing your final answer, with the `chunk_id` values from search results that support each claim. Every answer must be preceded by `rag_cite` — pass an empty list when nothing in the knowledge base grounds it.
Use chunk_ids exactly as they appear in the search response — copy the full UUID verbatim. Do not abbreviate, paraphrase, or reconstruct chunk_ids from memory; the tool matches them as opaque strings.
@ -33,7 +33,7 @@ The user may attach images to their question. An attached image is part of the q
4. Identify the chunk IDs that support your answer and call `rag_cite` with them
5. Then write a concise answer based strictly on the cited content
You MUST call `rag_cite` with at least one chunk ID before producing your final answer, **unless** you are refusing for lack of information (see below). Answers without citations are considered ungrounded.
You MUST call `rag_cite` before producing your final answer, every time, with no exceptions. Pass the chunk IDs that support the answer, or an empty list if none do. An answer not preceded by `rag_cite` is a protocol violation, not merely an ungrounded answer.
## Guidelines
@ -42,7 +42,7 @@ You MUST call `rag_cite` with at least one chunk ID before producing your final
- If multiple results are relevant, synthesize them coherently
- Be concise and direct — avoid elaboration unless asked
- If the search tool tells you the search limit is reached, stop searching and answer with what you have
- If the retrieved documents do not directly address the question, say: "I cannot find enough information in the knowledge base to answer this question." Do not guess or infer from tangentially related content. In this refusal case do **not** call `rag_cite` — there is nothing to cite.
- If the retrieved documents do not directly address the question, say: "I cannot find enough information in the knowledge base to answer this question." Do not guess or infer from tangentially related content. Refusing does not exempt you from `rag_cite` — call it with an empty list to record that nothing grounds the answer.
- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `rag_cite` tool separately to register citations.
## When search returns irrelevant results

17
haiku_rag_slim/haiku/rag/capabilities/ledger.py
View file

@ -133,10 +133,11 @@ class CapabilityEvidenceRecord(BaseModel):
) -> 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.
Citing again cannot narrow what a question already declared: calls merge
while no evidence outcome has followed the standing declaration, whatever
epoch they arrive at, so an empty second thought leaves a grounded question
grounded. Only genuinely newer evidence starts a declaration afresh, since
what the model saw in between may be what it is now citing.
"""
if self.question is None:
raise ValueError(
@ -145,14 +146,16 @@ class CapabilityEvidenceRecord(BaseModel):
)
self._reject_regression(epoch, "A declaration")
current = self.declaration
if current is not None and (current.question, current.epoch) == (
self.question,
epoch,
if (
current is not None
and current.question == self.question
and self.latest_evidence_epoch <= current.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
)
current.epoch = max(current.epoch, epoch)
else:
self.declaration = CitationDeclaration(
question=self.question, epoch=epoch, refs=list(refs)

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

@ -0,0 +1,233 @@
from dataclasses import dataclass
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 (
ModelMessage,
ModelResponse,
ToolCallPart,
UserPromptPart,
)
from pydantic_ai.models import ModelRequestContext
from pydantic_ai.run import AgentRunResult
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"
CITATION_REDIRECT_TAG = "[haiku.rag/citation-redirect]"
"""Tag the redirect carries, so a question can tell it has already been asked.
Not the wording: a user writing "record what grounded the answer you already gave"
in their own question would otherwise read as a redirect we had sent, silently
switching enforcement off for that question.
"""
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 "
f"exactly as you gave it. {CITATION_REDIRECT_TAG}"
)
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.
"""
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 that ends the question is the last opportunity: one carrying no
tool calls, or one whose call is an output tool, which is how a structured
answer arrives and which finishes the run just the same. 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.
Endings this cannot see — a host running ``end_strategy="early"`` can finish
on text beside a function call — are caught by ``after_run``, which can still
record the outcome even though it can no longer ask for a citation.
"""
if not _ends_the_question(response, request_context):
return response
evidence = discover_evidence(ctx)
question = question_in_progress(evidence)
if not _has_evidence_to_declare(evidence):
return response
records = [found.record for found in evidence]
if citation_status(records, question=question) != "missing":
return response
if _already_asked(ctx.messages, question):
return response
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.
Recorded once per question: a resumption of the same question decides
again, and one question is one outcome.
"""
outer = getattr(ctx.deps, "state", None)
if not isinstance(outer, dict):
return
state = CitationPolicyState.model_validate(outer.get(STATE_NAMESPACE) or {})
if question not in state.violations:
state.violations.append(question)
outer[STATE_NAMESPACE] = state.model_dump(mode="json")
async def after_run(
self, ctx: RunContext[Any], *, result: AgentRunResult[Any]
) -> AgentRunResult[Any]:
"""Record a question that finished undeclared, whatever ended it.
The backstop for an ending ``after_model_request`` cannot recognise. Nothing
can be asked of the model now, so this only records: a question that reached
the end of its run without a declaration is a violation, and one that was
asked but never answered is the same.
"""
evidence = discover_evidence(ctx)
question = question_in_progress(evidence)
if not _has_evidence_to_declare(evidence):
return result
records = [found.record for found in evidence]
if citation_status(records, question=question) == "missing":
self._record_violation(ctx, question)
return result
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 _ends_the_question(
response: ModelResponse, request_context: ModelRequestContext
) -> bool:
"""Whether this response finishes the question rather than continuing it.
An output tool call is a ``ToolCallPart`` like any other, but it carries the
final answer and ends the run, so treating every tool call as intermediate let a
structured answer finish undeclared.
"""
calls = [part for part in response.parts if isinstance(part, ToolCallPart)]
if not calls:
return True
output_tools = {
tool.name for tool in request_context.model_request_parameters.output_tools
}
return any(call.tool_name in output_tools for call in calls)
def _already_asked(messages: list[ModelMessage], question: int) -> bool:
"""Whether this question has already been asked to declare its grounding.
Read from the history rather than remembered on the instance, which a
resumption's ``for_run`` would forget — the same question would then be asked
twice. It also makes the right call when a redirect was enqueued but the run
ended before it reached the model: nothing is in the history, so it is asked
again, which is what the model needs.
Matched on the machine tag rather than the wording, so a question that happens
to contain the phrase cannot pass as a redirect we sent.
"""
return any(
isinstance(part, UserPromptPart)
and isinstance(part.content, str)
and CITATION_REDIRECT_TAG in part.content
for message in messages[question:]
for part in message.parts
)
def _has_evidence_to_declare(evidence: list[DiscoveredEvidence]) -> bool:
"""Whether anything exists that this answer could have been grounded on.
Either this question produced an evidence outcome, or the conversation has
already cited something — which stays available to a later answer, in a capsule
if a compactor is registered and in full if not. Requiring a fresh outcome
exempted exactly the follow-up that reuses earlier evidence, which is the case
enforcement exists for.
A conversation that has neither has nothing to declare: a greeting, an aside.
Read from the ledger rather than from ``state.searches``, which a new question
clears, so an answer grounded on code execution or a document read counts too.
"""
question = question_in_progress(evidence)
return any(
found.record.latest_evidence_epoch > question or found.record.occurrences
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",
"CITATION_REDIRECT_TAG",
"REDIRECT",
"REDIRECT_HINT",
"STATE_NAMESPACE",
"CitationPolicyCapability",
"CitationPolicyState",
"create_capability",
]

46
tests/capabilities/test_capabilities.py
View file

@ -295,7 +295,7 @@ async def test_run_error_closes_resources_and_propagates(temp_db_path):
@pytest.mark.asyncio
async def test_search_and_empty_citation_limits(temp_db_path):
async def test_a_spent_search_budget_fails_the_tool(temp_db_path):
config = AppConfig()
config.qa.max_searches = 0
capability = create_rag(db_path=temp_db_path, config=config)
@ -304,9 +304,6 @@ async def test_search_and_empty_citation_limits(temp_db_path):
with pytest.raises(ToolFailed, match="Search limit reached"):
await capability._search("anything", None)
with pytest.raises(ModelRetry, match="chunk_ids was empty"):
await capability._cite([])
@pytest.mark.asyncio
async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db_path):
@ -1242,3 +1239,44 @@ async def test_a_capability_fetches_its_own_evidences_pictures(temp_db_path):
client.document_item_repository.get_picture_bytes.assert_awaited_once_with(
"doc-1", "#/pictures/0"
)
@pytest.mark.asyncio
async def test_citing_nothing_is_a_valid_declaration(temp_db_path):
"""A model with nothing to cite must be able to say so.
Refusing the call left silence as the only way to express it, which is
indistinguishable from forgetting to cite at all.
"""
capability = create_rag(db_path=temp_db_path, config=AppConfig())
capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
capability.epoch = 5
result = await capability._cite([])
record = capability.state.evidence
assert "no" in result.lower()
assert record.declaration is not None
assert record.declaration.refs == []
assert citation_status([record], question=0) == "ungrounded"
assert capability.state.citations == []
@pytest.mark.asyncio
async def test_citing_nothing_after_citing_something_keeps_it_grounded(temp_db_path):
"""Declaring again cannot narrow what a question already declared."""
capability = create_rag(db_path=temp_db_path, config=AppConfig())
capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
capability.epoch = 5
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
capability.rag = client
await capability._cite(["chunk-1"])
await capability._cite([])
record = capability.state.evidence
assert citation_status([record], question=0) == "grounded"

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

@ -0,0 +1,497 @@
from dataclasses import dataclass, field
from typing import Any, cast
from unittest.mock import patch
import pytest
from pydantic import BaseModel
from pydantic_ai import Agent, DeferredToolResults
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 (
CITATION_REDIRECT_TAG,
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"
@pytest.mark.asyncio
async def test_a_follow_up_answered_from_retained_evidence_is_enforced(temp_db_path):
"""The multi-turn case is the one enforcement exists for.
A follow-up about something already cited needs no new search — the evidence is
still on the wire, whether in a capsule or in full — so requiring a fresh
evidence outcome let exactly those answers through undeclared.
"""
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
turns = iter(
[
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
[TextPart("first answer")],
[TextPart("a follow-up answered from what is already here")],
[TextPart("a follow-up answered from what is already here")],
]
)
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 what colour is the box in it?",
deps=deps,
message_history=first.all_messages(),
)
assert [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]
@pytest.mark.asyncio
async def test_a_conversation_that_never_cited_anything_is_still_left_alone(
temp_db_path,
):
"""A greeting has nothing to declare, and no evidence exists to declare from."""
_, _, 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_resumed_question_is_not_redirected_twice(temp_db_path):
"""Once per question has to mean once, across every run of that question.
Tracking it on the run instance forgot it at the next `for_run`, so resuming an
interrupted question asked for the citation again.
"""
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
turns = iter(
[
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
[TextPart("uncited")],
[TextPart("uncited again")],
[TextPart("uncited a third time")],
[TextPart("uncited a fourth time")],
[TextPart("uncited a fifth time")],
]
)
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)
# The same question again, continued rather than asked anew.
await agent.run(
deps=deps,
message_history=[
*first.all_messages(),
ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-9")]),
],
deferred_tool_results=DeferredToolResults(
calls={"call-9": "external result"}
),
)
redirects = [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]
assert len(redirects) == 1
@pytest.mark.asyncio
async def test_a_user_quoting_the_redirect_does_not_suppress_enforcement(temp_db_path):
"""Prose is not proof that we asked: a user can write any phrase.
Matching the wording let a question that merely mentioned it pass as already
asked, which silently switches enforcement off.
"""
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
turns = iter(
[
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
[TextPart("uncited")],
[TextPart("uncited again")],
]
)
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()]
)
with patch.object(RAGCapability, "_search", stub_search):
await agent.run(
"Please record what grounded the answer you already gave, in your notes.",
deps=Deps(),
)
assert [p for p in prompts_of(sent[-1]) if CITATION_REDIRECT_TAG in p]
class Answer(BaseModel):
"""A structured output, which the model returns through an output tool."""
text: str
@pytest.mark.asyncio
async def test_a_structured_output_answer_does_not_escape_enforcement(temp_db_path):
"""An output tool call is a `ToolCallPart` too, and it ends the run.
Treating every tool call as intermediate let a model search, skip citing, emit
its structured answer and finish with neither a redirect nor a violation.
"""
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
turns = iter(
[
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
[
ToolCallPart(
"final_result", {"text": "uncited structured answer"}, "out"
)
],
[
ToolCallPart(
"final_result", {"text": "uncited structured answer"}, "out"
)
],
]
)
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,
output_type=Answer,
capabilities=[rag, create_policy()],
)
deps = Deps()
with patch.object(RAGCapability, "_search", stub_search):
await agent.run("what does the supervisor do?", deps=deps)
# The cite tool is available here, so the redirect is what must happen; the
# backstop recording a violation would pass an `or` even with detection broken.
assert [p for p in prompts_of(sent[-1]) if CITATION_REDIRECT_TAG in p]
@pytest.mark.asyncio
async def test_a_question_asked_once_and_still_undeclared_is_recorded(temp_db_path):
"""Being asked is not an outcome; the question still ended undeclared.
Returning early on the redirect marker meant a question that was asked, ignored,
and then finished — with the cite tool possibly gone by that point — was neither
redirected again nor recorded anywhere.
"""
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
turns = iter(
[
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
[TextPart("uncited")],
[TextPart("still uncited after being asked")],
]
)
async def model(_messages, _info):
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):
await agent.run("what does the supervisor do?", deps=deps)
assert deps.state["citation_policy"]["violations"] == [0]

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)

35
tests/capabilities/test_evidence_ledger.py
View file

@ -240,3 +240,38 @@ def test_a_question_starts_clear_of_the_one_before_it():
assert record.latest_evidence_epoch == 0
assert record.declaration is None
assert citation_status([record], question=9) == "missing"
def test_an_empty_citation_after_a_grounded_one_cannot_narrow_it():
"""Citing again must not weaken a declaration, at any epoch.
Merging only within one epoch meant a second thought a request later replaced
the refs with nothing and reported the question ungrounded.
"""
record = CapabilityEvidenceRecord(question=0)
record.declare([rag_ref()], epoch=3)
record.declare([], epoch=5)
assert record.declaration is not None
assert [ref.chunk_id for ref in record.declaration.refs] == ["c1"]
assert citation_status([record], question=0) == "grounded"
def test_a_declaration_after_newer_evidence_starts_afresh():
"""Evidence the model has since seen may be what it is now citing."""
record = CapabilityEvidenceRecord(question=0)
record.declare([rag_ref("first")], epoch=3)
record.note_evidence(4)
record.declare([rag_ref("second")], epoch=5)
assert record.declaration is not None
assert [ref.chunk_id for ref in record.declaration.refs] == ["second"]
def test_an_empty_citation_after_newer_evidence_is_ungrounded():
record = CapabilityEvidenceRecord(question=0)
record.declare([rag_ref()], epoch=3)
record.note_evidence(4)
record.declare([], epoch=5)
assert citation_status([record], question=0) == "ungrounded"