Give the evidence capabilities one typed state
RAGState and AnalysisState each declared the same five fields, so the generic base could not name them: StateT was bound to BaseModel, and every access went through cast(Any, state), a getattr by string, or a loop clearing fields by name so it could skip the one only AnalysisState has. EvidenceState declares them once. RAGState adds nothing, AnalysisState adds executions and overrides begin_invocation to clear them. StateT binds to EvidenceState, which removes all ten casts and both state-shape getattrs; the three getattr(ctx.deps, "state") probes stay, since those check a host-supplied object rather than our own state. discover_evidence reached into capability.state for two fields. It now asks through evidence_record() and citation_index(), alongside the evidence_tool_names() and cite_available accessors it already used. The eval runner's _RagLikeState protocol and the chat app's getattr reads described this shape from outside and are gone. Compatibility is semantic JSON-object equivalence, not bytes: field names and nesting are unchanged, so a dict stored by 0.75.0 loads and re-dumps equal, but deriving from a shared base reorders AnalysisState's keys. Nothing serializes, hashes or string-compares this state — every carry point re-validates by key.
This commit is contained in:
Yiorgis Gozadinos
parent
78b3231dcd
commit
cab510d0d5
9 changed files with 248 additions and 84 deletions
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 == []
|
||||
|
|
|
|||
Loading…
Reference in a new issue