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

Merge pull request #278 from ggozad/fix/ask-tool-statedelta

Browse source 
ask tool now emits StateDeltaEvent with client-aware baseline
This commit is contained in:
Yiorgis Gozadinos 2026-02-17 13:53:01 +02:00 committed by GitHub
commit 0919d5babf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 271 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -1,6 +1,10 @@
# Changelog
## [Unreleased]
### Changed
- **AG-UI state sync**: `ask` tool now emits `StateDeltaEvent` (JSON Patch) instead of `StateSnapshotEvent`, consistent with the `search` tool
## [0.30.0] - 2026-02-16
### Added

15
haiku_rag_slim/haiku/rag/tools/context.py
View file

@ -62,6 +62,7 @@ class ToolContext(BaseModel):
state_key: str | None = None
_namespaces: dict[str, BaseModel] = PrivateAttr(default_factory=dict)
_client_snapshot: dict[str, Any] | None = PrivateAttr(default=None)
def register(self, namespace: str, state: BaseModel) -> None:
"""Register state for a namespace.
@ -121,6 +122,16 @@ class ToolContext(BaseModel):
"""List all registered namespaces."""
return list(self._namespaces.keys())
@property
def client_snapshot(self) -> dict[str, Any] | None:
"""Snapshot captured after the last restore_state_snapshot call.
Represents what the client has, before any server-side overrides.
Tools use this as the baseline for delta computation so that
server-side changes (e.g. background summarization) are included.
"""
return self._client_snapshot
def dump_namespaces(self) -> dict[str, dict[str, Any]]:
"""Serialize all namespace states to a dictionary.
@ -150,6 +161,9 @@ class ToolContext(BaseModel):
validates them via the namespace model, and updates the state
in place. Fields not present in *data* are left unchanged.
After restoring, captures a snapshot as ``client_snapshot`` so
tools can compute deltas against what the client actually has.
Args:
data: Flat dict as produced by build_state_snapshot().
"""
@ -163,6 +177,7 @@ class ToolContext(BaseModel):
updated = state.model_validate(current)
for field_name in matching:
setattr(state, field_name, getattr(updated, field_name))
self._client_snapshot = self.build_state_snapshot()
def load_namespace(self, namespace: str, state_type: type[T], data: dict) -> T:
"""Deserialize and register state for a namespace.

34
haiku_rag_slim/haiku/rag/tools/qa.py
View file

@ -1,7 +1,6 @@
import math
from collections.abc import Callable
from ag_ui.core import EventType, StateSnapshotEvent
from pydantic import BaseModel, Field
from pydantic_ai import FunctionToolset, RunContext, ToolReturn
@ -23,6 +22,7 @@ from haiku.rag.tools.session import (
SESSION_NAMESPACE,
SessionContext,
SessionState,
compute_combined_state_delta,
)
PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7
@ -245,9 +245,14 @@ def create_qa_toolset(
tool_context = ctx.deps.tool_context
state_key: str | None = None
client_snapshot: dict | None = None
if tool_context is not None:
state_key = tool_context.state_key
if tool_context.namespaces:
client_snapshot = (
tool_context.client_snapshot or tool_context.build_state_snapshot()
)
qa_result = await run_qa_core(
client=client,
@ -259,21 +264,20 @@ def create_qa_toolset(
on_qa_complete=on_ask_complete,
)
if tool_context is not None and tool_context.namespaces:
snapshot = tool_context.build_state_snapshot()
if state_key:
snapshot = {state_key: snapshot}
answer_text = qa_result.answer
if qa_result.citations:
citation_refs = " ".join(f"[{c.index}]" for c in qa_result.citations)
answer_text = f"{answer_text}\n\nSources: {citation_refs}"
state_event = StateSnapshotEvent(
type=EventType.STATE_SNAPSHOT,
snapshot=snapshot,
if client_snapshot is not None and tool_context is not None:
new_snapshot = tool_context.build_state_snapshot()
state_event = compute_combined_state_delta(
client_snapshot, new_snapshot, state_key=state_key
)
return ToolReturn(return_value=answer_text, metadata=[state_event])
if state_event is not None:
answer_text = qa_result.answer
if qa_result.citations:
citation_refs = " ".join(
f"[{c.index}]" for c in qa_result.citations
)
answer_text = f"{answer_text}\n\nSources: {citation_refs}"
return ToolReturn(return_value=answer_text, metadata=[state_event])
return qa_result

215
tests/agents/chat/test_chat_agent.py
View file

@ -1,7 +1,7 @@
from pathlib import Path
import pytest
from ag_ui.core import StateDeltaEvent, StateSnapshotEvent
from ag_ui.core import StateDeltaEvent
from haiku.rag.agents.chat import (
AGUI_STATE_KEY,
@ -23,7 +23,7 @@ from haiku.rag.tools.session import SESSION_NAMESPACE, SessionContext, SessionSt
def extract_state_from_result(result, state_key: str = AGUI_STATE_KEY) -> dict | None:
"""Extract emitted state from agent result's tool return metadata.
For deltas, applies the patch to an empty state to get the final state.
Applies the JSON Patch delta to an empty state to get the final state.
"""
import jsonpatch
@ -32,10 +32,7 @@ def extract_state_from_result(result, state_key: str = AGUI_STATE_KEY) -> dict |
for part in message.parts:
if hasattr(part, "metadata") and part.metadata:
for meta in part.metadata:
if isinstance(meta, StateSnapshotEvent):
return meta.snapshot.get(state_key)
elif isinstance(meta, StateDeltaEvent):
# Apply delta to empty state to get final state
if isinstance(meta, StateDeltaEvent):
empty_state = {
state_key: ChatSessionState().model_dump(mode="json")
}
@ -1422,3 +1419,209 @@ async def test_summarization_task_cleanup_on_completion():
# Task should be cleaned up
assert key not in _summarization_tasks
@pytest.mark.asyncio
async def test_ask_tool_returns_qa_result_when_no_state_delta(temp_db_path):
"""Test ask tool falls through to QAResult when state delta is empty.
When run_qa_core produces no state changes (client_snapshot == new_snapshot),
compute_combined_state_delta returns None and the tool returns QAResult directly.
"""
from unittest.mock import AsyncMock, patch
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from haiku.rag.tools.models import QAResult
from haiku.rag.tools.qa import create_qa_toolset
async with HaikuRAG(temp_db_path, create=True) as client:
context = ToolContext()
prepare_chat_context(context, features=["search", "qa"])
context.state_key = AGUI_STATE_KEY
deps = ChatDeps(config=Config, client=client, tool_context=context)
toolset = create_qa_toolset(Config)
agent = Agent(
TestModel(call_tools=["ask"]),
deps_type=ChatDeps,
toolsets=[toolset], # ty: ignore[invalid-argument-type]
)
# Mock run_qa_core to return a result WITHOUT modifying any state.
# Since client_snapshot == new_snapshot, the delta is None.
mock_result = QAResult(
question="test", answer="The answer", confidence=0.9, citations=[]
)
with patch(
"haiku.rag.tools.qa.run_qa_core",
new_callable=AsyncMock,
return_value=mock_result,
):
result = await agent.run(
"test question",
deps=deps, # ty: ignore[invalid-argument-type]
)
# The tool should have returned the QAResult (not a ToolReturn)
assert "The answer" in result.output
@pytest.mark.asyncio
async def test_ask_tool_delta_without_citations(temp_db_path):
"""Test ask tool emits StateDeltaEvent without citation sources when citations are empty.
When the QA result has no citations but state DID change,
the tool returns a ToolReturn with a delta but no 'Sources:' line.
"""
from unittest.mock import AsyncMock, patch
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from haiku.rag.tools.models import QAResult
from haiku.rag.tools.qa import (
QA_SESSION_NAMESPACE,
QASessionState,
create_qa_toolset,
)
async with HaikuRAG(temp_db_path, create=True) as client:
context = ToolContext()
prepare_chat_context(context, features=["search", "qa"])
context.state_key = AGUI_STATE_KEY
deps = ChatDeps(config=Config, client=client, tool_context=context)
toolset = create_qa_toolset(Config)
agent = Agent(
TestModel(call_tools=["ask"]),
deps_type=ChatDeps,
toolsets=[toolset], # ty: ignore[invalid-argument-type]
)
mock_result = QAResult(
question="test", answer="No citations answer", confidence=0.9, citations=[]
)
async def mock_qa_core(*args, **kwargs):
# Simulate state mutation (qa_history entry) but no citations
ctx = kwargs.get("context")
if ctx is not None:
qa_state = ctx.get(QA_SESSION_NAMESPACE, QASessionState)
if qa_state is not None:
qa_state.qa_history.append(
QAHistoryEntry(question="test", answer="No citations answer")
)
return mock_result
with patch(
"haiku.rag.tools.qa.run_qa_core",
new_callable=AsyncMock,
side_effect=mock_qa_core,
):
result = await agent.run(
"test question",
deps=deps, # ty: ignore[invalid-argument-type]
)
# Should have the answer but NO "Sources:" line
assert "No citations answer" in result.output
assert "Sources:" not in result.output
# Should have emitted a StateDeltaEvent (qa_history changed)
state = extract_state_from_result(result)
assert state is not None
assert len(state.get("qa_history", [])) > 0
@pytest.mark.asyncio
async def test_ask_tool_delta_with_citations(temp_db_path):
"""Test ask tool emits StateDeltaEvent with 'Sources:' line when citations are present.
When the QA result has citations and state changed, the tool returns
a ToolReturn with a delta and appended 'Sources: [1] [2]' line.
"""
from unittest.mock import AsyncMock, patch
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from haiku.rag.tools.models import QAResult
from haiku.rag.tools.qa import (
QA_SESSION_NAMESPACE,
QASessionState,
create_qa_toolset,
)
async with HaikuRAG(temp_db_path, create=True) as client:
context = ToolContext()
prepare_chat_context(context, features=["search", "qa"])
context.state_key = AGUI_STATE_KEY
deps = ChatDeps(config=Config, client=client, tool_context=context)
toolset = create_qa_toolset(Config)
agent = Agent(
TestModel(call_tools=["ask"]),
deps_type=ChatDeps,
toolsets=[toolset], # ty: ignore[invalid-argument-type]
)
citations = [
Citation(
index=1,
document_id="doc-1",
chunk_id="chunk-a",
document_uri="test.md",
document_title="Test",
content="content",
),
Citation(
index=2,
document_id="doc-1",
chunk_id="chunk-b",
document_uri="test.md",
document_title="Test",
content="content",
),
]
mock_result = QAResult(
question="test",
answer="Cited answer",
confidence=0.95,
citations=citations,
)
async def mock_qa_core(*args, **kwargs):
ctx = kwargs.get("context")
if ctx is not None:
qa_state = ctx.get(QA_SESSION_NAMESPACE, QASessionState)
if qa_state is not None:
qa_state.qa_history.append(
QAHistoryEntry(
question="test",
answer="Cited answer",
citations=citations,
)
)
return mock_result
with patch(
"haiku.rag.tools.qa.run_qa_core",
new_callable=AsyncMock,
side_effect=mock_qa_core,
):
result = await agent.run(
"test question",
deps=deps, # ty: ignore[invalid-argument-type]
)
# Should have the answer WITH "Sources:" line
assert "Cited answer" in result.output
assert "Sources:" in result.output
assert "[1]" in result.output
assert "[2]" in result.output
# Should have emitted a StateDeltaEvent
state = extract_state_from_result(result)
assert state is not None

24
tests/tools/test_context.py
View file

@ -444,6 +444,30 @@ def test_restore_state_snapshot_ignores_unknown_fields():
assert ns1.value == 10
def test_restore_state_snapshot_captures_client_snapshot():
"""restore_state_snapshot stores the restored state as client_snapshot.
This baseline is used by tools to compute deltas against what the
client actually has, so server-side changes (e.g. background
summarization) appear in the delta.
"""
ctx = ToolContext()
ctx.register("ns1", TestState(value=0))
ctx.register("ns2", TestStateWithList(items=[]))
assert ctx.client_snapshot is None
ctx.restore_state_snapshot({"value": 10, "items": ["a"]})
assert ctx.client_snapshot == {"value": 10, "items": ["a"]}
# Mutating state after restore doesn't affect the captured snapshot
ns1 = ctx.get("ns1", TestState)
assert ns1 is not None
ns1.value = 99
assert ctx.client_snapshot == {"value": 10, "items": ["a"]}
# --- prepare_context tests ---