Merge pull request #535 from ggozad/feat/evidence-compaction
Record cited evidence, and make history compaction an explicit capability
This commit is contained in:
Yiorgis Gozadinos
GitHub
commit
89b3a9186c
21 changed files with 2858 additions and 136 deletions
18
CHANGELOG.md
18
CHANGELOG.md
|
|
@ -1,12 +1,30 @@
|
|||
# Changelog
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- `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()`.
|
||||
- `haiku.rag.tools.search.decode_picture()`.
|
||||
|
||||
### Changed
|
||||
|
||||
- `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.
|
||||
|
||||
### Removed
|
||||
|
||||
- `PRIOR_TURN_NOTICE` and `_compact_old_tool_returns` from `haiku.rag.capabilities._base`, and `RAGCapabilityBase.turn_start`.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `cross-encoder` reranking no longer ties the scores of strongly-relevant candidates, which left their order to the sort. Scores remain 0-1.
|
||||
- `haiku-rag` and `haiku-ingester` CLI startup no longer imports `lancedb`, `pyarrow` and `pydantic_ai`.
|
||||
- `haiku.rag.store` no longer re-exports `Store`; import it from `haiku.rag.store.engine`.
|
||||
- `docling-local` reuses one docling `DocumentConverter` per set of conversion options instead of building one per document, so local layout, table and OCR models are no longer loaded per document. Conversions through a shared converter are serialized.
|
||||
- A resumed run keeps the searches, citations and executions of the question in progress instead of clearing them.
|
||||
- Each page image attached to a search result is preceded by a line giving its position and the chunk id it came from. `build_binary_parts_from_results` is now `build_image_content_from_results` and returns those labels interleaved with the pictures.
|
||||
|
||||
## [0.73.0] - 2026-08-06
|
||||
|
||||
|
|
|
|||
|
|
@ -18,6 +18,9 @@ from starlette.requests import Request
|
|||
from starlette.responses import JSONResponse, Response, StreamingResponse
|
||||
from starlette.routing import Route
|
||||
|
||||
from haiku.rag.capabilities.compaction import (
|
||||
create_capability as create_compaction,
|
||||
)
|
||||
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState, create_capability
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import load_yaml_config
|
||||
|
|
@ -81,7 +84,9 @@ capability = create_capability(db_path=db_path, config=Config, defer_loading=Fal
|
|||
agent = Agent(
|
||||
get_model(Config.qa.model, Config),
|
||||
instructions=AGENT_PREAMBLE,
|
||||
capabilities=[capability],
|
||||
# Conversations here are multi-turn, so earlier questions are reduced to the
|
||||
# evidence they cited rather than carried whole.
|
||||
capabilities=[capability, create_compaction()],
|
||||
deps_type=AppDeps,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -133,21 +133,21 @@ Two approaches are benchmarked separately:
|
|||
| Embedding Model | Reranker | Cases | MAP |
|
||||
|------------------------------------------|------------------------------------------------------|------:|-------:|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | none | 3045 | 0.9774 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | none | 3045 | 0.9709 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | none | 3045 | 0.9798 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `nvidia/llama-nemotron-rerank-vl-1b-v2` (multimodal) | 3045 | 0.9913 |
|
||||
|
||||
*The reranked row uses `reranking.multimodal: true`: picture chunks reach the vision reranker as images alongside their description text. Measured on haiku.rag main post-v0.67.3 (multimodal reranking ships in the next release).*
|
||||
*The nemotron row without a reranker is measured on this release. The reranked row uses `reranking.multimodal: true`: picture chunks reach the vision reranker as images alongside their description text, measured on haiku.rag main post-v0.67.3.*
|
||||
|
||||
##### QA accuracy + citation retrieval
|
||||
|
||||
| Embedding Model | Target | Capability model | Cases | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------------------|-----------------|-----------------------------------|------:|-------------|------------------|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 1409 | 0.89 | — |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.93 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.94 | 0.78 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3039 | 0.9263 | 0.9761 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3040 | 0.9362 | 0.9343 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Qwen3.6-35B-A3B-NVFP4` | 3045 | 0.95 | 0.93 |
|
||||
|
||||
*Measured on haiku.rag v0.52.0, no reranker, judged by `vllm:Qwen3.6-35B-A3B-NVFP4`. Qwen3-VL covered 1409 / 3045 cases.*
|
||||
*Both nemotron `Gemma-4` rows are measured on this release, no reranker, judged by `vllm:Qwen3.6-35B-A3B-NVFP4` with thinking on, and exclude the cases that errored (6 of 3045 for `rag-capability`, 5 for `analysis-capability`). The `rag-capability` row cites at 99.64% with a mean of 1.08 citations per case, at a median 4.7s per case against 5.0s for `analysis-capability`. Citation coverage is what moved on this release: 4.9% of analysis cases register no citation, against 26.3% before, at unchanged searches and code executions per case. The remaining rows are from haiku.rag v0.52.0, where Qwen3-VL covered 1409 / 3045 cases.*
|
||||
|
||||
#### Text embedder + VLM picture descriptions
|
||||
|
||||
|
|
|
|||
|
|
@ -46,6 +46,8 @@ async with HaikuRAG("my.lancedb") as client:
|
|||
|
||||
## State
|
||||
|
||||
When dependencies expose a state dictionary, `AnalysisState` is stored under `"analysis"`. It contains the document filter, code execution log, searches, and citations. Per-run searches and executions reset automatically; the filter and citation index persist.
|
||||
When dependencies expose a state dictionary, `AnalysisState` is stored under `"analysis"`. It contains the document filter, code execution log, searches, citations, and the `evidence` record of what was retrieved and cited per question. Searches and executions are cleared when a new question starts, and a resumed question keeps them; the filter, citation index and evidence record persist.
|
||||
|
||||
This capability does not alter the message history either. Register the [compaction capability](index.md#multi-turn-conversations) to compact earlier questions.
|
||||
|
||||
The capability lazily opens both LanceDB and the sandbox only after it is loaded and a tool requires them. Resources close at the end of the agent run.
|
||||
|
|
|
|||
|
|
@ -1,13 +1,14 @@
|
|||
# Capabilities
|
||||
|
||||
haiku.rag provides two native [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/):
|
||||
haiku.rag provides native [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/):
|
||||
|
||||
| Capability | Use it for |
|
||||
|---|---|
|
||||
| [`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. |
|
||||
|
||||
Both 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.
|
||||
|
||||
## Compose an agent
|
||||
|
||||
|
|
@ -34,6 +35,38 @@ agent = Agent(
|
|||
)
|
||||
```
|
||||
|
||||
## Multi-turn conversations
|
||||
|
||||
Every question adds its search results to the history, so requests grow turn after
|
||||
turn, and can degrade answers or exceed a provider's limits as they do. Register the
|
||||
compaction capability to replace earlier questions' evidence with the evidence that
|
||||
was actually cited:
|
||||
|
||||
```python
|
||||
from haiku.rag.capabilities.compaction import create_capability as compaction
|
||||
from haiku.rag.capabilities.rag import create_capability as rag
|
||||
|
||||
agent = Agent(
|
||||
"openai:gpt-5",
|
||||
capabilities=[rag(db_path="my.lancedb"), compaction()],
|
||||
)
|
||||
```
|
||||
|
||||
Cited text and cited page images are kept in full, grouped by the question that
|
||||
cited them, and stay citable by the same chunk ids. Everything else earlier becomes a
|
||||
short receipt. Registering the capability is the only switch: leave it out and the
|
||||
transcript reaches the model untouched. There is nothing to configure.
|
||||
|
||||
Compaction rewrites the request, never the stored history, so `all_messages()` still
|
||||
holds everything the run gathered. Retained evidence still grows with the
|
||||
conversation — this reduces what a request carries, it does not bound it. A host that
|
||||
needs more aggressive pruning can compact its own requests further, on the wire only.
|
||||
|
||||
Resuming a question (deferred tool results, an interruption, a suspension) requires
|
||||
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.
|
||||
|
||||
## 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.
|
||||
|
|
|
|||
|
|
@ -37,16 +37,19 @@ class RAGState(BaseModel):
|
|||
citation_index: dict[str, Citation]
|
||||
citations: list[str]
|
||||
document_filter: str | None
|
||||
evidence: CapabilityEvidenceRecord
|
||||
searches: dict[str, list[SearchResult]]
|
||||
```
|
||||
|
||||
`document_filter` persists between runs. Current citations and searches reset for each run, while the citation index remains available to the host application.
|
||||
`document_filter`, `citation_index` and `evidence` persist across runs. Citations and searches are cleared when a new question starts; a run that resumes a question keeps the evidence it is still answering from.
|
||||
|
||||
`evidence` records which chunks this capability retrieved and cited, and in which question. `haiku.rag.capabilities.ledger.citation_status(records, question=...)` derives `missing`, `grounded` or `ungrounded` from it, across capabilities.
|
||||
|
||||
State is ordinary application state; the capability does not depend on AG-UI. An AG-UI application can expose it using Pydantic AI's standard adapter.
|
||||
|
||||
## Context management
|
||||
|
||||
Large RAG tool results from earlier user turns are replaced with a short marker before model requests. Tool-call pairing and current-turn evidence are retained. This prevents long conversations from repeatedly sending old retrieved content.
|
||||
This capability does not alter the message history. To stop long conversations resending old retrieved content, register the [compaction capability](index.md#multi-turn-conversations) alongside it.
|
||||
|
||||
## Domain context and vision
|
||||
|
||||
|
|
|
|||
|
|
@ -15,8 +15,6 @@ from pydantic_ai.messages import (
|
|||
ModelResponse,
|
||||
ToolCallPart,
|
||||
ToolReturn,
|
||||
ToolReturnPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
from pydantic_ai.models import ModelRequestContext
|
||||
from pydantic_ai.run import AgentRunResult
|
||||
|
|
@ -24,11 +22,12 @@ from pydantic_ai.tools import ToolDefinition
|
|||
from pydantic_ai.toolsets import AgentToolset
|
||||
|
||||
from haiku.rag.capabilities._tools import CodeExecutionEntry, search_corpus
|
||||
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord, EvidenceRef
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
from haiku.rag.store.models.citation import Citation, resolve_citations
|
||||
from haiku.rag.tools.search import build_binary_parts_from_results
|
||||
from haiku.rag.tools.search import build_image_content_from_results
|
||||
|
||||
CITATION_GRACE_REQUESTS = 2
|
||||
"""Requests calling this capability's tools that its cite tool outlives the rest by.
|
||||
|
|
@ -70,47 +69,42 @@ def resolve_db_path(db_path: Path | None, config: AppConfig) -> Path:
|
|||
|
||||
|
||||
def _clear_invocation_state(state: BaseModel) -> None:
|
||||
"""Drop the working evidence of the previous question.
|
||||
|
||||
Only ever called when a new question starts. A resumption keeps it: the
|
||||
results belong to the question still being answered, and dropping them leaves
|
||||
a later citation unable to resolve against the expanded result the model saw,
|
||||
recording no provenance for it.
|
||||
"""
|
||||
for field_name in ("citations", "searches", "executions"):
|
||||
value = getattr(state, field_name, None)
|
||||
if hasattr(value, "clear"):
|
||||
value.clear()
|
||||
|
||||
|
||||
def _compact_old_tool_returns(
|
||||
messages: list[ModelMessage], tool_names: frozenset[str]
|
||||
) -> list[ModelMessage]:
|
||||
"""Remove bulky prior-turn evidence while retaining the current turn.
|
||||
def _is_resumption(prompt: Any, messages: list[ModelMessage]) -> bool:
|
||||
"""Whether this run continues a question rather than asking a new one.
|
||||
|
||||
Tool call and return parts remain paired; only the old return payload is
|
||||
replaced. This keeps provider histories valid and preserves all evidence
|
||||
gathered since the most recent user prompt.
|
||||
Two signals, either of which is enough, because getting this wrong hands the
|
||||
model a notice where its own evidence should be:
|
||||
|
||||
- no prompt: how pydantic-ai resumes for interruptions and suspensions.
|
||||
- an unfinished tail: the history ends with a request the model has not
|
||||
answered, or with a response whose tool calls have no returns yet. Deferred
|
||||
tool results may arrive *with* a prompt, so the prompt alone is not enough.
|
||||
|
||||
A settled history ends with the previous answer, so a genuinely new question
|
||||
is not mistaken for a continuation. The framework's own first-new-message
|
||||
index would be better than either signal, but it is not public here.
|
||||
"""
|
||||
latest_user_message = -1
|
||||
for index, message in enumerate(messages):
|
||||
if isinstance(message, ModelRequest) and any(
|
||||
isinstance(part, UserPromptPart) for part in message.parts
|
||||
):
|
||||
latest_user_message = index
|
||||
|
||||
if latest_user_message < 0:
|
||||
return messages
|
||||
|
||||
compacted = list(messages)
|
||||
for index, message in enumerate(messages[:latest_user_message]):
|
||||
if not isinstance(message, ModelRequest):
|
||||
continue
|
||||
parts = [
|
||||
replace(
|
||||
part,
|
||||
content="[Prior-turn RAG tool output removed; citations remain in state.]",
|
||||
)
|
||||
if isinstance(part, ToolReturnPart) and part.tool_name in tool_names
|
||||
else part
|
||||
for part in message.parts
|
||||
]
|
||||
if parts != message.parts:
|
||||
compacted[index] = replace(message, parts=parts)
|
||||
return compacted
|
||||
if prompt is None:
|
||||
return True
|
||||
if not messages:
|
||||
return False
|
||||
last = messages[-1]
|
||||
if isinstance(last, ModelRequest):
|
||||
return True
|
||||
return any(isinstance(part, ToolCallPart) for part in last.parts)
|
||||
|
||||
|
||||
def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) -> bool:
|
||||
|
|
@ -142,13 +136,38 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
search_count: int = field(default=0, repr=False)
|
||||
request_count: int = field(default=0, repr=False)
|
||||
grace_requests_used: int = field(default=0, repr=False)
|
||||
epoch: int = field(default=0, repr=False)
|
||||
|
||||
async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]":
|
||||
"""Start a run's own copy, and settle which question it is answering.
|
||||
|
||||
A new question takes the message count as its identity, which every
|
||||
participant derives identically from the same history. A resumption keeps
|
||||
the identity already recorded: the question is the one in progress, and
|
||||
adopting the current count would relabel it as a new one and judge its
|
||||
declarations against the wrong question. A resumption with no recorded
|
||||
identity is a state this design does not produce, so it is reported rather
|
||||
than guessed at. With no history at all there is nothing in progress: an
|
||||
absent prompt is then an instructions-only first question, which takes an
|
||||
identity like any other.
|
||||
"""
|
||||
outer = getattr(ctx.deps, "state", None)
|
||||
outer_state = outer if isinstance(outer, dict) else None
|
||||
raw_state = outer_state.get(self.state_namespace) if outer_state else None
|
||||
resuming = _is_resumption(ctx.prompt, ctx.messages)
|
||||
continuing = resuming and bool(ctx.messages)
|
||||
state = self.state_type.model_validate(raw_state or {})
|
||||
_clear_invocation_state(state)
|
||||
record = cast(CapabilityEvidenceRecord, cast(Any, state).evidence)
|
||||
if continuing and record.question is None:
|
||||
raise RuntimeError(
|
||||
f"The {self.state_namespace} capability is resuming a question with "
|
||||
"no stored question identity. Capabilities cannot be added, removed "
|
||||
"or migrated while a question is unfinished, and the run's state "
|
||||
"must be carried between its runs."
|
||||
)
|
||||
if not continuing:
|
||||
_clear_invocation_state(state)
|
||||
record.begin_question(len(ctx.messages))
|
||||
run_capability = replace(
|
||||
self,
|
||||
state=state,
|
||||
|
|
@ -159,6 +178,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
search_count=0,
|
||||
request_count=0,
|
||||
grace_requests_used=0,
|
||||
epoch=0,
|
||||
)
|
||||
run_capability._sync_state()
|
||||
return run_capability
|
||||
|
|
@ -171,9 +191,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
async def before_model_request(
|
||||
self, ctx: RunContext[Any], request_context: ModelRequestContext
|
||||
) -> ModelRequestContext:
|
||||
request_context.messages = _compact_old_tool_returns(
|
||||
request_context.messages, self.tool_names
|
||||
)
|
||||
self.epoch = len(ctx.messages)
|
||||
if instruction := self._budget_notice():
|
||||
current_request = request_context.messages[-1]
|
||||
if isinstance(current_request, ModelRequest):
|
||||
|
|
@ -217,7 +235,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
)
|
||||
if spent := self._spent_tool_names():
|
||||
names = ", ".join(sorted(spent))
|
||||
if remaining := sorted(self._evidence_tool_names() - spent):
|
||||
if remaining := sorted(self.evidence_tool_names() - spent):
|
||||
return (
|
||||
f"The {self.state_namespace} capability has spent its budget "
|
||||
f"for {names}; further calls to them fail. Gather any further "
|
||||
|
|
@ -255,8 +273,13 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
if tool.capability_id != self.id or tool.name == self._cite_tool_name
|
||||
]
|
||||
|
||||
def _evidence_tool_names(self) -> set[str]:
|
||||
"""Tools that can bring new evidence into the run."""
|
||||
def evidence_tool_names(self) -> set[str]:
|
||||
"""Tools that can bring new evidence into the run.
|
||||
|
||||
Public because compaction needs to know whose output on the wire is
|
||||
evidence: a cite acknowledgement is a receipt of the model's own action and
|
||||
must survive, while a code execution that reached the corpus is evidence.
|
||||
"""
|
||||
return {f"{self.state_namespace}_search"}
|
||||
|
||||
def _spent_tool_names(self) -> set[str]:
|
||||
|
|
@ -306,6 +329,18 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
self.rag = rag
|
||||
return self.rag
|
||||
|
||||
async def get_picture_bytes(self, document_id: str, self_ref: str) -> bytes | None:
|
||||
"""Fetch a picture of this capability's evidence, for whoever re-attaches it.
|
||||
|
||||
Public because compaction rehydrates cited pictures and this capability
|
||||
already holds the connection they came from; bytes are never kept in state.
|
||||
"""
|
||||
async with self.rag_lock:
|
||||
rag = await self._ensure_rag()
|
||||
return await rag.document_item_repository.get_picture_bytes(
|
||||
document_id, self_ref
|
||||
)
|
||||
|
||||
async def _close(self) -> None:
|
||||
if self.rag is not None:
|
||||
await self.rag.__aexit__(None, None, None)
|
||||
|
|
@ -326,6 +361,41 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
finally:
|
||||
self._sync_state()
|
||||
|
||||
def _evidence_record(self) -> CapabilityEvidenceRecord:
|
||||
assert self.state is not None
|
||||
return cast(CapabilityEvidenceRecord, cast(Any, self.state).evidence)
|
||||
|
||||
def _note_evidence(self) -> None:
|
||||
"""Record an outcome the model can ground an answer on.
|
||||
|
||||
Includes an empty search result and a failed execution that still printed
|
||||
output: negative evidence grounds a refusal. Excludes a spent budget, which
|
||||
yields nothing to ground anything on.
|
||||
"""
|
||||
self._evidence_record().note_evidence(self.epoch)
|
||||
|
||||
def _declare(self, citations: list[Citation]) -> None:
|
||||
"""Record what the model cited, once the ids have resolved.
|
||||
|
||||
Declaring earlier would let a call naming only unresolvable ids read as a
|
||||
grounded answer.
|
||||
"""
|
||||
state = cast(Any, self.state)
|
||||
retrieved = {
|
||||
result.chunk_id
|
||||
for results in state.searches.values()
|
||||
for result in results
|
||||
if result.chunk_id
|
||||
}
|
||||
self._evidence_record().declare(
|
||||
[
|
||||
EvidenceRef(capability=self.state_namespace, chunk_id=c.chunk_id)
|
||||
for c in citations
|
||||
],
|
||||
epoch=self.epoch,
|
||||
retrieved_now=retrieved,
|
||||
)
|
||||
|
||||
async def _search(self, query: str, limit: int | None) -> str | ToolReturn:
|
||||
assert self.state is not None
|
||||
self.search_count += 1
|
||||
|
|
@ -343,7 +413,8 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
)
|
||||
state = cast(Any, self.state)
|
||||
state.searches[query] = results
|
||||
if self.vision and (parts := build_binary_parts_from_results(results)):
|
||||
self._note_evidence()
|
||||
if self.vision and (parts := build_image_content_from_results(results)):
|
||||
return ToolReturn(return_value=formatted, content=parts)
|
||||
return formatted
|
||||
|
||||
|
|
@ -390,6 +461,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
"Copy chunk_ids verbatim from search results."
|
||||
)
|
||||
self._register_citations(citations)
|
||||
self._declare(citations)
|
||||
resolved = {citation.chunk_id for citation in citations}
|
||||
unresolved = [cid for cid in missing if cid not in resolved]
|
||||
if unresolved:
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from haiku.rag.capabilities._base import (
|
|||
RAGCapabilityBase,
|
||||
resolve_db_path,
|
||||
)
|
||||
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
|
|
@ -29,6 +30,7 @@ class AnalysisState(BaseModel):
|
|||
executions: list[CodeExecutionEntry] = Field(default_factory=list)
|
||||
citation_index: dict[str, Citation] = Field(default_factory=dict)
|
||||
citations: list[str] = Field(default_factory=list)
|
||||
evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord)
|
||||
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
|
||||
|
||||
|
||||
|
|
@ -85,11 +87,11 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
|
|||
self.sandbox = None
|
||||
await super()._close()
|
||||
|
||||
def _evidence_tool_names(self) -> set[str]:
|
||||
def evidence_tool_names(self) -> set[str]:
|
||||
# Searching from inside the sandbox does not count against
|
||||
# `qa.max_searches`, so code execution outlives a spent search budget
|
||||
# as a way to reach new evidence.
|
||||
return super()._evidence_tool_names() | {"analysis_execute_code"}
|
||||
return super().evidence_tool_names() | {"analysis_execute_code"}
|
||||
|
||||
def _spent_tool_names(self) -> set[str]:
|
||||
spent = super()._spent_tool_names()
|
||||
|
|
@ -107,6 +109,8 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
|
|||
)
|
||||
sandbox = await self._ensure_sandbox()
|
||||
result = await sandbox.execute(code)
|
||||
if result.success or result.stdout:
|
||||
self._note_evidence()
|
||||
if sandbox._search_results:
|
||||
existing = self.state.searches.get("_sandbox", [])
|
||||
seen = {item.chunk_id for item in existing}
|
||||
|
|
|
|||
445
haiku_rag_slim/haiku/rag/capabilities/compaction.py
Normal file
445
haiku_rag_slim/haiku/rag/capabilities/compaction.py
Normal file
|
|
@ -0,0 +1,445 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import Any, cast
|
||||
|
||||
from pydantic_ai import RunContext
|
||||
from pydantic_ai.capabilities import AbstractCapability, WrapModelRequestHandler
|
||||
from pydantic_ai.messages import (
|
||||
BinaryContent,
|
||||
ModelMessage,
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
ToolReturnPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
from pydantic_ai.models import ModelRequestContext
|
||||
|
||||
from haiku.rag.capabilities._base import RAGCapabilityBase
|
||||
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.tools.search import RETRIEVED_IMAGE_TAG, decode_picture
|
||||
|
||||
CAPABILITY_ID = "haiku-rag-evidence-compaction"
|
||||
|
||||
CAPSULE_HEADER = (
|
||||
"[Evidence cited earlier in this conversation, kept so later questions can "
|
||||
"rely on it. Cite these chunk_ids directly when you use them.]"
|
||||
)
|
||||
|
||||
RECEIPT = (
|
||||
"[Evidence retrieved for an earlier question, no longer shown. It does not "
|
||||
"count as cited for the current question.]"
|
||||
)
|
||||
|
||||
ENTRY_SEPARATOR = "\n\n"
|
||||
|
||||
|
||||
def group_label(position: int) -> str:
|
||||
"""Name a group by its position among the groups, not by question number.
|
||||
|
||||
A question identity is a message count, so a header built from it would present
|
||||
an index as a turn number, and an ordinal over the groups is not the
|
||||
conversation's ordinal either whenever a question in between cited nothing. The
|
||||
label claims only what it is: a grouping, newest first.
|
||||
"""
|
||||
return f"[Cited evidence group {position}]"
|
||||
|
||||
|
||||
def picture_label(chunk_id: str, self_ref: str) -> str:
|
||||
return (
|
||||
f"Page image retrieved from the knowledge base for cited evidence "
|
||||
f"[{chunk_id}] ({self_ref}). Not provided by the user. {RETRIEVED_IMAGE_TAG}"
|
||||
)
|
||||
|
||||
|
||||
@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.
|
||||
|
||||
Addressed by owner, document and reference, because a reference such as
|
||||
``#/pictures/0`` repeats across documents and capabilities. The label travels
|
||||
with it so it can never be emitted without its image.
|
||||
"""
|
||||
|
||||
capability: str
|
||||
chunk_id: str
|
||||
document_id: str
|
||||
self_ref: str
|
||||
label: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Capsule:
|
||||
"""Everything the compactor would insert, and nothing about where it goes."""
|
||||
|
||||
text: str = ""
|
||||
pictures: tuple[RetainedPicture, ...] = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Entry:
|
||||
capability: str
|
||||
chunk_id: str
|
||||
question: int
|
||||
citation: Citation
|
||||
|
||||
def render(self) -> str:
|
||||
title = self.citation.document_title
|
||||
uri = self.citation.document_uri
|
||||
source = f'"{title}"' if title else uri
|
||||
if title and uri and uri != title:
|
||||
source = f"{source} ({uri})"
|
||||
return f"[{self.chunk_id}] Source: {source}\n{self.citation.content}"
|
||||
|
||||
|
||||
def _eligible_entries(evidence: Sequence[DiscoveredEvidence]) -> list[_Entry]:
|
||||
"""Cited evidence with content, newest citing question first.
|
||||
|
||||
Evidence cited in several questions belongs to the most recent one, so it is
|
||||
rendered once and grouped where the model last used it.
|
||||
|
||||
An occurrence and its canonical ``Citation`` are written by the same call, so a
|
||||
cited chunk without one is not a state this design produces. Rendering the rest
|
||||
regardless would quietly drop evidence an answer rested on, so it is reported.
|
||||
"""
|
||||
entries = []
|
||||
for discovered in evidence:
|
||||
for chunk_id, occurrence in discovered.record.occurrences.items():
|
||||
if not occurrence.cited_in_questions:
|
||||
continue
|
||||
citation = discovered.citations.get(chunk_id)
|
||||
if citation is None:
|
||||
raise ValueError(
|
||||
f"{discovered.capability} cited {chunk_id} in question(s) "
|
||||
f"{occurrence.cited_in_questions} but has no citation record "
|
||||
"for it, so its content cannot be retained."
|
||||
)
|
||||
entries.append(
|
||||
_Entry(
|
||||
capability=discovered.capability,
|
||||
chunk_id=chunk_id,
|
||||
question=max(occurrence.cited_in_questions),
|
||||
citation=citation,
|
||||
)
|
||||
)
|
||||
entries.sort(key=lambda entry: (-entry.question, entry.capability, entry.chunk_id))
|
||||
return entries
|
||||
|
||||
|
||||
def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule:
|
||||
"""Render every cited piece of evidence, grouped by the question that cited it.
|
||||
|
||||
Everything cited is kept whole and everything else is dropped. There is no
|
||||
character budget: what a model can hold is the model's business, and a knob for
|
||||
it would only half-rescue models that fail on long conversations regardless.
|
||||
|
||||
A host that needs earlier evidence pruned can compact further on top, on the wire
|
||||
only. Removing or reordering the stored history breaks the message counts that
|
||||
question identities and epochs are derived from, and the next record written is
|
||||
refused.
|
||||
|
||||
Pure: no I/O and no message history, so what goes on the wire stays separable
|
||||
from what it should contain. Picture bytes are fetched by the caller, which is
|
||||
why a picture travels with its label rather than beside it.
|
||||
"""
|
||||
entries = _eligible_entries(evidence)
|
||||
if not entries:
|
||||
return Capsule()
|
||||
|
||||
lines = [CAPSULE_HEADER]
|
||||
pictures: list[RetainedPicture] = []
|
||||
seen: set[tuple[str, str, str]] = set()
|
||||
position = 0
|
||||
current_question: int | None = None
|
||||
for entry in entries:
|
||||
if entry.question != current_question:
|
||||
position += 1
|
||||
current_question = entry.question
|
||||
lines.append(group_label(position))
|
||||
lines.append(entry.render())
|
||||
for self_ref in entry.citation.picture_refs:
|
||||
# Overlapping chunks cite one figure, and a provider counts it twice.
|
||||
# Identity is owner plus document plus reference, so the same reference
|
||||
# in another document stays a different picture.
|
||||
identity = (entry.capability, entry.citation.document_id, self_ref)
|
||||
if identity in seen:
|
||||
continue
|
||||
seen.add(identity)
|
||||
pictures.append(
|
||||
RetainedPicture(
|
||||
capability=entry.capability,
|
||||
chunk_id=entry.chunk_id,
|
||||
document_id=entry.citation.document_id,
|
||||
self_ref=self_ref,
|
||||
label=picture_label(entry.chunk_id, self_ref),
|
||||
)
|
||||
)
|
||||
return Capsule(text=ENTRY_SEPARATOR.join(lines), pictures=tuple(pictures))
|
||||
|
||||
|
||||
def _strip_our_pictures(part: UserPromptPart) -> UserPromptPart | None:
|
||||
"""Drop the pictures we attached, together with the labels describing them.
|
||||
|
||||
Ours is a label carrying the machine tag immediately followed by an image —
|
||||
both halves required. Position alone is not ownership, since several tools'
|
||||
results can arrive in one request; prose alone is not either, because a user can
|
||||
write any phrase, and treating one as proof removed a user's own picture along
|
||||
with their text. A label is only ever dropped with its picture: left behind it
|
||||
would tell the model a figure is present when it is gone.
|
||||
"""
|
||||
if isinstance(part.content, str):
|
||||
return part
|
||||
items = list(part.content)
|
||||
kept: list[Any] = []
|
||||
index = 0
|
||||
while index < len(items):
|
||||
item = items[index]
|
||||
following = items[index + 1] if index + 1 < len(items) else None
|
||||
is_ours = (
|
||||
isinstance(item, str)
|
||||
and RETRIEVED_IMAGE_TAG in item
|
||||
and isinstance(following, BinaryContent)
|
||||
)
|
||||
if is_ours:
|
||||
index += 2
|
||||
continue
|
||||
kept.append(item)
|
||||
index += 1
|
||||
return replace(part, content=kept) if kept else None
|
||||
|
||||
|
||||
def compact_history(
|
||||
messages: list[ModelMessage],
|
||||
*,
|
||||
boundary: int,
|
||||
owned_tools: frozenset[str],
|
||||
capsule_text: str,
|
||||
capsule_images: Sequence[str | BinaryContent] = (),
|
||||
) -> list[ModelMessage]:
|
||||
"""Replace earlier questions' evidence with the capsule, on a copy.
|
||||
|
||||
``boundary`` is how many messages existed when the current question arrived, so
|
||||
everything below it belongs to an earlier one. It comes from the recorded
|
||||
question identity rather than from message shape: mid-question a user-role part
|
||||
is as likely to be page images or an injected notice, and reading either as the
|
||||
next question strips evidence the model is still answering from.
|
||||
|
||||
The newest earlier return carries the capsule and every other becomes a receipt,
|
||||
so exactly one capsule exists by construction. Returns are never removed, only
|
||||
rewritten, which keeps each one paired with its call. Nothing outside this
|
||||
capability's evidence tools is touched — not a cite acknowledgement, not another
|
||||
capability's output, not a picture the user attached.
|
||||
"""
|
||||
if boundary <= 0:
|
||||
return messages
|
||||
|
||||
carrier = _newest_owned_return(messages, boundary, owned_tools)
|
||||
compacted = list(messages)
|
||||
for index, message in enumerate(messages[:boundary]):
|
||||
if not isinstance(message, ModelRequest):
|
||||
continue
|
||||
parts: list[Any] = []
|
||||
for position, part in enumerate(message.parts):
|
||||
if isinstance(part, ToolReturnPart) and part.tool_name in owned_tools:
|
||||
carries = (index, position) == carrier
|
||||
body = capsule_text or RECEIPT if carries else RECEIPT
|
||||
parts.append(replace(part, content=body))
|
||||
elif isinstance(part, UserPromptPart):
|
||||
if (kept := _strip_our_pictures(part)) is not None:
|
||||
parts.append(kept)
|
||||
else:
|
||||
parts.append(part)
|
||||
if carrier is not None and index == carrier[0] and capsule_images:
|
||||
parts.append(UserPromptPart(content=list(capsule_images)))
|
||||
if not parts:
|
||||
# A request with no parts is not a message; whatever emptied it was not
|
||||
# ours to remove after all.
|
||||
continue
|
||||
if parts != message.parts:
|
||||
compacted[index] = replace(message, parts=parts)
|
||||
return compacted
|
||||
|
||||
|
||||
def _newest_owned_return(
|
||||
messages: list[ModelMessage], boundary: int, owned_tools: frozenset[str]
|
||||
) -> tuple[int, int] | None:
|
||||
"""Where the last of our evidence returns is, as message and part.
|
||||
|
||||
The part matters: a model can call search twice in one response, so one request
|
||||
can hold several of our returns, and giving the capsule to each duplicates the
|
||||
whole of it.
|
||||
"""
|
||||
for index in range(min(boundary, len(messages)) - 1, -1, -1):
|
||||
message = messages[index]
|
||||
if not isinstance(message, ModelRequest):
|
||||
continue
|
||||
for position in range(len(message.parts) - 1, -1, -1):
|
||||
part = message.parts[position]
|
||||
if isinstance(part, ToolReturnPart) and part.tool_name in owned_tools:
|
||||
return index, position
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class EvidenceCompactionCapability(AbstractCapability[Any]):
|
||||
"""Rewrites the history from what the evidence capabilities recorded.
|
||||
|
||||
Registering it is what turns compaction on: a host that leaves it out gets an
|
||||
untouched transcript, which is why it has no enable flag. It reads the evidence
|
||||
capabilities through the run's registry and holds no reference to any of them,
|
||||
so a host running one capability, both, or neither needs no wiring change.
|
||||
|
||||
Registering two is rejected by pydantic-ai before the run starts, since they
|
||||
would share this capability's id.
|
||||
"""
|
||||
|
||||
built_for: tuple[str | None, int] | None = field(default=None, repr=False)
|
||||
capsule: Capsule = field(default_factory=Capsule, repr=False)
|
||||
images: tuple[str | BinaryContent, ...] = field(default=(), repr=False)
|
||||
|
||||
async def for_run(self, ctx: RunContext[Any]) -> "EvidenceCompactionCapability":
|
||||
"""Give the run its own build cache, so concurrent runs cannot share one."""
|
||||
return replace(self, built_for=None, capsule=Capsule(), images=())
|
||||
|
||||
async def wrap_model_request(
|
||||
self,
|
||||
ctx: RunContext[Any],
|
||||
*,
|
||||
request_context: ModelRequestContext,
|
||||
handler: WrapModelRequestHandler,
|
||||
) -> ModelResponse:
|
||||
"""Rewrite the request, never the stored history.
|
||||
|
||||
Deliberately not ``before_model_request``: that hook's result is assigned
|
||||
back onto the run's message history, which would destroy the host's record
|
||||
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)
|
||||
if boundary > 0:
|
||||
await self._build_once(ctx, evidence)
|
||||
request_context.messages = compact_history(
|
||||
request_context.messages,
|
||||
boundary=boundary,
|
||||
owned_tools=frozenset().union(
|
||||
*(found.tool_names for found in evidence)
|
||||
),
|
||||
capsule_text=self.capsule.text,
|
||||
capsule_images=self.images,
|
||||
)
|
||||
return await handler(request_context)
|
||||
|
||||
async def _build_once(
|
||||
self, ctx: RunContext[Any], evidence: Sequence[DiscoveredEvidence]
|
||||
) -> None:
|
||||
"""Build the capsule once per model request, however often the hook runs.
|
||||
|
||||
Keyed on the run and its step rather than persisted: a stored key would
|
||||
freeze one question's capsule across the next.
|
||||
"""
|
||||
key = (ctx.run_id, ctx.run_step)
|
||||
if key == self.built_for:
|
||||
return
|
||||
self.capsule = build_capsule(evidence)
|
||||
self.images = await self._rehydrate(ctx)
|
||||
self.built_for = key
|
||||
|
||||
async def _rehydrate(self, ctx: RunContext[Any]) -> tuple[str | BinaryContent, ...]:
|
||||
"""Fetch the cited pictures through the capability that retrieved them.
|
||||
|
||||
Bytes are never stored in state, and the owner already holds an open
|
||||
connection. A picture that cannot be fetched, for any reason, or that will
|
||||
not decode, is emitted with neither its image nor its label: a label can
|
||||
never outlive what it describes, and a figure the model has already been
|
||||
given in text is not worth failing a question over.
|
||||
"""
|
||||
owners = {
|
||||
capability.state_namespace: capability
|
||||
for capability in ctx.capabilities.values()
|
||||
if isinstance(capability, RAGCapabilityBase)
|
||||
}
|
||||
content: list[str | BinaryContent] = []
|
||||
for retained in self.capsule.pictures:
|
||||
# Indexed, not looked up defensively: the capsule was built from these
|
||||
# same capabilities in this same call, so a missing owner is a broken
|
||||
# invariant rather than a picture to skip.
|
||||
owner = owners[retained.capability]
|
||||
try:
|
||||
data = await owner.get_picture_bytes(
|
||||
retained.document_id, retained.self_ref
|
||||
)
|
||||
except Exception:
|
||||
# A read that fails costs this picture, not the answer.
|
||||
continue
|
||||
if data is None:
|
||||
continue
|
||||
picture = decode_picture(data, retained.self_ref)
|
||||
if picture is None:
|
||||
continue
|
||||
content.append(retained.label)
|
||||
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."""
|
||||
return EvidenceCompactionCapability(
|
||||
id=CAPABILITY_ID,
|
||||
description=(
|
||||
"Replaces earlier questions' evidence on the wire with a capsule of "
|
||||
"what was cited."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CAPABILITY_ID",
|
||||
"CAPSULE_HEADER",
|
||||
"RECEIPT",
|
||||
"Capsule",
|
||||
"DiscoveredEvidence",
|
||||
"EvidenceCompactionCapability",
|
||||
"RetainedPicture",
|
||||
"build_capsule",
|
||||
"compact_history",
|
||||
"create_capability",
|
||||
"group_label",
|
||||
"picture_label",
|
||||
]
|
||||
213
haiku_rag_slim/haiku/rag/capabilities/ledger.py
Normal file
213
haiku_rag_slim/haiku/rag/capabilities/ledger.py
Normal file
|
|
@ -0,0 +1,213 @@
|
|||
from collections.abc import Iterable
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
CitationStatus = Literal["missing", "grounded", "ungrounded"]
|
||||
|
||||
|
||||
class EvidenceRef(BaseModel):
|
||||
"""One piece of evidence, identified by its owner as well as its chunk.
|
||||
|
||||
A chunk id alone is not an identity: the same id can be reported by more than
|
||||
one capability, and ownership is what tells compaction whose output it may
|
||||
touch.
|
||||
"""
|
||||
|
||||
capability: str
|
||||
chunk_id: str
|
||||
|
||||
|
||||
class EvidenceOccurrence(BaseModel):
|
||||
"""Which questions retrieved a piece of evidence, and which cited it."""
|
||||
|
||||
capability: str
|
||||
chunk_id: str
|
||||
retrieved_in_questions: list[int] = Field(default_factory=list)
|
||||
cited_in_questions: list[int] = Field(default_factory=list)
|
||||
|
||||
|
||||
class CitationDeclaration(BaseModel):
|
||||
"""What a question declared as its grounding, and when.
|
||||
|
||||
Bound to a question *and* an epoch: the epoch outlives a question, so matching
|
||||
it alone would let a question that gathered no evidence inherit the previous
|
||||
declaration and read as compliant having declared nothing.
|
||||
"""
|
||||
|
||||
question: int
|
||||
epoch: int
|
||||
refs: list[EvidenceRef] = Field(default_factory=list)
|
||||
|
||||
|
||||
class CapabilityEvidenceRecord(BaseModel):
|
||||
"""What one evidence capability wrote, in its own state namespace.
|
||||
|
||||
Holds no content: ``Citation`` in ``citation_index`` is the canonical record
|
||||
and already persists content, document id and picture refs. This is the index
|
||||
over it that compaction needs and the transcript cannot provide.
|
||||
|
||||
Single-writer by construction. A record shared between capabilities would be
|
||||
overwritten by whichever of them synced its state last; merging happens in the
|
||||
transient views built by ``citation_status`` and the optional capabilities.
|
||||
|
||||
``question`` is the number of messages that existed when the question arrived,
|
||||
and ``epoch`` the number when an outcome occurred. Both are derived from the
|
||||
conversation rather than counted locally, so every participant computes the
|
||||
same values without sharing a counter. ``question`` is unset until a run
|
||||
establishes it, so a record a host merely created is distinguishable from one
|
||||
that has been through a question.
|
||||
"""
|
||||
|
||||
occurrences: dict[str, EvidenceOccurrence] = Field(default_factory=dict)
|
||||
question: int | None = None
|
||||
latest_evidence_epoch: int = 0
|
||||
declaration: CitationDeclaration | None = None
|
||||
|
||||
def _reject_regression(self, count: int, what: str) -> None:
|
||||
"""Refuse a message count below one already recorded in this question.
|
||||
|
||||
Identities and epochs are both message counts, and every comparison
|
||||
between them assumes the conversation only grows while a question is
|
||||
being answered. One capability truncating or reordering the history
|
||||
breaks that, and each way of recording it has to refuse the same way: an
|
||||
unchecked evidence outcome freezes every later declaration as stale, while
|
||||
an unchecked declaration replaces a newer one with an older one and
|
||||
revives the answer it grounded.
|
||||
"""
|
||||
recorded = max(
|
||||
self.question or 0,
|
||||
self.latest_evidence_epoch,
|
||||
self.declaration.epoch if self.declaration else 0,
|
||||
)
|
||||
if count < recorded:
|
||||
raise ValueError(
|
||||
f"{what} at message count {count} is behind {recorded}, which is "
|
||||
"already recorded: message history must be append-only for "
|
||||
"question identities and epochs to hold."
|
||||
)
|
||||
|
||||
def begin_question(self, identity: int) -> None:
|
||||
"""Take the identity of a question that has just arrived.
|
||||
|
||||
The evidence epoch and the declaration describe the question that has
|
||||
just ended, and outlive their meaning the moment the next one begins:
|
||||
kept, they let the last question's citations ground this one, and hold a
|
||||
horizon this question's own declarations cannot pass. Clearing them
|
||||
confines those comparisons to one question, so a host whose stored
|
||||
history shifted between two of them is answered rather than refused.
|
||||
|
||||
Identities themselves must still separate one question from the next.
|
||||
Occurrences outlive the question that recorded them and carry the
|
||||
identities that cited them, which compaction groups and orders the capsule
|
||||
by: a reused identity merges two questions into one group, and a lower one
|
||||
renders later evidence as though it were cited earlier.
|
||||
"""
|
||||
if self.question is not None and identity <= self.question:
|
||||
raise ValueError(
|
||||
f"A question at message count {identity} is not past question "
|
||||
f"{self.question}, which is already answered: identities separate "
|
||||
"one question from the next and are compared as recency."
|
||||
)
|
||||
self.question = identity
|
||||
self.latest_evidence_epoch = 0
|
||||
self.declaration = None
|
||||
|
||||
def note_evidence(self, epoch: int) -> None:
|
||||
"""Record that the model has seen an evidence outcome.
|
||||
|
||||
Called for anything an answer could rest on, including a search that
|
||||
returned nothing and a failed execution that still printed output — a
|
||||
fruitless search grounds a refusal. Not called for a failure that yields
|
||||
no evidence at all, such as an exhausted budget.
|
||||
"""
|
||||
self._reject_regression(epoch, "Evidence")
|
||||
self.latest_evidence_epoch = epoch
|
||||
|
||||
def declare(
|
||||
self,
|
||||
refs: list[EvidenceRef],
|
||||
*,
|
||||
epoch: int,
|
||||
retrieved_now: set[str] | None = None,
|
||||
) -> None:
|
||||
"""Record validated citations for the current question.
|
||||
|
||||
Repeated calls at the same epoch merge, so citing again cannot narrow what
|
||||
was already declared: an empty call after a grounded one leaves it
|
||||
grounded. A call at a later epoch declares afresh, because evidence the
|
||||
model saw in between may be what it is now citing.
|
||||
"""
|
||||
if self.question is None:
|
||||
raise ValueError(
|
||||
"Citations cannot be declared before a run establishes the "
|
||||
"question identity."
|
||||
)
|
||||
self._reject_regression(epoch, "A declaration")
|
||||
current = self.declaration
|
||||
if current is not None and (current.question, current.epoch) == (
|
||||
self.question,
|
||||
epoch,
|
||||
):
|
||||
known = {(ref.capability, ref.chunk_id) for ref in current.refs}
|
||||
current.refs.extend(
|
||||
ref for ref in refs if (ref.capability, ref.chunk_id) not in known
|
||||
)
|
||||
else:
|
||||
self.declaration = CitationDeclaration(
|
||||
question=self.question, epoch=epoch, refs=list(refs)
|
||||
)
|
||||
|
||||
for ref in refs:
|
||||
occurrence = self.occurrences.setdefault(
|
||||
ref.chunk_id,
|
||||
EvidenceOccurrence(capability=ref.capability, chunk_id=ref.chunk_id),
|
||||
)
|
||||
if self.question not in occurrence.cited_in_questions:
|
||||
occurrence.cited_in_questions.append(self.question)
|
||||
if (
|
||||
retrieved_now
|
||||
and ref.chunk_id in retrieved_now
|
||||
and self.question not in occurrence.retrieved_in_questions
|
||||
):
|
||||
occurrence.retrieved_in_questions.append(self.question)
|
||||
|
||||
|
||||
def citation_status(
|
||||
records: Iterable[CapabilityEvidenceRecord], *, question: int
|
||||
) -> CitationStatus:
|
||||
"""Derived, never stored, so refs and status cannot contradict.
|
||||
|
||||
A declaration is current only for the question it was made in, and only if it
|
||||
followed the newest evidence outcome of *every* capability: a question where
|
||||
one capability cited and another then searched without citing is not grounded.
|
||||
Strictly later, since a citation made in the same request as an outcome cannot
|
||||
have read it.
|
||||
|
||||
A grounding *violation* is not one of these: that is an enforcement outcome
|
||||
recorded by the policy capability, not something the model declared.
|
||||
"""
|
||||
records = list(records)
|
||||
horizon = max((record.latest_evidence_epoch for record in records), default=0)
|
||||
current = [
|
||||
record.declaration
|
||||
for record in records
|
||||
if record.declaration is not None
|
||||
and record.declaration.question == question
|
||||
and record.declaration.epoch > horizon
|
||||
]
|
||||
if not current:
|
||||
return "missing"
|
||||
return (
|
||||
"grounded" if any(declaration.refs for declaration in current) else "ungrounded"
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CapabilityEvidenceRecord",
|
||||
"CitationDeclaration",
|
||||
"CitationStatus",
|
||||
"EvidenceOccurrence",
|
||||
"EvidenceRef",
|
||||
"citation_status",
|
||||
]
|
||||
|
|
@ -12,6 +12,7 @@ from haiku.rag.capabilities._base import (
|
|||
RAGCapabilityBase,
|
||||
resolve_db_path,
|
||||
)
|
||||
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
|
@ -33,6 +34,7 @@ _instructions_path = Path(__file__).parent / "instructions" / "rag.md"
|
|||
class RAGState(BaseModel):
|
||||
citation_index: dict[str, Citation] = Field(default_factory=dict)
|
||||
citations: list[str] = Field(default_factory=list)
|
||||
evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord)
|
||||
document_filter: str | None = None
|
||||
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import uuid
|
||||
from collections.abc import Iterable, Sequence
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
|
@ -25,6 +26,7 @@ from textual.worker import Worker
|
|||
|
||||
from haiku.rag.capabilities._base import RAGCapabilityBase
|
||||
from haiku.rag.capabilities.analysis import AnalysisState
|
||||
from haiku.rag.capabilities.compaction import create_capability as create_compaction
|
||||
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState
|
||||
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
|
||||
from haiku.rag.chat.widgets.image_select import ImageAdded
|
||||
|
|
@ -153,7 +155,9 @@ class ChatApp(App):
|
|||
self._model,
|
||||
deps_type=ChatDeps,
|
||||
instructions=AGENT_PREAMBLE,
|
||||
capabilities=self._capabilities,
|
||||
# A chat is multi-turn by definition, so earlier questions are reduced
|
||||
# to the evidence they cited rather than carried whole.
|
||||
capabilities=[*self._capabilities, create_compaction()],
|
||||
)
|
||||
self._state = {}
|
||||
for capability in self._capabilities:
|
||||
|
|
@ -205,7 +209,12 @@ class ChatApp(App):
|
|||
await chat_history.show_thinking()
|
||||
|
||||
message = None
|
||||
deps = ChatDeps(state=self._state)
|
||||
# The run gets a copy: state and message history have to advance together.
|
||||
# A cancelled or failed run discards its messages, and state that advanced
|
||||
# anyway would leave the next question deriving its identity from a shorter
|
||||
# history than the evidence already recorded — refused as non-append-only,
|
||||
# with the conversation stuck until it is cleared.
|
||||
deps = ChatDeps(state=deepcopy(self._state))
|
||||
|
||||
try:
|
||||
async with self._agent.run_stream_events(
|
||||
|
|
|
|||
|
|
@ -211,7 +211,7 @@ class SearchResult(BaseModel):
|
|||
parts.append(f"Type: {primary_label}")
|
||||
|
||||
# Surface picture captions when present. Order matches the binary
|
||||
# attachments emitted by build_binary_parts_from_results, so the model
|
||||
# attachments emitted by build_image_content_from_results, so the model
|
||||
# can correlate caption ↔ attached image by position (BinaryContent
|
||||
# identifiers don't survive serialization to the OpenAI vision API).
|
||||
if self.picture_captions:
|
||||
|
|
|
|||
|
|
@ -10,11 +10,35 @@ from haiku.rag.config.models import AppConfig
|
|||
from haiku.rag.store.models import SearchResult
|
||||
from haiku.rag.tools.context import RAGDeps
|
||||
|
||||
RETRIEVED_IMAGE_TAG = "[haiku.rag/retrieved-image]"
|
||||
"""Tag every label we attach to a retrieved picture ends with.
|
||||
|
||||
def build_binary_parts_from_results(
|
||||
Identifies our own pictures on the wire without inferring ownership from position,
|
||||
which is wrong as soon as two tools' results arrive in one request. Deliberately not
|
||||
a phrase: a user writing "retrieved from the knowledge base for my report" above
|
||||
their own picture had it removed, along with their text.
|
||||
"""
|
||||
|
||||
|
||||
def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
|
||||
"""Wrap picture bytes for the wire, or return nothing if they will not decode.
|
||||
|
||||
The model adapter renders one vision placeholder per ``BinaryContent``, so
|
||||
emitting one for an image the server cannot decode leaves the processor with an
|
||||
off-by-one count.
|
||||
"""
|
||||
try:
|
||||
with Image.open(BytesIO(data)) as image:
|
||||
image.verify()
|
||||
except Exception:
|
||||
return None
|
||||
return BinaryContent(data=data, media_type="image/png", identifier=self_ref)
|
||||
|
||||
|
||||
def build_image_content_from_results(
|
||||
results: list[SearchResult],
|
||||
) -> list[BinaryContent]:
|
||||
"""Decode and validate picture bytes attached to search results.
|
||||
) -> list[str | BinaryContent]:
|
||||
"""Decode and validate picture bytes attached to search results, labelled.
|
||||
|
||||
Dedup keyed on ``(document_id, self_ref)`` so the same picture in
|
||||
different chunks is sent once. Pictures that fail
|
||||
|
|
@ -22,8 +46,18 @@ def build_binary_parts_from_results(
|
|||
vision placeholder per ``BinaryContent``, so emitting one for an
|
||||
image the server can't decode leaves the processor with an
|
||||
off-by-one count.
|
||||
|
||||
Every picture is preceded by a line naming the result it belongs to.
|
||||
``ToolReturn.content`` reaches the model as a user-role message, so
|
||||
retrieved pictures are otherwise indistinguishable from ones the user
|
||||
attached, and models narrate them as part of the question: unlabelled,
|
||||
gemma4-26b answered about a figure from an unrelated document, and with a
|
||||
single note ahead of the batch it still called them "images in the prompt".
|
||||
The label also names the chunk to cite for a figure, which
|
||||
``BinaryContent.identifier`` cannot do — it does not survive serialization
|
||||
to the vision API.
|
||||
"""
|
||||
parts: list[BinaryContent] = []
|
||||
collected: list[tuple[str | None, str, BinaryContent]] = []
|
||||
seen: set[tuple[str | None, str]] = set()
|
||||
for result in results:
|
||||
if not result.image_data:
|
||||
|
|
@ -32,21 +66,22 @@ def build_binary_parts_from_results(
|
|||
key = (result.document_id, self_ref)
|
||||
if key in seen:
|
||||
continue
|
||||
data = base64.b64decode(b64)
|
||||
try:
|
||||
with Image.open(BytesIO(data)) as img:
|
||||
img.verify()
|
||||
except Exception:
|
||||
picture = decode_picture(base64.b64decode(b64), self_ref)
|
||||
if picture is None:
|
||||
continue
|
||||
parts.append(
|
||||
BinaryContent(
|
||||
data=data,
|
||||
media_type="image/png",
|
||||
identifier=self_ref,
|
||||
)
|
||||
)
|
||||
collected.append((result.chunk_id, self_ref, picture))
|
||||
seen.add(key)
|
||||
return parts
|
||||
|
||||
content: list[str | BinaryContent] = []
|
||||
total = len(collected)
|
||||
for position, (chunk_id, self_ref, picture) in enumerate(collected, 1):
|
||||
content.append(
|
||||
f"Page image {position} of {total}, retrieved from the knowledge base "
|
||||
f"for search result [{chunk_id}] ({self_ref}). "
|
||||
f"Not provided by the user. {RETRIEVED_IMAGE_TAG}"
|
||||
)
|
||||
content.append(picture)
|
||||
return content
|
||||
|
||||
|
||||
def create_search_toolset(
|
||||
|
|
@ -134,9 +169,9 @@ def create_search_toolset(
|
|||
if not config.qa.model.vision:
|
||||
return text
|
||||
|
||||
binary_parts = build_binary_parts_from_results(results_list)
|
||||
if binary_parts:
|
||||
return ToolReturn(return_value=text, content=binary_parts)
|
||||
image_content = build_image_content_from_results(results_list)
|
||||
if image_content:
|
||||
return ToolReturn(return_value=text, content=image_content)
|
||||
return text
|
||||
|
||||
toolset: FunctionToolset[RAGDeps] = FunctionToolset()
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from typing import Any, cast
|
|||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import Agent, ModelRetry, RunContext, ToolFailed
|
||||
from pydantic_ai import Agent, DeferredToolResults, ModelRetry, RunContext, ToolFailed
|
||||
from pydantic_ai.messages import (
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
|
|
@ -20,10 +20,13 @@ from pydantic_ai.usage import RunUsage
|
|||
from haiku.rag.capabilities._base import (
|
||||
CITATION_GRACE_REQUESTS,
|
||||
_called_own_tool,
|
||||
_compact_old_tool_returns,
|
||||
)
|
||||
from haiku.rag.capabilities.analysis import AnalysisCapability, AnalysisState
|
||||
from haiku.rag.capabilities.analysis import create_capability as create_analysis
|
||||
from haiku.rag.capabilities.ledger import (
|
||||
CapabilityEvidenceRecord,
|
||||
citation_status,
|
||||
)
|
||||
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGCapability, RAGState
|
||||
from haiku.rag.capabilities.rag import create_capability as create_rag
|
||||
from haiku.rag.config.models import AppConfig, PromptsConfig
|
||||
|
|
@ -308,7 +311,7 @@ async def test_search_and_empty_citation_limits(temp_db_path):
|
|||
@pytest.mark.asyncio
|
||||
async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db_path):
|
||||
capability = create_rag(db_path=temp_db_path, config=AppConfig())
|
||||
capability.state = RAGState()
|
||||
capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
|
||||
client = AsyncMock()
|
||||
client.get_chunk_by_id.side_effect = [
|
||||
Chunk(id="chunk-1", document_id="doc-1", content="first"),
|
||||
|
|
@ -334,7 +337,7 @@ async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db
|
|||
@pytest.mark.asyncio
|
||||
async def test_cite_reports_unresolved_ids_on_partial_success(temp_db_path):
|
||||
capability = create_rag(db_path=temp_db_path, config=AppConfig())
|
||||
capability.state = RAGState()
|
||||
capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
|
||||
client = AsyncMock()
|
||||
client.get_chunk_by_id.side_effect = [
|
||||
Chunk(id="chunk-1", document_id="doc-1", content="first"),
|
||||
|
|
@ -364,6 +367,7 @@ async def test_cite_repairs_chunk_ids_damaged_in_transcription(temp_db_path):
|
|||
unrelated = "9c2cd07e-5a3f-45a6-968d-cbd6f06ab57b"
|
||||
capability = create_rag(db_path=temp_db_path, config=AppConfig())
|
||||
capability.state = RAGState(
|
||||
evidence=CapabilityEvidenceRecord(question=0),
|
||||
searches={
|
||||
"q": [
|
||||
SearchResult(
|
||||
|
|
@ -374,7 +378,7 @@ async def test_cite_repairs_chunk_ids_damaged_in_transcription(temp_db_path):
|
|||
document_uri="test://document",
|
||||
)
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
client = AsyncMock()
|
||||
client.get_chunk_by_id.return_value = None
|
||||
|
|
@ -461,6 +465,58 @@ async def test_analysis_execution_limit_fails_the_tool(temp_db_path):
|
|||
await capability._execute_code("print('done')")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_spent_execution_budget_is_not_evidence(temp_db_path):
|
||||
"""Nothing was produced to ground an answer on, so nothing is recorded."""
|
||||
config = AppConfig()
|
||||
config.analysis.max_executions = 0
|
||||
capability = create_analysis(db_path=temp_db_path, config=config)
|
||||
capability.state = AnalysisState()
|
||||
capability.epoch = 5
|
||||
|
||||
with pytest.raises(ToolFailed):
|
||||
await capability._execute_code("print('done')")
|
||||
|
||||
assert capability.state.evidence.latest_evidence_epoch == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("success", "stdout", "expected_epoch"),
|
||||
[
|
||||
pytest.param(True, "42", 5, id="succeeded"),
|
||||
pytest.param(True, "", 5, id="succeeded without output"),
|
||||
pytest.param(False, "42", 5, id="failed after printing"),
|
||||
pytest.param(False, "", 0, id="failed without printing"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_a_code_execution_the_model_can_read_is_evidence(
|
||||
temp_db_path, success, stdout, expected_epoch
|
||||
):
|
||||
"""A raised error with nothing printed grounds nothing, so it is not evidence.
|
||||
|
||||
A failure that printed first does ground an answer, and so does a successful
|
||||
run whose outcome is that it printed nothing.
|
||||
"""
|
||||
capability = create_analysis(db_path=temp_db_path, config=AppConfig())
|
||||
capability.state = AnalysisState(evidence=CapabilityEvidenceRecord(question=0))
|
||||
capability.epoch = 5
|
||||
sandbox = AsyncMock(spec=Sandbox)
|
||||
sandbox._search_results = []
|
||||
sandbox.execute.return_value = SandboxResult(
|
||||
stdout=stdout, stderr="" if success else "boom", success=success
|
||||
)
|
||||
capability.sandbox = sandbox
|
||||
|
||||
if success:
|
||||
await capability._execute_code("print(42)")
|
||||
else:
|
||||
with pytest.raises(ToolFailed):
|
||||
await capability._execute_code("print(42)")
|
||||
|
||||
assert capability.state.evidence.latest_evidence_epoch == expected_epoch
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spent_search_budget_is_announced_but_keeps_the_tool(rag_db):
|
||||
"""A spent budget is announced; the tool stays declared to avoid a dead run.
|
||||
|
|
@ -555,7 +611,7 @@ async def test_spent_search_notice_points_at_code_while_it_has_budget(temp_db_pa
|
|||
notice = capability._budget_notice()
|
||||
assert notice is not None
|
||||
assert "analysis_execute_code" in notice
|
||||
assert capability._evidence_tool_names() <= capability._spent_tool_names()
|
||||
assert capability.evidence_tool_names() <= capability._spent_tool_names()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -570,7 +626,7 @@ async def test_spent_search_notice_tells_rag_to_answer(temp_db_path):
|
|||
|
||||
assert notice is not None
|
||||
assert "rag_search" in notice
|
||||
assert capability._evidence_tool_names() == {"rag_search"}
|
||||
assert capability.evidence_tool_names() == {"rag_search"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -746,7 +802,9 @@ async def test_native_agent_composition_initializes_host_state(temp_db_path):
|
|||
result = await agent.run("Hello", deps=deps)
|
||||
|
||||
assert result.output == "success (no tool calls)"
|
||||
assert deps.state["rag"] == RAGState().model_dump(mode="json")
|
||||
assert deps.state["rag"] == RAGState(
|
||||
evidence=CapabilityEvidenceRecord(question=0)
|
||||
).model_dump(mode="json")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -787,41 +845,400 @@ async def test_deferred_capability_loads_native_tools(temp_db_path):
|
|||
assert "rag_search" in loaded_payloads[0]
|
||||
|
||||
|
||||
def test_prior_turn_tool_results_are_compacted_but_current_evidence_is_kept():
|
||||
messages = [
|
||||
ModelRequest(parts=[UserPromptPart("old question")]),
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {}, "old-call")]),
|
||||
def _resuming_deps() -> Deps:
|
||||
"""State as a resumption always finds it: the question already identified.
|
||||
|
||||
A run that resumes has been through ``for_run`` before, so the identity of the
|
||||
question in progress is stored. Fabricating the history without it is a state
|
||||
the design does not produce, and is rejected rather than guessed at.
|
||||
"""
|
||||
return Deps(
|
||||
state={
|
||||
"rag": RAGState(evidence=CapabilityEvidenceRecord(question=0)).model_dump(
|
||||
mode="json"
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _in_flight_history() -> list[Any]:
|
||||
"""A question already asked and searched, still awaiting its answer."""
|
||||
return [
|
||||
ModelRequest(parts=[UserPromptPart("what does the supervisor do?")]),
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "s"}, "call-1")]),
|
||||
ModelRequest(
|
||||
parts=[ToolReturnPart("rag_search", "large old evidence", "old-call")]
|
||||
),
|
||||
ModelRequest(parts=[UserPromptPart("current question")]),
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {}, "current-call")]),
|
||||
ModelRequest(
|
||||
parts=[ToolReturnPart("rag_search", "current evidence", "current-call")]
|
||||
parts=[ToolReturnPart("rag_search", "EVIDENCE FOR THE LIVE TURN", "call-1")]
|
||||
),
|
||||
]
|
||||
|
||||
compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"}))
|
||||
|
||||
old_return = compacted[2].parts[0]
|
||||
current_return = compacted[5].parts[0]
|
||||
assert isinstance(old_return, ToolReturnPart)
|
||||
assert "removed" in str(old_return.content)
|
||||
assert isinstance(current_return, ToolReturnPart)
|
||||
assert current_return.content == "current evidence"
|
||||
def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord:
|
||||
return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"])
|
||||
|
||||
|
||||
def test_tool_results_are_unchanged_when_history_has_no_user_prompt():
|
||||
messages = [
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {}, "current-call")]),
|
||||
ModelRequest(
|
||||
parts=[ToolReturnPart("rag_search", "current evidence", "current-call")]
|
||||
),
|
||||
async def _stub_search(self, query: str, _limit: int | None) -> str:
|
||||
"""Record a result the way the real search does, so citing resolves."""
|
||||
cast(Any, self.state).searches[query] = [
|
||||
SearchResult(content="evidence", score=1.0, chunk_id="chunk-1")
|
||||
]
|
||||
self._note_evidence()
|
||||
return "EVIDENCE"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_question_takes_its_own_identity_and_both_capabilities_agree(
|
||||
temp_db_path,
|
||||
):
|
||||
"""Identity is derived from the conversation, so no counter is shared."""
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
analysis = create_analysis(
|
||||
db_path=temp_db_path, config=AppConfig(), defer_loading=False
|
||||
)
|
||||
|
||||
async def model(_messages, _info):
|
||||
return ModelResponse(parts=[TextPart("answer")])
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, analysis])
|
||||
deps = Deps()
|
||||
|
||||
first = await agent.run("first question", deps=deps)
|
||||
first_identity = _record(deps, "rag").question
|
||||
await agent.run("second question", deps=deps, message_history=first.all_messages())
|
||||
second_identity = _record(deps, "rag").question
|
||||
|
||||
assert first_identity == 0
|
||||
assert second_identity is not None and second_identity > 0
|
||||
assert _record(deps, "analysis").question == second_identity
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_resumption_keeps_the_identity_of_the_question_in_progress(
|
||||
temp_db_path,
|
||||
):
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
|
||||
async def model(_messages, _info):
|
||||
return ModelResponse(parts=[TextPart("answer")])
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
||||
deps = Deps(
|
||||
state={
|
||||
"rag": RAGState(evidence=CapabilityEvidenceRecord(question=7)).model_dump(
|
||||
mode="json"
|
||||
)
|
||||
}
|
||||
)
|
||||
history = [
|
||||
*_in_flight_history(),
|
||||
ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]),
|
||||
]
|
||||
|
||||
compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"}))
|
||||
await agent.run(
|
||||
"carry on",
|
||||
deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}),
|
||||
message_history=history,
|
||||
deps=deps,
|
||||
)
|
||||
|
||||
assert compacted is messages
|
||||
current_return = compacted[1].parts[0]
|
||||
assert isinstance(current_return, ToolReturnPart)
|
||||
assert current_return.content == "current evidence"
|
||||
assert _record(deps, "rag").question == 7
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resuming_without_a_stored_identity_fails_instead_of_guessing(
|
||||
temp_db_path,
|
||||
):
|
||||
"""Adopting the message count would relabel a question already in progress.
|
||||
|
||||
Every declaration and epoch comparison in it would then be judged against the
|
||||
wrong question, silently. This state is not one the design produces, so it is
|
||||
reported rather than repaired.
|
||||
"""
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
|
||||
async def model(_messages, _info): # pragma: no cover - never reached
|
||||
return ModelResponse(parts=[TextPart("answer")])
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
||||
|
||||
with pytest.raises(RuntimeError, match="no stored question identity"):
|
||||
await agent.run(
|
||||
"carry on",
|
||||
deferred_tool_results=DeferredToolResults(
|
||||
calls={"call-2": "external result"}
|
||||
),
|
||||
message_history=[
|
||||
*_in_flight_history(),
|
||||
ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]),
|
||||
],
|
||||
deps=Deps(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_citing_after_searching_grounds_the_question(temp_db_path):
|
||||
"""The whole rule, end to end, with no compactor and no policy capability."""
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
calls = iter(
|
||||
[
|
||||
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
||||
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
|
||||
[TextPart("answer")],
|
||||
]
|
||||
)
|
||||
|
||||
async def model(_messages, _info):
|
||||
return ModelResponse(parts=next(calls))
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
||||
deps = Deps()
|
||||
|
||||
with patch.object(RAGCapability, "_search", _stub_search):
|
||||
await agent.run("what does the supervisor do?", deps=deps)
|
||||
|
||||
record = _record(deps, "rag")
|
||||
question = record.question
|
||||
assert question is not None
|
||||
assert record.declaration is not None
|
||||
assert [ref.chunk_id for ref in record.declaration.refs] == ["chunk-1"]
|
||||
assert record.occurrences["chunk-1"].retrieved_in_questions == [question]
|
||||
assert citation_status([record], question=question) == "grounded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_searching_after_citing_leaves_the_question_uncited(temp_db_path):
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
calls = iter(
|
||||
[
|
||||
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
||||
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
|
||||
[ToolCallPart("rag_search", {"query": "again"}, "call-3")],
|
||||
[TextPart("answer")],
|
||||
]
|
||||
)
|
||||
|
||||
async def model(_messages, _info):
|
||||
return ModelResponse(parts=next(calls))
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
||||
deps = Deps()
|
||||
|
||||
with patch.object(RAGCapability, "_search", _stub_search):
|
||||
await agent.run("what does the supervisor do?", deps=deps)
|
||||
|
||||
record = _record(deps, "rag")
|
||||
question = record.question
|
||||
assert question is not None
|
||||
assert record.declaration is not None
|
||||
assert citation_status([record], question=question) == "missing"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_citation_in_the_same_request_as_its_search_is_not_current(
|
||||
temp_db_path,
|
||||
):
|
||||
"""Two calls in one response share an epoch, and citing must follow seeing."""
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
calls = iter(
|
||||
[
|
||||
[
|
||||
ToolCallPart("rag_search", {"query": "supervisor"}, "call-1"),
|
||||
ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2"),
|
||||
],
|
||||
[TextPart("answer")],
|
||||
]
|
||||
)
|
||||
|
||||
async def model(_messages, _info):
|
||||
return ModelResponse(parts=next(calls))
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
||||
deps = Deps()
|
||||
|
||||
with patch.object(RAGCapability, "_search", _stub_search):
|
||||
await agent.run("what does the supervisor do?", deps=deps)
|
||||
|
||||
record = _record(deps, "rag")
|
||||
question = record.question
|
||||
assert question is not None
|
||||
assert record.declaration is not None
|
||||
assert record.declaration.epoch == record.latest_evidence_epoch
|
||||
assert citation_status([record], question=question) == "missing"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_evidence_cited_in_two_questions_keeps_both_in_the_record(temp_db_path):
|
||||
"""Occurrences outlive the question that wrote them, through the state dict."""
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
calls = iter(
|
||||
[
|
||||
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
||||
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
|
||||
[TextPart("first answer")],
|
||||
[ToolCallPart("rag_search", {"query": "supervisor again"}, "call-3")],
|
||||
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-4")],
|
||||
[TextPart("second answer")],
|
||||
]
|
||||
)
|
||||
|
||||
async def model(_messages, _info):
|
||||
return ModelResponse(parts=next(calls))
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
||||
deps = Deps()
|
||||
|
||||
with patch.object(RAGCapability, "_search", _stub_search):
|
||||
first = await agent.run("who supervises?", deps=deps)
|
||||
first_question = _record(deps, "rag").question
|
||||
await agent.run(
|
||||
"and who supervises them?", deps=deps, message_history=first.all_messages()
|
||||
)
|
||||
|
||||
record = _record(deps, "rag")
|
||||
assert record.occurrences["chunk-1"].cited_in_questions == [
|
||||
first_question,
|
||||
record.question,
|
||||
]
|
||||
assert record.question != first_question
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_run_with_no_prompt_and_no_history_starts_a_question(temp_db_path):
|
||||
"""An instructions-only run is a first question, not a resumption.
|
||||
|
||||
There is no question in progress to keep an identity for, so nothing is
|
||||
missing and the run proceeds with a fresh one.
|
||||
"""
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
|
||||
async def model(_messages, _info):
|
||||
return ModelResponse(parts=[TextPart("answer")])
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
||||
deps = Deps()
|
||||
|
||||
await agent.run(deps=deps)
|
||||
|
||||
assert _record(deps, "rag").question == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_citing_without_searching_grounds_the_question(temp_db_path):
|
||||
"""A direct chunk-id citation stands on its own, with no evidence outcome.
|
||||
|
||||
Epochs count messages and so start above zero, which is what lets a
|
||||
declaration made in the first request still beat an empty evidence horizon.
|
||||
"""
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
calls = iter(
|
||||
[
|
||||
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-1")],
|
||||
[TextPart("answer")],
|
||||
]
|
||||
)
|
||||
|
||||
async def model(_messages, _info):
|
||||
return ModelResponse(parts=next(calls))
|
||||
|
||||
client = AsyncMock()
|
||||
client.get_chunk_by_id.return_value = Chunk(
|
||||
id="chunk-1", document_id="doc-1", content="evidence"
|
||||
)
|
||||
client.get_document_by_id.return_value = None
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
||||
deps = Deps()
|
||||
|
||||
with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=client)):
|
||||
await agent.run("cite chunk-1", deps=deps)
|
||||
|
||||
record = _record(deps, "rag")
|
||||
question = record.question
|
||||
assert question is not None
|
||||
assert record.latest_evidence_epoch == 0
|
||||
assert record.declaration is not None
|
||||
assert record.declaration.epoch > 0
|
||||
assert record.occurrences["chunk-1"].retrieved_in_questions == []
|
||||
assert citation_status([record], question=question) == "grounded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_host_seeded_record_does_not_pass_for_a_resumption(temp_db_path):
|
||||
"""A default record is truthy, so its presence cannot stand in for identity.
|
||||
|
||||
Seeding one is what a host does when it has no state to send, and taking it
|
||||
at face value would silently answer as question zero.
|
||||
"""
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
|
||||
async def model(_messages, _info): # pragma: no cover - never reached
|
||||
return ModelResponse(parts=[TextPart("answer")])
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
||||
|
||||
with pytest.raises(RuntimeError, match="no stored question identity"):
|
||||
await agent.run(
|
||||
"carry on",
|
||||
message_history=_in_flight_history(),
|
||||
deps=Deps(state={"rag": RAGState().model_dump(mode="json")}),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_resumption_keeps_the_evidence_the_question_already_gathered(
|
||||
temp_db_path,
|
||||
):
|
||||
"""Clearing it would lose the results the model is still answering from.
|
||||
|
||||
A citation after the resumption then records no provenance, and cannot resolve
|
||||
against the expanded search result the model actually saw.
|
||||
"""
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
calls = iter(
|
||||
[
|
||||
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
||||
[TextPart("partial answer")],
|
||||
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-3")],
|
||||
[TextPart("answer")],
|
||||
]
|
||||
)
|
||||
|
||||
async def model(_messages, _info):
|
||||
return ModelResponse(parts=next(calls))
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
||||
deps = Deps()
|
||||
|
||||
with patch.object(RAGCapability, "_search", _stub_search):
|
||||
interrupted = await agent.run("what does the supervisor do?", deps=deps)
|
||||
identity = _record(deps, "rag").question
|
||||
assert identity is not None
|
||||
await agent.run(
|
||||
deferred_tool_results=DeferredToolResults(
|
||||
calls={"call-2": "external result"}
|
||||
),
|
||||
message_history=[
|
||||
*interrupted.all_messages(),
|
||||
ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]),
|
||||
],
|
||||
deps=deps,
|
||||
)
|
||||
|
||||
record = _record(deps, "rag")
|
||||
assert record.question == identity
|
||||
assert record.occurrences["chunk-1"].retrieved_in_questions == [identity]
|
||||
assert citation_status([record], question=identity) == "grounded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_capability_fetches_its_own_evidences_pictures(temp_db_path):
|
||||
"""Compaction rehydrates through the owner, which already holds the connection."""
|
||||
capability = create_rag(db_path=temp_db_path, config=AppConfig())
|
||||
client = AsyncMock()
|
||||
client.document_item_repository.get_picture_bytes.return_value = b"picture-bytes"
|
||||
capability.rag = client
|
||||
|
||||
data = await capability.get_picture_bytes("doc-1", "#/pictures/0")
|
||||
|
||||
assert data == b"picture-bytes"
|
||||
client.document_item_repository.get_picture_bytes.assert_awaited_once_with(
|
||||
"doc-1", "#/pictures/0"
|
||||
)
|
||||
|
|
|
|||
386
tests/capabilities/test_evidence_capsule.py
Normal file
386
tests/capabilities/test_evidence_capsule.py
Normal file
|
|
@ -0,0 +1,386 @@
|
|||
from dataclasses import dataclass, field, replace
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.exceptions import UserError
|
||||
from pydantic_ai.messages import ModelResponse, TextPart
|
||||
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.ledger import (
|
||||
CapabilityEvidenceRecord,
|
||||
EvidenceOccurrence,
|
||||
)
|
||||
from haiku.rag.capabilities.rag import create_capability as create_rag
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
|
||||
@dataclass
|
||||
class Deps:
|
||||
state: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
def citation(chunk_id: str, content: str = "evidence body", pictures=()) -> Citation:
|
||||
return Citation(
|
||||
document_id=f"doc-of-{chunk_id}",
|
||||
chunk_id=chunk_id,
|
||||
document_uri=f"test://{chunk_id}",
|
||||
document_title=f"Title {chunk_id}",
|
||||
content=content,
|
||||
picture_refs=list(pictures),
|
||||
)
|
||||
|
||||
|
||||
def replace_citation(cited: Citation, **changes: Any) -> Citation:
|
||||
return cited.model_copy(update=changes)
|
||||
|
||||
|
||||
def discovered(
|
||||
capability: str = "rag",
|
||||
*,
|
||||
cited: dict[str, list[int]] | None = None,
|
||||
contents: dict[str, str] | None = None,
|
||||
pictures: dict[str, list[str]] | None = None,
|
||||
) -> DiscoveredEvidence:
|
||||
"""One capability's records, as the compactor would find them."""
|
||||
cited = cited or {}
|
||||
contents = contents or {}
|
||||
pictures = pictures or {}
|
||||
record = CapabilityEvidenceRecord(question=max(max(cited.values(), default=[0])))
|
||||
for chunk_id, questions in cited.items():
|
||||
record.occurrences[chunk_id] = EvidenceOccurrence(
|
||||
capability=capability,
|
||||
chunk_id=chunk_id,
|
||||
retrieved_in_questions=list(questions),
|
||||
cited_in_questions=list(questions),
|
||||
)
|
||||
return DiscoveredEvidence(
|
||||
capability=capability,
|
||||
record=record,
|
||||
citations={
|
||||
chunk_id: citation(
|
||||
chunk_id,
|
||||
contents.get(chunk_id, "evidence body"),
|
||||
pictures.get(chunk_id, ()),
|
||||
)
|
||||
for chunk_id in cited
|
||||
},
|
||||
tool_names=frozenset({f"{capability}_search"}),
|
||||
)
|
||||
|
||||
|
||||
def test_nothing_cited_produces_no_capsule():
|
||||
capsule = build_capsule([discovered()])
|
||||
|
||||
assert capsule.text == ""
|
||||
assert capsule.pictures == ()
|
||||
|
||||
|
||||
def test_cited_evidence_is_grouped_newest_question_first():
|
||||
capsule = build_capsule([discovered(cited={"old": [2], "new": [8]})])
|
||||
|
||||
assert capsule.text.index(group_label(1)) < capsule.text.index(group_label(2))
|
||||
assert capsule.text.index("[new]") < capsule.text.index("[old]")
|
||||
assert CAPSULE_HEADER in capsule.text
|
||||
|
||||
|
||||
def test_an_entry_is_rendered_once_in_its_most_recent_citing_group():
|
||||
capsule = build_capsule([discovered(cited={"reused": [2, 8], "only-old": [2]})])
|
||||
|
||||
assert capsule.text.count("[reused]") == 1
|
||||
assert capsule.text.index("[reused]") < capsule.text.index("[only-old]")
|
||||
|
||||
|
||||
def test_evidence_cited_in_one_question_forms_one_group():
|
||||
capsule = build_capsule([discovered(cited={"a": [4], "b": [4]})])
|
||||
|
||||
assert group_label(1) in capsule.text
|
||||
assert group_label(2) not in capsule.text
|
||||
|
||||
|
||||
def test_every_cited_entry_is_kept_whole():
|
||||
"""No budget: a long citation is retained in full rather than truncated."""
|
||||
body = "L" * 20_000
|
||||
capsule = build_capsule([discovered(cited={"long": [4]}, contents={"long": body})])
|
||||
|
||||
assert body in capsule.text
|
||||
|
||||
|
||||
def test_both_capabilities_share_one_capsule():
|
||||
capsule = build_capsule(
|
||||
[
|
||||
discovered("rag", cited={"from-rag": [4]}),
|
||||
discovered("analysis", cited={"from-analysis": [6]}),
|
||||
]
|
||||
)
|
||||
|
||||
assert capsule.text.count(CAPSULE_HEADER) == 1
|
||||
assert "[from-rag]" in capsule.text
|
||||
assert "[from-analysis]" in capsule.text
|
||||
|
||||
|
||||
def test_the_same_chunk_id_under_two_capabilities_is_kept_apart():
|
||||
capsule = build_capsule(
|
||||
[
|
||||
discovered("rag", cited={"shared": [4]}, contents={"shared": "rag body"}),
|
||||
discovered(
|
||||
"analysis", cited={"shared": [4]}, contents={"shared": "analysis body"}
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
assert "rag body" in capsule.text
|
||||
assert "analysis body" in capsule.text
|
||||
|
||||
|
||||
def test_cited_evidence_with_no_canonical_citation_is_an_error():
|
||||
"""Both are written by the same call, so divergence is not a valid state.
|
||||
|
||||
Rendering the rest would quietly drop evidence an answer rested on, against
|
||||
the one guarantee this capsule makes.
|
||||
"""
|
||||
evidence = discovered(cited={"present": [4]})
|
||||
evidence.record.occurrences["absent"] = EvidenceOccurrence(
|
||||
capability="rag", chunk_id="absent", cited_in_questions=[4]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="absent"):
|
||||
build_capsule([evidence])
|
||||
|
||||
|
||||
def test_retrieved_but_uncited_evidence_is_not_kept():
|
||||
evidence = discovered(cited={"cited": [4]})
|
||||
evidence.record.occurrences["seen-only"] = EvidenceOccurrence(
|
||||
capability="rag", chunk_id="seen-only", retrieved_in_questions=[4]
|
||||
)
|
||||
|
||||
capsule = build_capsule([evidence])
|
||||
|
||||
assert "[cited]" in capsule.text
|
||||
assert "seen-only" not in capsule.text
|
||||
|
||||
|
||||
def test_a_source_is_named_once_when_the_title_is_the_uri():
|
||||
"""Real corpora set both to the document id, which reads as a stutter."""
|
||||
evidence = discovered(cited={"a": [4]})
|
||||
evidence.citations["a"].document_title = "2410.11843v5"
|
||||
evidence.citations["a"].document_uri = "2410.11843v5"
|
||||
|
||||
capsule = build_capsule([evidence])
|
||||
|
||||
assert 'Source: "2410.11843v5"' in capsule.text
|
||||
assert "2410.11843v5)" not in capsule.text
|
||||
|
||||
|
||||
def test_pictures_of_cited_evidence_are_all_retained_newest_first():
|
||||
capsule = build_capsule(
|
||||
[
|
||||
discovered(
|
||||
cited={"a": [2], "c": [6]},
|
||||
pictures={"a": ["#/pictures/0"], "c": ["#/pictures/1", "#/pictures/2"]},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert [picture.self_ref for picture in capsule.pictures] == [
|
||||
"#/pictures/1",
|
||||
"#/pictures/2",
|
||||
"#/pictures/0",
|
||||
]
|
||||
assert capsule.pictures[0].document_id == "doc-of-c"
|
||||
assert capsule.pictures[0].capability == "rag"
|
||||
|
||||
|
||||
def test_a_picture_of_uncited_evidence_is_not_retained():
|
||||
found = discovered(cited={"cited": [4]}, pictures={"cited": ["#/pictures/0"]})
|
||||
found.record.occurrences["seen-only"] = EvidenceOccurrence(
|
||||
capability="rag", chunk_id="seen-only", retrieved_in_questions=[4]
|
||||
)
|
||||
evidence = replace(
|
||||
found,
|
||||
citations={
|
||||
**found.citations,
|
||||
"seen-only": citation("seen-only", pictures=["#/pictures/9"]),
|
||||
},
|
||||
)
|
||||
|
||||
capsule = build_capsule([evidence])
|
||||
|
||||
assert [picture.self_ref for picture in capsule.pictures] == ["#/pictures/0"]
|
||||
|
||||
|
||||
def test_one_picture_cited_through_two_chunks_is_attached_once():
|
||||
"""Overlapping chunks of one document share a figure, counted twice by a provider."""
|
||||
found = discovered(
|
||||
cited={"first": [4], "second": [4]},
|
||||
pictures={"first": ["#/pictures/1"], "second": ["#/pictures/1"]},
|
||||
)
|
||||
shared = {
|
||||
chunk_id: replace_citation(cited, document_id="doc-shared")
|
||||
for chunk_id, cited in found.citations.items()
|
||||
}
|
||||
|
||||
capsule = build_capsule([replace(found, citations=shared)])
|
||||
|
||||
assert len(capsule.pictures) == 1
|
||||
assert capsule.pictures[0].chunk_id == "first"
|
||||
|
||||
|
||||
def test_the_same_reference_in_two_documents_is_kept_twice():
|
||||
"""``#/pictures/1`` means a different figure in a different document.
|
||||
|
||||
One capability throughout, so only the document differs: dropping the document
|
||||
from the identity would have to fail this.
|
||||
"""
|
||||
capsule = build_capsule(
|
||||
[
|
||||
discovered(
|
||||
"rag",
|
||||
cited={"a": [4], "b": [4]},
|
||||
pictures={"a": ["#/pictures/1"], "b": ["#/pictures/1"]},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert len(capsule.pictures) == 2
|
||||
assert {picture.capability for picture in capsule.pictures} == {"rag"}
|
||||
assert {picture.document_id for picture in capsule.pictures} == {
|
||||
"doc-of-a",
|
||||
"doc-of-b",
|
||||
}
|
||||
|
||||
|
||||
def test_a_picture_label_names_the_chunk_it_belongs_to():
|
||||
capsule = build_capsule(
|
||||
[discovered(cited={"a": [4]}, pictures={"a": ["#/pictures/0"]})]
|
||||
)
|
||||
|
||||
label = capsule.pictures[0].label
|
||||
assert "[a]" in label
|
||||
assert "#/pictures/0" in label
|
||||
assert "knowledge base" in label
|
||||
assert "Not provided by the user" in label
|
||||
|
||||
|
||||
def _spy_discovery(found: list[list[DiscoveredEvidence]]):
|
||||
"""Discover from ``before_run``, the earliest point the registry is reliable."""
|
||||
original = EvidenceCompactionCapability.before_run
|
||||
|
||||
async def spy(self, ctx):
|
||||
await original(self, ctx)
|
||||
found.append(self.discover(ctx))
|
||||
|
||||
return patch.object(EvidenceCompactionCapability, "before_run", spy)
|
||||
|
||||
|
||||
async def _answer(_messages, _info):
|
||||
return ModelResponse(parts=[TextPart("answer")])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_compactor_discovers_both_evidence_capabilities(temp_db_path):
|
||||
"""Discovery runs one way through the registry, so nothing needs wiring."""
|
||||
compactor = create_compaction()
|
||||
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
|
||||
)
|
||||
found: list[list[DiscoveredEvidence]] = []
|
||||
|
||||
with _spy_discovery(found):
|
||||
agent = Agent(
|
||||
FunctionModel(_answer),
|
||||
deps_type=Deps,
|
||||
capabilities=[rag, analysis, compactor],
|
||||
)
|
||||
await agent.run("a question", deps=Deps())
|
||||
|
||||
assert {evidence.capability: set(evidence.tool_names) for evidence in found[0]} == {
|
||||
"rag": {"rag_search"},
|
||||
"analysis": {"analysis_search", "analysis_execute_code"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discovery_sees_the_run_instances_not_the_registered_ones(temp_db_path):
|
||||
"""A registered capability holds no state; only its per-run copy does."""
|
||||
compactor = create_compaction()
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
found: list[list[DiscoveredEvidence]] = []
|
||||
|
||||
with _spy_discovery(found):
|
||||
agent = Agent(
|
||||
FunctionModel(_answer), deps_type=Deps, capabilities=[rag, compactor]
|
||||
)
|
||||
await agent.run("a question", deps=Deps())
|
||||
|
||||
assert rag.state is None
|
||||
assert found[0][0].record.question == 0
|
||||
|
||||
|
||||
def test_two_compactors_fail_fast(temp_db_path):
|
||||
"""Each would rewrite the same history and each would build its own capsule.
|
||||
|
||||
They share this capability's id, so pydantic-ai refuses at construction and
|
||||
nothing here has to police it.
|
||||
"""
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
|
||||
with pytest.raises(UserError, match="unique within a run"):
|
||||
Agent(
|
||||
FunctionModel(_answer),
|
||||
deps_type=Deps,
|
||||
capabilities=[rag, create_compaction(), create_compaction()],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_compactor_alone_discovers_nothing_and_still_runs():
|
||||
found: list[list[DiscoveredEvidence]] = []
|
||||
|
||||
with _spy_discovery(found):
|
||||
agent = Agent(
|
||||
FunctionModel(_answer), deps_type=Deps, capabilities=[create_compaction()]
|
||||
)
|
||||
result = await agent.run("a question", deps=Deps())
|
||||
|
||||
assert found == [[]]
|
||||
assert result.output == "answer"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_deferred_capability_the_model_never_loaded_has_an_empty_record(
|
||||
temp_db_path,
|
||||
):
|
||||
"""It is still discovered, because every registered capability gets a run copy.
|
||||
|
||||
Nothing was retrieved under it, so its record contributes no entries and the
|
||||
compactor needs no special case for it.
|
||||
"""
|
||||
deferred = create_rag(db_path=temp_db_path, config=AppConfig())
|
||||
found: list[list[DiscoveredEvidence]] = []
|
||||
|
||||
with _spy_discovery(found):
|
||||
agent = Agent(
|
||||
FunctionModel(_answer),
|
||||
deps_type=Deps,
|
||||
capabilities=[deferred, create_compaction()],
|
||||
)
|
||||
await agent.run("a question", deps=Deps())
|
||||
|
||||
assert deferred.defer_loading is True
|
||||
assert deferred.state is None
|
||||
assert [evidence.capability for evidence in found[0]] == ["rag"]
|
||||
assert found[0][0].record.occurrences == {}
|
||||
assert build_capsule(found[0]).text == ""
|
||||
242
tests/capabilities/test_evidence_ledger.py
Normal file
242
tests/capabilities/test_evidence_ledger.py
Normal file
|
|
@ -0,0 +1,242 @@
|
|||
import pytest
|
||||
|
||||
from haiku.rag.capabilities.ledger import (
|
||||
CapabilityEvidenceRecord,
|
||||
CitationDeclaration,
|
||||
EvidenceRef,
|
||||
citation_status,
|
||||
)
|
||||
|
||||
|
||||
def rag_ref(chunk_id: str = "c1") -> EvidenceRef:
|
||||
return EvidenceRef(capability="rag", chunk_id=chunk_id)
|
||||
|
||||
|
||||
def test_a_record_survives_the_state_round_trip():
|
||||
"""Capability state is persisted as JSON, so the schema must survive it.
|
||||
|
||||
A dict keyed by ``(capability, chunk_id)`` does not: the key serialises to
|
||||
``"rag,c1"`` and fails revalidation as a tuple.
|
||||
"""
|
||||
record = CapabilityEvidenceRecord(question=4)
|
||||
record.note_evidence(5)
|
||||
record.declare([rag_ref()], epoch=7, retrieved_now={"c1"})
|
||||
|
||||
restored = CapabilityEvidenceRecord.model_validate(record.model_dump(mode="json"))
|
||||
|
||||
assert restored == record
|
||||
assert citation_status([restored], question=4) == "grounded"
|
||||
assert restored.occurrences["c1"].cited_in_questions == [4]
|
||||
assert restored.occurrences["c1"].retrieved_in_questions == [4]
|
||||
|
||||
|
||||
def test_no_declaration_reads_as_missing():
|
||||
assert citation_status([CapabilityEvidenceRecord()], question=0) == "missing"
|
||||
assert citation_status([], question=0) == "missing"
|
||||
|
||||
|
||||
def test_refs_make_it_grounded_and_no_refs_make_it_ungrounded():
|
||||
grounded = CapabilityEvidenceRecord(question=0)
|
||||
grounded.declare([rag_ref()], epoch=1)
|
||||
|
||||
ungrounded = CapabilityEvidenceRecord(question=0)
|
||||
ungrounded.declare([], epoch=1)
|
||||
|
||||
assert citation_status([grounded], question=0) == "grounded"
|
||||
assert citation_status([ungrounded], question=0) == "ungrounded"
|
||||
|
||||
|
||||
def test_an_earlier_questions_declaration_is_never_current():
|
||||
"""A later question inherits nothing: it has declared nothing yet."""
|
||||
record = CapabilityEvidenceRecord(question=2)
|
||||
record.declare([rag_ref()], epoch=3)
|
||||
assert citation_status([record], question=2) == "grounded"
|
||||
|
||||
record.begin_question(8)
|
||||
|
||||
assert citation_status([record], question=8) == "missing"
|
||||
|
||||
|
||||
def test_a_citation_in_the_same_request_as_the_evidence_is_not_current():
|
||||
"""Citing must follow seeing: equal epochs mean one request."""
|
||||
record = CapabilityEvidenceRecord(question=0)
|
||||
record.note_evidence(5)
|
||||
record.declare([rag_ref()], epoch=5)
|
||||
|
||||
assert citation_status([record], question=0) == "missing"
|
||||
|
||||
record.declare([rag_ref()], epoch=7)
|
||||
|
||||
assert citation_status([record], question=0) == "grounded"
|
||||
|
||||
|
||||
def test_evidence_from_another_capability_after_citing_makes_it_uncited():
|
||||
"""Currency spans capabilities, which only works because epochs are global."""
|
||||
cited = CapabilityEvidenceRecord(question=0)
|
||||
cited.note_evidence(3)
|
||||
cited.declare([rag_ref()], epoch=5)
|
||||
searched_after = CapabilityEvidenceRecord(question=0)
|
||||
searched_after.note_evidence(7)
|
||||
|
||||
assert citation_status([cited], question=0) == "grounded"
|
||||
assert citation_status([cited, searched_after], question=0) == "missing"
|
||||
|
||||
|
||||
def test_declarations_at_the_same_epoch_merge():
|
||||
record = CapabilityEvidenceRecord(question=0)
|
||||
record.declare([rag_ref("c1")], epoch=3)
|
||||
record.declare([rag_ref("c2")], epoch=3)
|
||||
|
||||
assert record.declaration is not None
|
||||
assert [ref.chunk_id for ref in record.declaration.refs] == ["c1", "c2"]
|
||||
|
||||
|
||||
def test_repeating_a_ref_at_the_same_epoch_does_not_duplicate_it():
|
||||
record = CapabilityEvidenceRecord(question=0)
|
||||
record.declare([rag_ref()], epoch=3)
|
||||
record.declare([rag_ref()], epoch=3)
|
||||
|
||||
assert record.declaration is not None
|
||||
assert len(record.declaration.refs) == 1
|
||||
|
||||
|
||||
def test_neither_cite_order_downgrades_a_grounded_declaration():
|
||||
grounded_then_empty = CapabilityEvidenceRecord(question=0)
|
||||
grounded_then_empty.declare([rag_ref()], epoch=3)
|
||||
grounded_then_empty.declare([], epoch=3)
|
||||
|
||||
empty_then_grounded = CapabilityEvidenceRecord(question=0)
|
||||
empty_then_grounded.declare([], epoch=3)
|
||||
empty_then_grounded.declare([rag_ref()], epoch=3)
|
||||
|
||||
assert citation_status([grounded_then_empty], question=0) == "grounded"
|
||||
assert citation_status([empty_then_grounded], question=0) == "grounded"
|
||||
|
||||
|
||||
def test_the_same_chunk_id_under_two_capabilities_stays_separate():
|
||||
rag = CapabilityEvidenceRecord(question=0)
|
||||
rag.declare([EvidenceRef(capability="rag", chunk_id="shared")], epoch=3)
|
||||
analysis = CapabilityEvidenceRecord(question=0)
|
||||
analysis.declare([EvidenceRef(capability="analysis", chunk_id="shared")], epoch=3)
|
||||
|
||||
assert rag.occurrences["shared"].capability == "rag"
|
||||
assert analysis.occurrences["shared"].capability == "analysis"
|
||||
|
||||
|
||||
def test_citing_the_same_chunk_in_two_questions_records_both():
|
||||
record = CapabilityEvidenceRecord(question=2)
|
||||
record.declare([rag_ref()], epoch=3, retrieved_now={"c1"})
|
||||
record.begin_question(8)
|
||||
record.declare([rag_ref()], epoch=9)
|
||||
|
||||
occurrence = record.occurrences["c1"]
|
||||
assert occurrence.cited_in_questions == [2, 8]
|
||||
assert occurrence.retrieved_in_questions == [2]
|
||||
|
||||
|
||||
def test_a_declaration_records_the_question_and_epoch_it_was_made_at():
|
||||
record = CapabilityEvidenceRecord(question=6)
|
||||
record.declare([rag_ref()], epoch=11)
|
||||
|
||||
assert record.declaration == CitationDeclaration(
|
||||
question=6, epoch=11, refs=[rag_ref()]
|
||||
)
|
||||
|
||||
|
||||
def test_a_fresh_record_has_no_question_identity():
|
||||
"""The identity is established by the run, and its absence must be detectable.
|
||||
|
||||
A default record is truthy, so its mere presence cannot stand in for having
|
||||
been through ``for_run``: a host that seeds one would otherwise pass the
|
||||
resumption check with a fabricated identity of zero.
|
||||
"""
|
||||
assert CapabilityEvidenceRecord().question is None
|
||||
|
||||
|
||||
def test_evidence_cannot_move_backwards_in_the_conversation():
|
||||
"""Epochs are message counts, and currency depends on them only growing.
|
||||
|
||||
Silently keeping the newer value would leave every later declaration stale
|
||||
for the rest of the conversation, permanently and invisibly.
|
||||
"""
|
||||
record = CapabilityEvidenceRecord(question=0)
|
||||
record.note_evidence(9)
|
||||
|
||||
with pytest.raises(ValueError, match="append-only"):
|
||||
record.note_evidence(4)
|
||||
|
||||
|
||||
def test_citing_before_a_run_establishes_the_question_is_refused():
|
||||
with pytest.raises(ValueError, match="question identity"):
|
||||
CapabilityEvidenceRecord().declare([rag_ref()], epoch=3)
|
||||
|
||||
|
||||
def test_a_declaration_cannot_move_backwards():
|
||||
"""Otherwise a stale citation replaces a newer one and revives the answer."""
|
||||
record = CapabilityEvidenceRecord(question=0)
|
||||
record.declare([rag_ref("newer")], epoch=5)
|
||||
|
||||
with pytest.raises(ValueError, match="append-only"):
|
||||
record.declare([rag_ref("older")], epoch=3)
|
||||
|
||||
assert record.declaration is not None
|
||||
assert [ref.chunk_id for ref in record.declaration.refs] == ["newer"]
|
||||
|
||||
|
||||
def test_evidence_cannot_predate_a_recorded_declaration():
|
||||
"""The declaration's epoch is a recorded message count like any other."""
|
||||
record = CapabilityEvidenceRecord(question=0)
|
||||
record.declare([rag_ref()], epoch=5)
|
||||
|
||||
with pytest.raises(ValueError, match="append-only"):
|
||||
record.note_evidence(3)
|
||||
|
||||
|
||||
def test_a_question_starts_behind_the_epochs_of_the_one_before_it():
|
||||
"""A question's identity is the history it arrives on, not a continuation.
|
||||
|
||||
Epochs are compared only within the question that recorded them, so a host
|
||||
whose stored history shifted between two questions is answered rather than
|
||||
refused.
|
||||
"""
|
||||
record = CapabilityEvidenceRecord(question=0)
|
||||
record.note_evidence(9)
|
||||
|
||||
record.begin_question(4)
|
||||
|
||||
assert record.question == 4
|
||||
|
||||
|
||||
def test_a_question_cannot_reuse_the_identity_of_the_one_before_it():
|
||||
"""Occurrences outlive their question and are ordered by identity.
|
||||
|
||||
Two questions sharing one identity merge into a single capsule group, and a
|
||||
lower one is rendered as though its evidence were cited earlier.
|
||||
"""
|
||||
record = CapabilityEvidenceRecord(question=4)
|
||||
|
||||
with pytest.raises(ValueError, match="already answered"):
|
||||
record.begin_question(4)
|
||||
|
||||
with pytest.raises(ValueError, match="already answered"):
|
||||
record.begin_question(3)
|
||||
|
||||
assert record.question == 4
|
||||
|
||||
|
||||
def test_a_question_starts_clear_of_the_one_before_it():
|
||||
"""Evidence and declarations describe a single question and end with it.
|
||||
|
||||
Carrying either into the next question makes it answerable by the last
|
||||
question's citations, and freezes its own declarations behind an epoch no
|
||||
message in it can reach.
|
||||
"""
|
||||
record = CapabilityEvidenceRecord(question=4)
|
||||
record.note_evidence(6)
|
||||
record.declare([rag_ref()], epoch=7)
|
||||
|
||||
record.begin_question(9)
|
||||
|
||||
assert record.latest_evidence_epoch == 0
|
||||
assert record.declaration is None
|
||||
assert citation_status([record], question=9) == "missing"
|
||||
729
tests/capabilities/test_evidence_wire.py
Normal file
729
tests/capabilities/test_evidence_wire.py
Normal file
|
|
@ -0,0 +1,729 @@
|
|||
import base64
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import Any, cast
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import Agent, DeferredToolResults, RunContext
|
||||
from pydantic_ai.messages import (
|
||||
BinaryContent,
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
TextPart,
|
||||
ToolCallPart,
|
||||
ToolReturnPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
from pydantic_ai.models import ModelRequestContext, ModelRequestParameters
|
||||
from pydantic_ai.models.function import FunctionModel
|
||||
from pydantic_ai.models.test import TestModel
|
||||
from pydantic_ai.usage import RunUsage
|
||||
|
||||
from haiku.rag.capabilities.compaction import (
|
||||
RECEIPT,
|
||||
Capsule,
|
||||
compact_history,
|
||||
picture_label,
|
||||
)
|
||||
from haiku.rag.capabilities.compaction import create_capability as create_compaction
|
||||
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
|
||||
from haiku.rag.capabilities.rag import RAGCapability, RAGState
|
||||
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
|
||||
|
||||
OWNED = frozenset({"rag_search"})
|
||||
PNG = BinaryContent(data=b"fake-image-bytes", media_type="image/png")
|
||||
|
||||
|
||||
def retrieved_image(chunk_id: str = "chunk-1", self_ref: str = "#/pictures/0"):
|
||||
"""A page image on the wire, labelled the way a search result attaches it."""
|
||||
return UserPromptPart(content=[picture_label(chunk_id, self_ref), PNG])
|
||||
|
||||
|
||||
def answered_question(question: str, *, evidence: str, images: bool = False):
|
||||
"""One settled question: prompt, search, result, answer."""
|
||||
returned: list[Any] = [ToolReturnPart("rag_search", evidence, "call-1")]
|
||||
if images:
|
||||
returned.append(retrieved_image())
|
||||
return [
|
||||
ModelRequest(parts=[UserPromptPart(question)]),
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
||||
ModelRequest(parts=returned),
|
||||
ModelResponse(parts=[TextPart("an answer")]),
|
||||
]
|
||||
|
||||
|
||||
def returns_of(messages) -> list[str]:
|
||||
return [
|
||||
str(part.content)
|
||||
for message in messages
|
||||
if isinstance(message, ModelRequest)
|
||||
for part in message.parts
|
||||
if isinstance(part, ToolReturnPart)
|
||||
]
|
||||
|
||||
|
||||
def images_of(messages) -> list[BinaryContent]:
|
||||
return [
|
||||
item
|
||||
for message in messages
|
||||
if isinstance(message, ModelRequest)
|
||||
for part in message.parts
|
||||
if isinstance(part, UserPromptPart) and not isinstance(part.content, str)
|
||||
for item in part.content
|
||||
if isinstance(item, BinaryContent)
|
||||
]
|
||||
|
||||
|
||||
def texts_of(messages) -> list[str]:
|
||||
return [
|
||||
item
|
||||
for message in messages
|
||||
if isinstance(message, ModelRequest)
|
||||
for part in message.parts
|
||||
if isinstance(part, UserPromptPart) and not isinstance(part.content, str)
|
||||
for item in part.content
|
||||
if isinstance(item, str)
|
||||
]
|
||||
|
||||
|
||||
def test_nothing_before_the_first_question_is_compacted():
|
||||
messages = answered_question("first", evidence="EVIDENCE")
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=0, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert compacted == messages
|
||||
|
||||
|
||||
def test_the_newest_earlier_return_carries_the_capsule():
|
||||
messages = [
|
||||
*answered_question("first", evidence="OLD EVIDENCE"),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert returns_of(compacted) == ["CAPSULE"]
|
||||
assert "OLD EVIDENCE" not in returns_of(compacted)
|
||||
|
||||
|
||||
def test_older_returns_become_receipts_and_only_the_newest_carries_the_capsule():
|
||||
messages = [
|
||||
*answered_question("first", evidence="OLDEST"),
|
||||
*answered_question("second", evidence="NEWER"),
|
||||
ModelRequest(parts=[UserPromptPart("third")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=8, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert returns_of(compacted) == [RECEIPT, "CAPSULE"]
|
||||
|
||||
|
||||
def test_the_current_question_keeps_its_own_evidence():
|
||||
messages = [
|
||||
*answered_question("first", evidence="OLD EVIDENCE"),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-2")]),
|
||||
ModelRequest(parts=[ToolReturnPart("rag_search", "LIVE EVIDENCE", "call-2")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert returns_of(compacted) == ["CAPSULE", "LIVE EVIDENCE"]
|
||||
|
||||
|
||||
def test_another_capabilitys_return_is_left_alone():
|
||||
messages = [
|
||||
ModelRequest(parts=[UserPromptPart("first")]),
|
||||
ModelResponse(parts=[ToolCallPart("other_tool", {}, "call-1")]),
|
||||
ModelRequest(parts=[ToolReturnPart("other_tool", "NOT OURS", "call-1")]),
|
||||
ModelResponse(parts=[TextPart("an answer")]),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert returns_of(compacted) == ["NOT OURS"]
|
||||
|
||||
|
||||
def test_a_cite_acknowledgement_survives():
|
||||
"""A receipt of the model's own action, not evidence."""
|
||||
messages = [
|
||||
ModelRequest(parts=[UserPromptPart("first")]),
|
||||
ModelResponse(
|
||||
parts=[ToolCallPart("rag_cite", {"chunk_ids": ["c1"]}, "call-1")]
|
||||
),
|
||||
ModelRequest(
|
||||
parts=[ToolReturnPart("rag_cite", "Registered 1 citation.", "c1")]
|
||||
),
|
||||
ModelResponse(parts=[TextPart("an answer")]),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert returns_of(compacted) == ["Registered 1 citation."]
|
||||
|
||||
|
||||
def test_an_uncited_earlier_image_is_dropped_with_its_label():
|
||||
messages = [
|
||||
*answered_question("first", evidence="OLD", images=True),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert images_of(compacted) == []
|
||||
assert texts_of(compacted) == []
|
||||
|
||||
|
||||
def test_cited_pictures_are_attached_beside_the_capsule():
|
||||
messages = [
|
||||
*answered_question("first", evidence="OLD", images=True),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
fresh = BinaryContent(data=b"cited-bytes", media_type="image/png")
|
||||
|
||||
compacted = compact_history(
|
||||
messages,
|
||||
boundary=4,
|
||||
owned_tools=OWNED,
|
||||
capsule_text="CAPSULE",
|
||||
capsule_images=[picture_label("cited-chunk", "#/pictures/3"), fresh],
|
||||
)
|
||||
|
||||
assert images_of(compacted) == [fresh]
|
||||
assert texts_of(compacted) == [picture_label("cited-chunk", "#/pictures/3")]
|
||||
carrier = [
|
||||
index
|
||||
for index, message in enumerate(compacted)
|
||||
if isinstance(message, ModelRequest)
|
||||
and any(
|
||||
isinstance(part, ToolReturnPart) and part.content == "CAPSULE"
|
||||
for part in message.parts
|
||||
)
|
||||
]
|
||||
attached = [
|
||||
index
|
||||
for index, message in enumerate(compacted)
|
||||
if isinstance(message, ModelRequest)
|
||||
and any(
|
||||
isinstance(part, UserPromptPart) and not isinstance(part.content, str)
|
||||
for part in message.parts
|
||||
)
|
||||
]
|
||||
assert carrier == attached
|
||||
|
||||
|
||||
def test_a_user_attached_image_is_never_dropped():
|
||||
"""The user's own picture is not ours to remove, even in an earlier question."""
|
||||
mine = UserPromptPart(content=["look at this", PNG])
|
||||
messages = [
|
||||
ModelRequest(parts=[mine]),
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
||||
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
|
||||
ModelResponse(parts=[TextPart("an answer")]),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert images_of(compacted) == [PNG]
|
||||
assert texts_of(compacted) == ["look at this"]
|
||||
|
||||
|
||||
def test_the_stored_messages_are_never_mutated():
|
||||
messages = [
|
||||
*answered_question("first", evidence="OLD EVIDENCE", images=True),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
before = [list(message.parts) for message in messages]
|
||||
|
||||
compact_history(messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE")
|
||||
|
||||
assert [list(message.parts) for message in messages] == before
|
||||
assert "OLD EVIDENCE" in returns_of(messages)
|
||||
|
||||
|
||||
def test_nothing_cited_leaves_only_receipts():
|
||||
messages = [
|
||||
*answered_question("first", evidence="OLD"),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=4, owned_tools=OWNED, capsule_text=""
|
||||
)
|
||||
|
||||
assert returns_of(compacted) == [RECEIPT]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Deps:
|
||||
state: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
def rag_and_compactor(temp_db_path):
|
||||
return (
|
||||
create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False),
|
||||
create_compaction(),
|
||||
)
|
||||
|
||||
|
||||
def in_flight_history() -> list[Any]:
|
||||
"""A question already asked and searched, still awaiting its answer."""
|
||||
return [
|
||||
ModelRequest(parts=[UserPromptPart("what does the supervisor do?")]),
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "s"}, "call-1")]),
|
||||
ModelRequest(
|
||||
parts=[ToolReturnPart("rag_search", "EVIDENCE FOR THE LIVE TURN", "call-1")]
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def resuming_deps(question: int = 0) -> Deps:
|
||||
"""State as a resumption always finds it: the question already identified."""
|
||||
return Deps(
|
||||
state={
|
||||
"rag": RAGState(
|
||||
evidence=CapabilityEvidenceRecord(question=question)
|
||||
).model_dump(mode="json")
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_the_compactor_the_history_is_untouched(temp_db_path):
|
||||
"""Omission is the switch: there is no flag to test, only absence."""
|
||||
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
|
||||
wire: list[list[Any]] = []
|
||||
|
||||
async def model(messages, _info):
|
||||
wire.append(list(messages))
|
||||
return ModelResponse(parts=[TextPart("answer")])
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
|
||||
settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])]
|
||||
|
||||
await agent.run("a different question", deps=Deps(), message_history=settled)
|
||||
|
||||
assert returns_of(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_with_the_compactor_a_new_question_compacts_the_previous_one(
|
||||
temp_db_path,
|
||||
):
|
||||
rag, compactor = rag_and_compactor(temp_db_path)
|
||||
wire: list[list[Any]] = []
|
||||
|
||||
async def model(messages, _info):
|
||||
wire.append(list(messages))
|
||||
return ModelResponse(parts=[TextPart("answer")])
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
|
||||
settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])]
|
||||
|
||||
await agent.run("a different question", deps=Deps(), message_history=settled)
|
||||
|
||||
assert returns_of(wire[-1]) == [RECEIPT]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"resume_kwargs",
|
||||
[
|
||||
pytest.param({}, id="no prompt"),
|
||||
pytest.param(
|
||||
{"deferred_tool_results": DeferredToolResults()}, id="deferred results"
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_resumed_question_keeps_the_evidence_it_is_answering_from(
|
||||
temp_db_path, resume_kwargs
|
||||
):
|
||||
"""The boundary is the stored identity of the question in progress.
|
||||
|
||||
An earlier question below it is compacted; the evidence the model is still
|
||||
answering from sits above it and survives. Deriving the boundary from message
|
||||
shape instead would put the live evidence below it and answer with a receipt
|
||||
where the search result should be.
|
||||
"""
|
||||
rag, compactor = rag_and_compactor(temp_db_path)
|
||||
wire: list[list[Any]] = []
|
||||
|
||||
async def model(messages, _info):
|
||||
wire.append(list(messages))
|
||||
return ModelResponse(parts=[TextPart("answer")])
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
|
||||
history = [
|
||||
*answered_question("an earlier question", evidence="EVIDENCE FOR THE OLD TURN"),
|
||||
*in_flight_history(),
|
||||
]
|
||||
|
||||
await agent.run(
|
||||
deps=resuming_deps(question=4), message_history=history, **resume_kwargs
|
||||
)
|
||||
|
||||
assert returns_of(wire[-1]) == [RECEIPT, "EVIDENCE FOR THE LIVE TURN"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compaction_never_reaches_the_stored_message_history(temp_db_path):
|
||||
"""Rewriting is for the wire; hosts keep the evidence they gathered."""
|
||||
rag, compactor = rag_and_compactor(temp_db_path)
|
||||
turns = iter(
|
||||
[
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
||||
ModelResponse(parts=[TextPart("first answer")]),
|
||||
ModelResponse(parts=[TextPart("second answer")]),
|
||||
]
|
||||
)
|
||||
|
||||
async def model(_messages, _info):
|
||||
return next(turns)
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
|
||||
deps = Deps()
|
||||
|
||||
with patch.object(
|
||||
RAGCapability, "_search", AsyncMock(return_value="REAL EVIDENCE")
|
||||
):
|
||||
first = await agent.run("old question", deps=deps)
|
||||
second = await agent.run(
|
||||
"current question", deps=deps, message_history=first.all_messages()
|
||||
)
|
||||
|
||||
assert "REAL EVIDENCE" in returns_of(second.all_messages())
|
||||
assert RECEIPT not in returns_of(second.all_messages())
|
||||
|
||||
|
||||
REAL_PNG = base64.b64decode(
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
|
||||
)
|
||||
|
||||
|
||||
async def _search_with_a_picture(self, query: str, _limit: int | None) -> str:
|
||||
"""Record a result carrying a page image, the way a real search does."""
|
||||
cast(Any, self.state).searches[query] = [
|
||||
SearchResult(
|
||||
content="evidence",
|
||||
score=1.0,
|
||||
chunk_id="chunk-1",
|
||||
document_id="doc-1",
|
||||
doc_item_refs=["#/pictures/0"],
|
||||
)
|
||||
]
|
||||
self._note_evidence()
|
||||
return "EVIDENCE"
|
||||
|
||||
|
||||
async def _cite_a_picture_chunk(temp_db_path, fetched: bytes | None):
|
||||
"""Two questions: cite a picture chunk, then ask something else."""
|
||||
rag, compactor = rag_and_compactor(temp_db_path)
|
||||
calls = iter(
|
||||
[
|
||||
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
||||
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
|
||||
[TextPart("first answer")],
|
||||
[TextPart("second answer")],
|
||||
]
|
||||
)
|
||||
wire: list[list[Any]] = []
|
||||
|
||||
async def model(messages, _info):
|
||||
wire.append(list(messages))
|
||||
return ModelResponse(parts=next(calls))
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
|
||||
deps = Deps()
|
||||
|
||||
with (
|
||||
patch.object(RAGCapability, "_search", _search_with_a_picture),
|
||||
patch.object(
|
||||
RAGCapability, "get_picture_bytes", AsyncMock(return_value=fetched)
|
||||
),
|
||||
):
|
||||
first = await agent.run("what does the supervisor do?", deps=deps)
|
||||
await agent.run(
|
||||
"and what else?", deps=deps, message_history=first.all_messages()
|
||||
)
|
||||
return wire
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_cited_picture_is_fetched_and_attached_with_its_label(temp_db_path):
|
||||
wire = await _cite_a_picture_chunk(temp_db_path, REAL_PNG)
|
||||
|
||||
assert [picture.data for picture in images_of(wire[-1])] == [REAL_PNG]
|
||||
assert texts_of(wire[-1]) == [picture_label("chunk-1", "#/pictures/0")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_picture_that_cannot_be_fetched_emits_neither_image_nor_label(
|
||||
temp_db_path,
|
||||
):
|
||||
"""A label without its picture tells the model a figure is there when it is not."""
|
||||
wire = await _cite_a_picture_chunk(temp_db_path, None)
|
||||
|
||||
assert images_of(wire[-1]) == []
|
||||
assert texts_of(wire[-1]) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_picture_that_will_not_decode_emits_neither_image_nor_label(
|
||||
temp_db_path,
|
||||
):
|
||||
"""One vision placeholder is rendered per attachment, so a corrupt one miscounts."""
|
||||
wire = await _cite_a_picture_chunk(temp_db_path, b"not-an-image")
|
||||
|
||||
assert images_of(wire[-1]) == []
|
||||
assert texts_of(wire[-1]) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_capsule_is_built_once_per_request_and_again_for_the_next(
|
||||
temp_db_path,
|
||||
):
|
||||
"""Two hook passes for one request must not rebuild; the next request must."""
|
||||
rag, compactor = rag_and_compactor(temp_db_path)
|
||||
deps = Deps()
|
||||
ctx = RunContext(
|
||||
deps=deps, model=TestModel(), usage=RunUsage(), run_id="run-1", run_step=1
|
||||
)
|
||||
run_rag = await rag.for_run(ctx)
|
||||
run_compactor = await compactor.for_run(ctx)
|
||||
cast(Any, run_rag.state).evidence.begin_question(4)
|
||||
ctx = replace(ctx, capabilities={"rag": run_rag, "compaction": run_compactor})
|
||||
builds = 0
|
||||
|
||||
def counting_build(evidence):
|
||||
nonlocal builds
|
||||
builds += 1
|
||||
return Capsule(text="CAPSULE")
|
||||
|
||||
async def handler(_request_context):
|
||||
return ModelResponse(parts=[TextPart("answer")])
|
||||
|
||||
request = ModelRequestContext(
|
||||
messages=[*answered_question("first", evidence="OLD")],
|
||||
model=TestModel(),
|
||||
model_request_parameters=ModelRequestParameters(),
|
||||
model_settings=None,
|
||||
)
|
||||
|
||||
with patch("haiku.rag.capabilities.compaction.build_capsule", counting_build):
|
||||
await run_compactor.wrap_model_request(
|
||||
ctx, request_context=request, handler=handler
|
||||
)
|
||||
await run_compactor.wrap_model_request(
|
||||
ctx, request_context=request, handler=handler
|
||||
)
|
||||
assert builds == 1
|
||||
|
||||
await run_compactor.wrap_model_request(
|
||||
replace(ctx, run_step=2), request_context=request, handler=handler
|
||||
)
|
||||
|
||||
assert builds == 2
|
||||
|
||||
|
||||
def test_a_user_quoting_our_wording_keeps_their_picture_and_their_text():
|
||||
"""Prose is not proof of ownership: a user can write any phrase.
|
||||
|
||||
Recognising our own pictures by a natural-language substring removed a user's
|
||||
image, its text, and with it the whole message part.
|
||||
"""
|
||||
quoted = UserPromptPart(
|
||||
content=[
|
||||
"Here is a page image retrieved from the knowledge base for my report",
|
||||
PNG,
|
||||
]
|
||||
)
|
||||
messages = [
|
||||
ModelRequest(parts=[quoted]),
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
||||
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
|
||||
ModelResponse(parts=[TextPart("an answer")]),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert images_of(compacted) == [PNG]
|
||||
assert texts_of(compacted) == [
|
||||
"Here is a page image retrieved from the knowledge base for my report"
|
||||
]
|
||||
assert all(message.parts for message in compacted)
|
||||
|
||||
|
||||
def test_a_label_of_ours_with_no_picture_after_it_is_kept():
|
||||
"""Only a genuine pair is ours to remove; a lone label is someone else's text."""
|
||||
lonely = UserPromptPart(
|
||||
content=[picture_label("chunk-1", "#/pictures/0"), "and more"]
|
||||
)
|
||||
messages = [
|
||||
ModelRequest(parts=[lonely]),
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
||||
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
|
||||
ModelResponse(parts=[TextPart("an answer")]),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert texts_of(compacted) == [
|
||||
picture_label("chunk-1", "#/pictures/0"),
|
||||
"and more",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_picture_whose_fetch_raises_costs_the_picture_not_the_answer(
|
||||
temp_db_path,
|
||||
):
|
||||
rag, compactor = rag_and_compactor(temp_db_path)
|
||||
calls = iter(
|
||||
[
|
||||
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
|
||||
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
|
||||
[TextPart("first answer")],
|
||||
[TextPart("second answer")],
|
||||
]
|
||||
)
|
||||
wire: list[list[Any]] = []
|
||||
|
||||
async def model(messages, _info):
|
||||
wire.append(list(messages))
|
||||
return ModelResponse(parts=next(calls))
|
||||
|
||||
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
|
||||
deps = Deps()
|
||||
|
||||
with (
|
||||
patch.object(RAGCapability, "_search", _search_with_a_picture),
|
||||
patch.object(
|
||||
RAGCapability,
|
||||
"get_picture_bytes",
|
||||
AsyncMock(side_effect=OSError("the read failed")),
|
||||
),
|
||||
):
|
||||
first = await agent.run("what does the supervisor do?", deps=deps)
|
||||
second = await agent.run(
|
||||
"and what else?", deps=deps, message_history=first.all_messages()
|
||||
)
|
||||
|
||||
assert second.output == "second answer"
|
||||
assert images_of(wire[-1]) == []
|
||||
assert texts_of(wire[-1]) == []
|
||||
|
||||
|
||||
def test_a_request_is_never_left_with_no_parts():
|
||||
"""Emptying a message would leave something that is not a message.
|
||||
|
||||
Our own pictures always travel with the tool return in their request, so this
|
||||
shape does not come from us — but a rewritten history can hold it, and a
|
||||
partless request is invalid whatever produced it.
|
||||
"""
|
||||
ours_alone = ModelRequest(
|
||||
parts=[UserPromptPart(content=[picture_label("chunk-1", "#/pictures/0"), PNG])]
|
||||
)
|
||||
messages = [
|
||||
ours_alone,
|
||||
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
|
||||
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
|
||||
ModelResponse(parts=[TextPart("an answer")]),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert all(message.parts for message in compacted)
|
||||
assert compacted[0] is ours_alone
|
||||
|
||||
|
||||
def test_two_owned_returns_in_one_request_yield_one_capsule():
|
||||
"""A model can search twice in one response, so a request can hold two returns.
|
||||
|
||||
Identifying the carrier by message alone gave every return in it the capsule,
|
||||
which duplicates the whole thing — and it is unbounded.
|
||||
"""
|
||||
messages = [
|
||||
ModelRequest(parts=[UserPromptPart("first")]),
|
||||
ModelResponse(
|
||||
parts=[
|
||||
ToolCallPart("rag_search", {"query": "a"}, "call-1"),
|
||||
ToolCallPart("rag_search", {"query": "b"}, "call-2"),
|
||||
]
|
||||
),
|
||||
ModelRequest(
|
||||
parts=[
|
||||
ToolReturnPart("rag_search", "FIRST EVIDENCE", "call-1"),
|
||||
ToolReturnPart("rag_search", "SECOND EVIDENCE", "call-2"),
|
||||
]
|
||||
),
|
||||
ModelResponse(parts=[TextPart("an answer")]),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
|
||||
compacted = compact_history(
|
||||
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
|
||||
)
|
||||
|
||||
assert returns_of(compacted) == [RECEIPT, "CAPSULE"]
|
||||
|
||||
|
||||
def test_the_capsule_is_attached_beside_the_newest_return_of_that_request():
|
||||
messages = [
|
||||
ModelRequest(parts=[UserPromptPart("first")]),
|
||||
ModelResponse(
|
||||
parts=[
|
||||
ToolCallPart("rag_search", {"query": "a"}, "call-1"),
|
||||
ToolCallPart("rag_search", {"query": "b"}, "call-2"),
|
||||
]
|
||||
),
|
||||
ModelRequest(
|
||||
parts=[
|
||||
ToolReturnPart("rag_search", "FIRST", "call-1"),
|
||||
ToolReturnPart("rag_search", "SECOND", "call-2"),
|
||||
]
|
||||
),
|
||||
ModelResponse(parts=[TextPart("an answer")]),
|
||||
ModelRequest(parts=[UserPromptPart("second")]),
|
||||
]
|
||||
fresh = BinaryContent(data=b"cited-bytes", media_type="image/png")
|
||||
|
||||
compacted = compact_history(
|
||||
messages,
|
||||
boundary=4,
|
||||
owned_tools=OWNED,
|
||||
capsule_text="CAPSULE",
|
||||
capsule_images=[picture_label("cited", "#/pictures/1"), fresh],
|
||||
)
|
||||
|
||||
assert images_of(compacted) == [fresh]
|
||||
assert returns_of(compacted) == [RECEIPT, "CAPSULE"]
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
|
@ -412,3 +413,55 @@ async def test_chat_app_open_failure_surfaces_real_error(tmp_path: Path):
|
|||
with pytest.raises(FileNotFoundError):
|
||||
async with app.run_test():
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_cancelled_run_does_not_advance_persisted_state(temp_db_path: Path):
|
||||
"""State and message history have to move together, or the thread bricks.
|
||||
|
||||
A cancelled run keeps whatever the tools wrote but discards the run's messages.
|
||||
If the state advanced, the next question derives its identity from the shorter
|
||||
history, lands behind the recorded evidence epoch, and is refused as
|
||||
non-append-only — leaving the conversation unusable until cleared.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
app, mock_client = _make_app(temp_db_path)
|
||||
|
||||
class CancellingRun:
|
||||
"""A run that writes evidence through the tools, then is cancelled."""
|
||||
|
||||
def __init__(self, deps):
|
||||
self._deps = deps
|
||||
|
||||
async def __aenter__(self):
|
||||
self._deps.state["rag"] = {
|
||||
"evidence": {"question": 0, "latest_evidence_epoch": 7}
|
||||
}
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_):
|
||||
return False
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
raise asyncio.CancelledError
|
||||
|
||||
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
|
||||
async with app.run_test():
|
||||
app._state = {
|
||||
"rag": {"evidence": {"question": 0, "latest_evidence_epoch": 0}}
|
||||
}
|
||||
before = deepcopy(app._state)
|
||||
|
||||
class Agent:
|
||||
def run_stream_events(self, *_, deps, **__):
|
||||
return CancellingRun(deps)
|
||||
|
||||
app._agent = Agent() # type: ignore[assignment]
|
||||
await app._run_agent("a question")
|
||||
|
||||
assert app._state == before
|
||||
assert app._messages == []
|
||||
|
|
|
|||
|
|
@ -402,8 +402,9 @@ async def test_search_tool_returns_multimodal_when_picture_present():
|
|||
assert isinstance(result.return_value, str)
|
||||
assert "Type: picture" in result.return_value or "rank 1" in result.return_value
|
||||
assert result.content is not None
|
||||
assert len(result.content) == 1
|
||||
part = result.content[0]
|
||||
images = [c for c in result.content if isinstance(c, BinaryContent)]
|
||||
assert len(images) == 1
|
||||
part = images[0]
|
||||
assert isinstance(part, BinaryContent)
|
||||
assert part.media_type == "image/png"
|
||||
assert part.identifier == "#/pictures/0"
|
||||
|
|
@ -457,11 +458,12 @@ async def test_search_tool_attaches_same_self_ref_from_different_documents():
|
|||
|
||||
assert isinstance(result, ToolReturn)
|
||||
assert result.content is not None
|
||||
assert len(result.content) == 2, (
|
||||
images = [c for c in result.content if isinstance(c, BinaryContent)]
|
||||
assert len(images) == 2, (
|
||||
"Both documents' figures must reach the model — dedup keyed on "
|
||||
"self_ref alone would drop doc-B's bytes."
|
||||
)
|
||||
payloads = {part.data for part in result.content} # type: ignore[attr-defined]
|
||||
payloads = {part.data for part in images}
|
||||
assert PICTURE_BYTES in payloads
|
||||
assert other_bytes in payloads
|
||||
|
||||
|
|
@ -905,7 +907,7 @@ async def test_search_tool_drops_invalid_image_bytes():
|
|||
|
||||
assert isinstance(result, ToolReturn)
|
||||
assert result.content is not None
|
||||
identifiers = {p.identifier for p in result.content} # type: ignore[attr-defined]
|
||||
identifiers = {p.identifier for p in result.content if isinstance(p, BinaryContent)}
|
||||
assert identifiers == {"#/pictures/0"}, (
|
||||
"Only the decodable PNG should reach the model — the corrupt "
|
||||
"ref must be dropped so we don't emit a placeholder for an "
|
||||
|
|
|
|||
|
|
@ -210,30 +210,35 @@ def search_config():
|
|||
return Config
|
||||
|
||||
|
||||
class TestBuildBinaryPartsFromResults:
|
||||
"""Picture bytes are attached once per (document, self_ref) pair."""
|
||||
def _png_b64():
|
||||
import base64
|
||||
from io import BytesIO
|
||||
|
||||
from PIL import Image as PILImage
|
||||
|
||||
buf = BytesIO()
|
||||
PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode()
|
||||
|
||||
|
||||
class TestBuildImageContentFromResults:
|
||||
"""Picture bytes are attached once per (document, self_ref) pair, and labelled."""
|
||||
|
||||
def test_results_without_image_data_contribute_nothing(self):
|
||||
from haiku.rag.tools.search import build_binary_parts_from_results
|
||||
from haiku.rag.tools.search import build_image_content_from_results
|
||||
|
||||
results = [
|
||||
SearchResult(content="text only", score=0.5, chunk_id="c1", image_data=None)
|
||||
]
|
||||
|
||||
assert build_binary_parts_from_results(results) == []
|
||||
assert build_image_content_from_results(results) == []
|
||||
|
||||
def test_duplicate_document_and_ref_is_attached_once(self):
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from pydantic_ai.messages import BinaryContent
|
||||
|
||||
from PIL import Image as PILImage
|
||||
from haiku.rag.tools.search import build_image_content_from_results
|
||||
|
||||
from haiku.rag.tools.search import build_binary_parts_from_results
|
||||
|
||||
buf = BytesIO()
|
||||
PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG")
|
||||
png = base64.b64encode(buf.getvalue()).decode()
|
||||
shared = {"#/pictures/0": png}
|
||||
shared = {"#/pictures/0": _png_b64()}
|
||||
results = [
|
||||
SearchResult(
|
||||
content="a",
|
||||
|
|
@ -251,6 +256,53 @@ class TestBuildBinaryPartsFromResults:
|
|||
),
|
||||
]
|
||||
|
||||
parts = build_binary_parts_from_results(results)
|
||||
content = build_image_content_from_results(results)
|
||||
|
||||
assert len(parts) == 1
|
||||
images = [item for item in content if isinstance(item, BinaryContent)]
|
||||
assert len(images) == 1
|
||||
|
||||
def test_each_image_is_labelled_with_the_result_it_belongs_to(self):
|
||||
"""Label every picture, not just the batch.
|
||||
|
||||
``ToolReturn.content`` reaches the model as a user-role message, and one
|
||||
leading note does not override that: with a single note on the wire,
|
||||
gemma4-26b still reasoned "the user also provided images in the prompt".
|
||||
A label adjacent to each picture also names the chunk to cite for it,
|
||||
which ``BinaryContent.identifier`` cannot do — it does not survive
|
||||
serialization to the vision API.
|
||||
"""
|
||||
from pydantic_ai.messages import BinaryContent
|
||||
|
||||
from haiku.rag.tools.search import build_image_content_from_results
|
||||
|
||||
results = [
|
||||
SearchResult(
|
||||
content="a",
|
||||
score=0.9,
|
||||
chunk_id="c1",
|
||||
document_id="doc-1",
|
||||
image_data={"#/pictures/0": _png_b64()},
|
||||
),
|
||||
SearchResult(
|
||||
content="b",
|
||||
score=0.8,
|
||||
chunk_id="c2",
|
||||
document_id="doc-2",
|
||||
image_data={"#/pictures/3": _png_b64()},
|
||||
),
|
||||
]
|
||||
|
||||
content = build_image_content_from_results(results)
|
||||
|
||||
# label, image, label, image — each picture preceded by its own line.
|
||||
assert [type(item) is str for item in content] == [True, False, True, False]
|
||||
assert isinstance(content[1], BinaryContent)
|
||||
assert isinstance(content[3], BinaryContent)
|
||||
|
||||
first, second = content[0], content[2]
|
||||
assert isinstance(first, str) and isinstance(second, str)
|
||||
assert "c1" in first and "#/pictures/0" in first
|
||||
assert "c2" in second and "#/pictures/3" in second
|
||||
assert "1 of 2" in first and "2 of 2" in second
|
||||
for label in (first, second):
|
||||
assert "not provided by the user" in label.lower()
|
||||
|
|
|
|||
Loading…
Reference in a new issue