diff --git a/CHANGELOG.md b/CHANGELOG.md index 0c37dfa7..29d34af2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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. diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py index 613ae439..67d95588 100644 --- a/evaluations/evaluations/capability_runner.py +++ b/evaluations/evaluations/capability_runner.py @@ -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() diff --git a/haiku_rag_slim/haiku/rag/capabilities/__init__.py b/haiku_rag_slim/haiku/rag/capabilities/__init__.py index 654224b6..7b705106 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/__init__.py +++ b/haiku_rag_slim/haiku/rag/capabilities/__init__.py @@ -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", diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 2f0bea64..b359bbc8 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -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: diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index b75fa06c..0b3c860c 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/capabilities/evidence.py b/haiku_rag_slim/haiku/rag/capabilities/evidence.py index 49f8f385..a7848dd2 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/evidence.py +++ b/haiku_rag_slim/haiku/rag/capabilities/evidence.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index f8ea44e9..c45ff6a7 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/haiku_rag_slim/haiku/rag/capabilities/rag.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index beccebcd..d46d6732 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -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) diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py index ef6a1cef..b1dfc769 100644 --- a/tests/capabilities/test_evidence_wire.py +++ b/tests/capabilities/test_evidence_wire.py @@ -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 == []