diff --git a/CHANGELOG.md b/CHANGELOG.md index eb65363f..2b68eb21 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/haiku_rag_slim/haiku/rag/tools/context.py b/haiku_rag_slim/haiku/rag/tools/context.py index 7f1dc0f8..4296e75d 100644 --- a/haiku_rag_slim/haiku/rag/tools/context.py +++ b/haiku_rag_slim/haiku/rag/tools/context.py @@ -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. diff --git a/haiku_rag_slim/haiku/rag/tools/qa.py b/haiku_rag_slim/haiku/rag/tools/qa.py index 0637932a..76d301ac 100644 --- a/haiku_rag_slim/haiku/rag/tools/qa.py +++ b/haiku_rag_slim/haiku/rag/tools/qa.py @@ -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 diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index a1ec782f..659ad1ff 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/tests/agents/chat/test_chat_agent.py @@ -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 diff --git a/tests/tools/test_context.py b/tests/tools/test_context.py index 14ad4de1..28aa827a 100644 --- a/tests/tools/test_context.py +++ b/tests/tools/test_context.py @@ -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 ---