Merge pull request #538 from ggozad/feat/citation-policy
Citation policy: require a declaration, accept an empty one
This commit is contained in:
commit
0a0ddbd3b6
13 changed files with 961 additions and 67 deletions
|
|
@ -3,6 +3,8 @@
|
||||||
|
|
||||||
### Added
|
### 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.
|
- `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.
|
- `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()`.
|
- `RAGCapabilityBase.evidence_tool_names()` and `get_picture_bytes()`.
|
||||||
|
|
@ -10,6 +12,7 @@
|
||||||
|
|
||||||
### Changed
|
### 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.
|
- `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.
|
- 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.
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -7,6 +7,7 @@ haiku.rag provides native [Pydantic AI capabilities](https://ai.pydantic.dev/cap
|
||||||
| [`RAGCapability`](rag.md) | Grounded document search and citations. |
|
| [`RAGCapability`](rag.md) | Grounded document search and citations. |
|
||||||
| [`AnalysisCapability`](analysis.md) | Corpus computation and structural analysis with sandboxed Python. |
|
| [`AnalysisCapability`](analysis.md) | Corpus computation and structural analysis with sandboxed Python. |
|
||||||
| `EvidenceCompactionCapability` | Optional. Shrinking a conversation's history to the evidence that was cited. |
|
| `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.
|
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
|
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.
|
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
|
## 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.
|
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.
|
||||||
|
|
|
||||||
|
|
@ -273,6 +273,16 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
||||||
if tool.capability_id != self.id or tool.name == self._cite_tool_name
|
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]:
|
def evidence_tool_names(self) -> set[str]:
|
||||||
"""Tools that can bring new evidence into the run.
|
"""Tools that can bring new evidence into the run.
|
||||||
|
|
||||||
|
|
@ -419,12 +429,17 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
||||||
return formatted
|
return formatted
|
||||||
|
|
||||||
async def _cite(self, chunk_ids: list[str]) -> str:
|
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
|
assert self.state is not None
|
||||||
if not chunk_ids:
|
if not chunk_ids:
|
||||||
raise ModelRetry(
|
self._declare([])
|
||||||
"No citations registered: chunk_ids was empty. Pass the chunk_ids "
|
return "Recorded: this answer cites no knowledge-base evidence."
|
||||||
"you want to cite, copied verbatim from search results."
|
|
||||||
)
|
|
||||||
|
|
||||||
all_results: list[SearchResult] = []
|
all_results: list[SearchResult] = []
|
||||||
state = cast(Any, self.state)
|
state = cast(Any, self.state)
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,6 @@
|
||||||
from collections.abc import Mapping, Sequence
|
from collections.abc import Sequence
|
||||||
from dataclasses import dataclass, field, replace
|
from dataclasses import dataclass, field, replace
|
||||||
from typing import Any, cast
|
from typing import Any
|
||||||
|
|
||||||
from pydantic_ai import RunContext
|
from pydantic_ai import RunContext
|
||||||
from pydantic_ai.capabilities import AbstractCapability, WrapModelRequestHandler
|
from pydantic_ai.capabilities import AbstractCapability, WrapModelRequestHandler
|
||||||
|
|
@ -15,7 +15,11 @@ from pydantic_ai.messages import (
|
||||||
from pydantic_ai.models import ModelRequestContext
|
from pydantic_ai.models import ModelRequestContext
|
||||||
|
|
||||||
from haiku.rag.capabilities._base import RAGCapabilityBase
|
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.store.models.citation import Citation
|
||||||
from haiku.rag.tools.search import RETRIEVED_IMAGE_TAG, decode_picture
|
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)
|
@dataclass(frozen=True)
|
||||||
class RetainedPicture:
|
class RetainedPicture:
|
||||||
"""A picture to re-attach, with the label that must accompany it.
|
"""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
|
of what was retrieved and break the message counts that question identities
|
||||||
and epochs are derived from.
|
and epochs are derived from.
|
||||||
"""
|
"""
|
||||||
evidence = self.discover(ctx)
|
evidence = discover_evidence(ctx)
|
||||||
boundary = max((found.record.question or 0 for found in evidence), default=0)
|
boundary = question_in_progress(evidence)
|
||||||
if boundary > 0:
|
if boundary > 0:
|
||||||
await self._build_once(ctx, evidence)
|
await self._build_once(ctx, evidence)
|
||||||
request_context.messages = compact_history(
|
request_context.messages = compact_history(
|
||||||
|
|
@ -395,28 +385,6 @@ class EvidenceCompactionCapability(AbstractCapability[Any]):
|
||||||
content.append(picture)
|
content.append(picture)
|
||||||
return tuple(content)
|
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:
|
def create_capability() -> EvidenceCompactionCapability:
|
||||||
"""Create the capability that compacts history from recorded evidence."""
|
"""Create the capability that compacts history from recorded evidence."""
|
||||||
|
|
@ -434,7 +402,6 @@ __all__ = [
|
||||||
"CAPSULE_HEADER",
|
"CAPSULE_HEADER",
|
||||||
"RECEIPT",
|
"RECEIPT",
|
||||||
"Capsule",
|
"Capsule",
|
||||||
"DiscoveredEvidence",
|
|
||||||
"EvidenceCompactionCapability",
|
"EvidenceCompactionCapability",
|
||||||
"RetainedPicture",
|
"RetainedPicture",
|
||||||
"build_capsule",
|
"build_capsule",
|
||||||
|
|
|
||||||
63
haiku_rag_slim/haiku/rag/capabilities/evidence.py
Normal file
63
haiku_rag_slim/haiku/rag/capabilities/evidence.py
Normal 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",
|
||||||
|
]
|
||||||
|
|
@ -3,9 +3,9 @@
|
||||||
You answer questions over a document knowledge base. Two common workflows:
|
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_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
|
## Tools
|
||||||
|
|
||||||
|
|
@ -25,7 +25,7 @@ Search the knowledge base directly (outside code execution). Each result has a `
|
||||||
### analysis_cite
|
### 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.
|
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:
|
Chunk IDs come from two places:
|
||||||
- The `chunk_id` field on `search` / `await search(...)` results
|
- 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.
|
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.
|
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
|
## 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`)
|
- 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.
|
- 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.
|
- 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.
|
||||||
|
|
|
||||||
|
|
@ -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.
|
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
|
### 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.
|
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
|
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
|
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
|
## 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
|
- If multiple results are relevant, synthesize them coherently
|
||||||
- Be concise and direct — avoid elaboration unless asked
|
- 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 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.
|
- 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
|
## When search returns irrelevant results
|
||||||
|
|
|
||||||
|
|
@ -133,10 +133,11 @@ class CapabilityEvidenceRecord(BaseModel):
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Record validated citations for the current question.
|
"""Record validated citations for the current question.
|
||||||
|
|
||||||
Repeated calls at the same epoch merge, so citing again cannot narrow what
|
Citing again cannot narrow what a question already declared: calls merge
|
||||||
was already declared: an empty call after a grounded one leaves it
|
while no evidence outcome has followed the standing declaration, whatever
|
||||||
grounded. A call at a later epoch declares afresh, because evidence the
|
epoch they arrive at, so an empty second thought leaves a grounded question
|
||||||
model saw in between may be what it is now citing.
|
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:
|
if self.question is None:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
|
|
@ -145,14 +146,16 @@ class CapabilityEvidenceRecord(BaseModel):
|
||||||
)
|
)
|
||||||
self._reject_regression(epoch, "A declaration")
|
self._reject_regression(epoch, "A declaration")
|
||||||
current = self.declaration
|
current = self.declaration
|
||||||
if current is not None and (current.question, current.epoch) == (
|
if (
|
||||||
self.question,
|
current is not None
|
||||||
epoch,
|
and current.question == self.question
|
||||||
|
and self.latest_evidence_epoch <= current.epoch
|
||||||
):
|
):
|
||||||
known = {(ref.capability, ref.chunk_id) for ref in current.refs}
|
known = {(ref.capability, ref.chunk_id) for ref in current.refs}
|
||||||
current.refs.extend(
|
current.refs.extend(
|
||||||
ref for ref in refs if (ref.capability, ref.chunk_id) not in known
|
ref for ref in refs if (ref.capability, ref.chunk_id) not in known
|
||||||
)
|
)
|
||||||
|
current.epoch = max(current.epoch, epoch)
|
||||||
else:
|
else:
|
||||||
self.declaration = CitationDeclaration(
|
self.declaration = CitationDeclaration(
|
||||||
question=self.question, epoch=epoch, refs=list(refs)
|
question=self.question, epoch=epoch, refs=list(refs)
|
||||||
|
|
|
||||||
233
haiku_rag_slim/haiku/rag/capabilities/policy.py
Normal file
233
haiku_rag_slim/haiku/rag/capabilities/policy.py
Normal 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",
|
||||||
|
]
|
||||||
|
|
@ -295,7 +295,7 @@ async def test_run_error_closes_resources_and_propagates(temp_db_path):
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@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 = AppConfig()
|
||||||
config.qa.max_searches = 0
|
config.qa.max_searches = 0
|
||||||
capability = create_rag(db_path=temp_db_path, config=config)
|
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"):
|
with pytest.raises(ToolFailed, match="Search limit reached"):
|
||||||
await capability._search("anything", None)
|
await capability._search("anything", None)
|
||||||
|
|
||||||
with pytest.raises(ModelRetry, match="chunk_ids was empty"):
|
|
||||||
await capability._cite([])
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db_path):
|
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(
|
client.document_item_repository.get_picture_bytes.assert_awaited_once_with(
|
||||||
"doc-1", "#/pictures/0"
|
"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
497
tests/capabilities/test_citation_policy.py
Normal 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]
|
||||||
|
|
@ -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.analysis import create_capability as create_analysis
|
||||||
from haiku.rag.capabilities.compaction import (
|
from haiku.rag.capabilities.compaction import (
|
||||||
CAPSULE_HEADER,
|
CAPSULE_HEADER,
|
||||||
DiscoveredEvidence,
|
|
||||||
EvidenceCompactionCapability,
|
EvidenceCompactionCapability,
|
||||||
build_capsule,
|
build_capsule,
|
||||||
group_label,
|
group_label,
|
||||||
)
|
)
|
||||||
from haiku.rag.capabilities.compaction import create_capability as create_compaction
|
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 (
|
from haiku.rag.capabilities.ledger import (
|
||||||
CapabilityEvidenceRecord,
|
CapabilityEvidenceRecord,
|
||||||
EvidenceOccurrence,
|
EvidenceOccurrence,
|
||||||
|
|
@ -77,6 +77,7 @@ def discovered(
|
||||||
for chunk_id in cited
|
for chunk_id in cited
|
||||||
},
|
},
|
||||||
tool_names=frozenset({f"{capability}_search"}),
|
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):
|
async def spy(self, ctx):
|
||||||
await original(self, ctx)
|
await original(self, ctx)
|
||||||
found.append(self.discover(ctx))
|
found.append(discover_evidence(ctx))
|
||||||
|
|
||||||
return patch.object(EvidenceCompactionCapability, "before_run", spy)
|
return patch.object(EvidenceCompactionCapability, "before_run", spy)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -240,3 +240,38 @@ def test_a_question_starts_clear_of_the_one_before_it():
|
||||||
assert record.latest_evidence_epoch == 0
|
assert record.latest_evidence_epoch == 0
|
||||||
assert record.declaration is None
|
assert record.declaration is None
|
||||||
assert citation_status([record], question=9) == "missing"
|
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"
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue