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

Merge pull request #563 from ggozad/refactor/evidence-state

Browse source 
Give the evidence capabilities one typed state
This commit is contained in:
Yiorgis Gozadinos 2026-08-20 10:27:07 +03:00 committed by GitHub
commit 9d64a0d9f0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 248 additions and 84 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

@ -3,6 +3,7 @@
### Added
- `haiku.rag.capabilities.EvidenceState`: the state base `RAGState` and `AnalysisState` derive from, with `begin_invocation()` for the per-question reset. `RAGCapabilityBase.evidence_record()` and `citation_index()` expose what a capability recorded, so a host reads it without reaching into `capability.state`.
- The `haiku.rag` package declares the `jina` extra, so `provider: jina-local` is supported by declaration rather than through `cross-encoder`'s transitive `transformers` and `torch`. Raises the full package's torch floor to 2.0.
- `providers.docling_serve.timeout` (default 300 seconds), forwarded to the docling-serve client's per-request timeout.

26
evaluations/evaluations/capability_runner.py
View file

@ -1,7 +1,7 @@
from collections.abc import Callable, Iterable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, NamedTuple, Protocol, cast
from typing import Any, NamedTuple
from pydantic_ai import Agent
from pydantic_ai.messages import (
@ -19,12 +19,10 @@ from pydantic_ai.models import Model
from pydantic_ai.capabilities import AbstractCapability
from evaluations.config import Turn
from haiku.rag.capabilities import RAGCapabilityBase
from haiku.rag.capabilities import EvidenceState, RAGCapabilityBase
from haiku.rag.capabilities.compaction import create_capability as create_compaction
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord, citation_status
from haiku.rag.capabilities.ledger import citation_status
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import Citation
CapabilityFactory = Callable[..., RAGCapabilityBase[Any]]
@ -40,14 +38,6 @@ def prefix_to_messages(turns: Iterable[Turn]) -> list[ModelMessage]:
return messages
class _RagLikeState(Protocol):
document_filter: str | None
citation_index: dict[str, Citation]
citations: list[str]
evidence: CapabilityEvidenceRecord
searches: dict[str, list[SearchResult]]
@dataclass
class CapabilityRunResult:
answer: str
@ -145,9 +135,8 @@ def _prepare_agent(
if request_limit is not None:
capability.request_limit = request_limit
state = capability.state_type()
typed = cast(_RagLikeState, state)
if document_filter is not None:
typed.document_filter = document_filter
state.document_filter = document_filter
capabilities: list[AbstractCapability] = [capability]
if compaction:
@ -163,9 +152,8 @@ def _prepare_agent(
def _state_after_run(
capability: RAGCapabilityBase[Any], deps: _EvalDeps
) -> _RagLikeState:
state = capability.state_type.model_validate(deps.state[capability.state_namespace])
return cast(_RagLikeState, state)
) -> EvidenceState:
return capability.state_type.model_validate(deps.state[capability.state_namespace])
async def run_capability_question(
@ -251,7 +239,7 @@ async def run_capability_conversation(
def _result_from_run(
answer: str, typed: _RagLikeState, traffic: ToolTraffic
answer: str, typed: EvidenceState, traffic: ToolTraffic
) -> CapabilityRunResult:
cited_chunk_ids: list[str] = list(typed.citations)
seen_cited: set[str] = set()

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

@ -1,12 +1,13 @@
"""Native Pydantic AI capabilities provided by haiku.rag."""
from haiku.rag.capabilities._base import RAGCapabilityBase
from haiku.rag.capabilities._base import EvidenceState, RAGCapabilityBase
from haiku.rag.capabilities.analysis import AnalysisCapability, AnalysisState
from haiku.rag.capabilities.rag import RAGCapability, RAGState
__all__ = [
"AnalysisCapability",
"AnalysisState",
"EvidenceState",
"RAGCapability",
"RAGCapabilityBase",
"RAGState",

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

@ -3,9 +3,9 @@ import os
from dataclasses import dataclass, field, replace
from difflib import get_close_matches
from pathlib import Path
from typing import Any, cast
from typing import Any
from pydantic import BaseModel
from pydantic import BaseModel, Field
from pydantic_ai import (
DeferredToolRequests,
ModelRetry,
@ -74,18 +74,34 @@ def resolve_db_path(db_path: Path | None, config: AppConfig) -> Path:
return config.storage.data_dir / "haiku.rag.lancedb"
def _clear_invocation_state(state: BaseModel) -> None:
"""Drop the working evidence of the previous question.
class EvidenceState(BaseModel):
"""What a capability accumulates while answering, carried between its runs.
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.
Hosts dump this to JSON and hand it back on the next turn, including over
AG-UI, so the field names and their nesting are a compatibility surface: keep
it flat and don't rename. Key order is not part of it — every carry point
re-validates by key.
"""
for field_name in ("citations", "searches", "executions"):
value = getattr(state, field_name, None)
if hasattr(value, "clear"):
value.clear()
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)
def begin_invocation(self) -> 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.
`document_filter` scopes the conversation, and `evidence` carries question
identity, so neither is working evidence.
"""
self.citations.clear()
self.searches.clear()
def _awaits_the_model(messages: list[ModelMessage]) -> bool:
@ -117,7 +133,7 @@ def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) -
@dataclass
class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
db_path: Path
config: AppConfig
state_type: type[StateT]
@ -163,7 +179,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
outer_state = outer if isinstance(outer, dict) else None
raw_state = outer_state.get(self.state_namespace) if outer_state else None
state = self.state_type.model_validate(raw_state or {})
record = cast(CapabilityEvidenceRecord, cast(Any, state).evidence)
record = state.evidence
continuing = record.in_progress
state_carried = record.question is not None
if not continuing and _awaits_the_model(ctx.messages):
@ -174,7 +190,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
"must be carried between its runs."
)
if not continuing:
_clear_invocation_state(state)
state.begin_invocation()
record.begin_question(len(ctx.messages))
run_capability = replace(
self,
@ -338,7 +354,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
if self.state is not None and not isinstance(
result.output, DeferredToolRequests
):
self._evidence_record().end_question()
self.evidence_record().end_question()
self._sync_state()
await self._close()
return result
@ -392,9 +408,15 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
finally:
self._sync_state()
def _evidence_record(self) -> CapabilityEvidenceRecord:
def evidence_record(self) -> CapabilityEvidenceRecord:
"""What this capability has retrieved and cited, per question."""
assert self.state is not None
return cast(CapabilityEvidenceRecord, cast(Any, self.state).evidence)
return self.state.evidence
def citation_index(self) -> dict[str, Citation]:
"""Citations registered so far, by chunk id."""
assert self.state is not None
return self.state.citation_index
def _note_evidence(self) -> None:
"""Record an outcome the model can ground an answer on.
@ -403,7 +425,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
output: negative evidence grounds a refusal. Excludes a spent budget, which
yields nothing to ground anything on.
"""
self._evidence_record().note_evidence(self.epoch)
self.evidence_record().note_evidence(self.epoch)
def _declare(self, citations: list[Citation]) -> None:
"""Record what the model cited, once the ids have resolved.
@ -411,14 +433,15 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
Declaring earlier would let a call naming only unresolvable ids read as a
grounded answer.
"""
state = cast(Any, self.state)
assert self.state is not None
state = self.state
retrieved = {
result.chunk_id
for results in state.searches.values()
for result in results
if result.chunk_id
}
self._evidence_record().declare(
self.evidence_record().declare(
[
EvidenceRef(capability=self.state_namespace, chunk_id=c.chunk_id)
for c in citations
@ -440,9 +463,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
await self._ensure_rag(),
query,
limit=limit,
document_filter=getattr(self.state, "document_filter", None),
document_filter=self.state.document_filter,
)
state = cast(Any, self.state)
state = self.state
state.searches[query] = results
self._note_evidence()
if self.vision and (parts := build_image_content_from_results(results)):
@ -463,7 +486,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
return "Recorded: this answer cites no knowledge-base evidence."
all_results: list[SearchResult] = []
state = cast(Any, self.state)
state = self.state
for results in state.searches.values():
all_results.extend(results)
known_ids = [result.chunk_id for result in all_results if result.chunk_id]
@ -514,7 +537,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
def _register_citations(self, citations: list[Citation]) -> None:
assert self.state is not None
state = cast(Any, self.state)
state = self.state
next_index = len(state.citation_index) + 1
for citation in citations:
if citation.chunk_id not in state.citation_index:

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

@ -3,7 +3,7 @@ from functools import cache
from pathlib import Path
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field
from pydantic import Field
from pydantic_ai import RunContext, ToolFailed
from pydantic_ai.messages import ToolReturn
from pydantic_ai.toolsets import FunctionToolset
@ -13,14 +13,12 @@ if TYPE_CHECKING:
from haiku.rag.capabilities._base import (
CodeExecutionEntry,
EvidenceState,
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
from haiku.rag.store.models.citation import Citation
STATE_NAMESPACE = "analysis"
_CAPABILITY_ID = "haiku-rag-analysis"
@ -28,13 +26,12 @@ _TOOL_NAMES = frozenset({"analysis_search", "analysis_execute_code", "analysis_c
_instructions_path = Path(__file__).parent / "instructions" / "analysis.md"
class AnalysisState(BaseModel):
document_filter: str | None = None
class AnalysisState(EvidenceState):
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)
def begin_invocation(self) -> None:
super().begin_invocation()
self.executions.clear()
@cache

26
haiku_rag_slim/haiku/rag/capabilities/evidence.py
View file

@ -1,6 +1,6 @@
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, cast
from typing import Any
from pydantic_ai import RunContext
@ -34,18 +34,20 @@ def discover_evidence(ctx: RunContext[Any]) -> list[DiscoveredEvidence]:
carrying state; the registered objects never do. That includes a deferred
capability the model has not loaded, whose record is simply empty.
"""
discovered = [
DiscoveredEvidence(
capability=capability.state_namespace,
record=cast(CapabilityEvidenceRecord, cast(Any, capability.state).evidence),
citations=cast(Any, capability.state).citation_index,
tool_names=frozenset(capability.evidence_tool_names()),
cite_available=capability.cite_available,
state_carried=capability.state_carried,
discovered = []
for capability in ctx.capabilities.values():
if not isinstance(capability, RAGCapabilityBase):
continue
discovered.append(
DiscoveredEvidence(
capability=capability.state_namespace,
record=capability.evidence_record(),
citations=capability.citation_index(),
tool_names=frozenset(capability.evidence_tool_names()),
cite_available=capability.cite_available,
state_carried=capability.state_carried,
)
)
for capability in ctx.capabilities.values()
if isinstance(capability, RAGCapabilityBase)
]
return sorted(discovered, key=lambda evidence: evidence.capability)

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

@ -3,7 +3,6 @@ from functools import cache
from pathlib import Path
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field
from pydantic_ai import RunContext
from pydantic_ai.messages import ToolReturn
from pydantic_ai.toolsets import FunctionToolset
@ -12,13 +11,11 @@ if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
from haiku.rag.capabilities._base import (
EvidenceState,
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
AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
@ -34,12 +31,8 @@ _TOOL_NAMES = frozenset({"rag_search", "rag_cite"})
_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)
class RAGState(EvidenceState):
"""The RAG capability carries nothing beyond the shared evidence fields."""
@cache

11
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -286,11 +286,9 @@ class ChatApp(App):
continue
state_type = RAGState if namespace == RAG_STATE_NAMESPACE else AnalysisState
state = state_type.model_validate(state_data)
cited_ids = getattr(state, "citations", [])
citation_index = getattr(state, "citation_index", {})
for cid in cited_ids:
if cid in citation_index:
citations.append(citation_index[cid])
for cid in state.citations:
if cid in state.citation_index:
citations.append(state.citation_index[cid])
if not citations:
return
@ -314,8 +312,7 @@ class ChatApp(App):
if analysis_data := self._state.get(ANALYSIS_STATE_NAMESPACE):
analysis_state = AnalysisState.model_validate(analysis_data)
executions = getattr(analysis_state, "executions", [])
successful = [e for e in executions if e.success]
successful = [e for e in analysis_state.executions if e.success]
if successful:
await chat_history.add_program(successful[-1].code)

162
tests/capabilities/test_evidence_wire.py
View file

@ -877,3 +877,165 @@ async def test_compaction_proceeds_when_the_capability_without_a_record_has_no_e
await agent.run("a follow-up", deps=rag_only, message_history=history)
assert returns_of(wire[-1]) == [RECEIPT]
# The exact JSON a 0.75.0 host stored or sent over AG-UI. Capability state is
# dumped and re-validated at four carry points, one of them a snapshot the
# browser client sends back on the next turn, so this layout is a wire format.
_STORED_RAG_STATE = {
"citation_index": {},
"citations": [],
"document_filter": None,
"evidence": {
"declaration": None,
"in_progress": False,
"latest_evidence_epoch": 0,
"occurrences": {},
"question": None,
},
"searches": {},
}
def test_stored_state_shape_is_unchanged():
"""A dict stored by an older version still loads, and dumps to the same keys.
Compatibility here is semantic JSON-object equivalence: the same keys, the
same nesting, the same values. Key *order* is not part of the contract —
deriving both states from a shared base reordered `AnalysisState`'s fields,
and nothing serializes, hashes or string-compares this state; every carry
point re-validates it by key.
"""
from haiku.rag.capabilities.rag import RAGState
state = RAGState.model_validate(_STORED_RAG_STATE)
assert state.model_dump(mode="json") == _STORED_RAG_STATE
def _populated(state_type):
"""A state with every field carrying real nested data."""
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import Citation
return state_type(
document_filter="uri LIKE 'test://%'",
citations=["chunk-1"],
citation_index={
"chunk-1": Citation(
chunk_id="chunk-1",
document_id="doc-1",
document_uri="test://doc",
content="cited text",
index=1,
)
},
searches={
"a query": [SearchResult(content="evidence", score=0.9, chunk_id="chunk-1")]
},
evidence=CapabilityEvidenceRecord(question=3, in_progress=True),
)
@pytest.mark.parametrize("namespace", ["rag", "analysis"])
def test_populated_state_round_trips(namespace):
"""Dump, reload, dump again: a host hands this dict back on the next turn,
so the second dump has to equal the first, nested values included."""
from haiku.rag.capabilities.analysis import AnalysisState
from haiku.rag.capabilities.rag import RAGState
state_type = RAGState if namespace == "rag" else AnalysisState
state = _populated(state_type)
first = state.model_dump(mode="json")
second = state_type.model_validate(first).model_dump(mode="json")
assert second == first
assert second["searches"]["a query"][0]["chunk_id"] == "chunk-1"
assert second["citation_index"]["chunk-1"]["document_uri"] == "test://doc"
assert second["evidence"]["question"] == 3
def test_analysis_state_loads_the_old_field_order():
"""A dict written before AnalysisState derived from the shared base lists its
keys in a different order and omits nothing; it still loads."""
from haiku.rag.capabilities.analysis import AnalysisState, CodeExecutionEntry
populated = _populated(AnalysisState)
populated.executions.append(CodeExecutionEntry(code="print(1)", stdout="1"))
dumped = populated.model_dump(mode="json")
old_order = {
key: dumped[key]
for key in (
"document_filter",
"executions",
"citation_index",
"citations",
"evidence",
"searches",
)
}
reloaded = AnalysisState.model_validate(old_order)
assert reloaded == populated
assert reloaded.model_dump(mode="json") == dumped
def test_both_states_are_evidence_states():
from haiku.rag.capabilities._base import EvidenceState
from haiku.rag.capabilities.analysis import AnalysisState
from haiku.rag.capabilities.rag import RAGState
assert issubclass(RAGState, EvidenceState)
assert issubclass(AnalysisState, EvidenceState)
def test_begin_invocation_drops_the_previous_question_working_set():
"""A new question starts from no citations and no results, while the evidence
record — which carries question identity — survives."""
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
from haiku.rag.capabilities.rag import RAGState
from haiku.rag.store.models.citation import Citation
citation = Citation(
chunk_id="chunk-a",
document_id="doc-1",
document_uri="test://doc",
content="cited text",
index=1,
)
state = RAGState(
citations=["chunk-a"],
citation_index={"chunk-a": citation},
searches={"query": []},
document_filter="uri LIKE 'x%'",
evidence=CapabilityEvidenceRecord(question=7),
)
state.begin_invocation()
assert state.citations == []
assert state.searches == {}
# Kept: the filter scopes the whole conversation, the record carries question
# identity, and citation_index keeps citation numbering continuous across
# questions.
assert state.document_filter == "uri LIKE 'x%'"
assert state.evidence.question == 7
assert state.citation_index == {"chunk-a": citation}
def test_begin_invocation_also_drops_analysis_executions():
from haiku.rag.capabilities.analysis import AnalysisState, CodeExecutionEntry
state = AnalysisState(
citations=["chunk-a"],
executions=[CodeExecutionEntry(code="print(1)", stdout="1")],
)
state.begin_invocation()
assert state.executions == []
assert state.citations == []