diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index a1aa585a..fad7d675 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -6,6 +6,7 @@ from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Any +import jsonpatch from ag_ui.core import EventType from pydantic_ai import ( Agent, @@ -106,6 +107,7 @@ class ChatApp(App): self._current_worker: Worker[None] | None = None self._message_history: list[ModelMessage] = [] self._document_filter: list[str] = [] + self._agui_state_snapshot: dict[str, Any] = {} def compose(self) -> "ComposeResult": """Compose the UI layout.""" @@ -185,20 +187,33 @@ class ChatApp(App): widget = self._tool_call_widgets[tool_call_id] chat_history.mark_tool_complete(widget) - # Extract citations from StateSnapshotEvent in tool metadata + # Extract citations from state events in tool metadata result = getattr(event, "result", None) metadata = getattr(result, "metadata", None) if result else None if metadata: for meta_event in metadata: - if ( - hasattr(meta_event, "type") - and meta_event.type == EventType.STATE_SNAPSHOT - ): + if not hasattr(meta_event, "type"): + continue + + if meta_event.type == EventType.STATE_SNAPSHOT: snapshot = getattr(meta_event, "snapshot", {}) + self._agui_state_snapshot = snapshot chat_state = snapshot.get(AGUI_STATE_KEY, snapshot) - self._last_citations = [ - Citation(**c) for c in chat_state["citations"] - ] + citations = chat_state.get("citations", []) + self._last_citations = [Citation(**c) for c in citations] + + elif meta_event.type == EventType.STATE_DELTA: + delta = getattr(meta_event, "delta", []) + if delta: + patch = jsonpatch.JsonPatch(delta) + self._agui_state_snapshot = patch.apply( + self._agui_state_snapshot + ) + chat_state = self._agui_state_snapshot.get( + AGUI_STATE_KEY, self._agui_state_snapshot + ) + citations = chat_state.get("citations", []) + self._last_citations = [Citation(**c) for c in citations] async def _event_stream_handler( self, @@ -252,6 +267,12 @@ class ChatApp(App): await chat_history.show_thinking() try: + # Initialize AGUI state snapshot from session state for delta application + if self.session_state: + self._agui_state_snapshot = { + AGUI_STATE_KEY: self.session_state.model_dump(mode="json") + } + deps = ChatDeps( client=self.client, config=self.config, @@ -304,6 +325,7 @@ class ChatApp(App): await chat_history.clear_messages() self._last_citations.clear() self._message_history.clear() + self._agui_state_snapshot = {} # Reset context lock and session state (reset to CLI value) self._context_locked = False self.session_state = ChatSessionState( diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index a842acc4..0f148860 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -282,6 +282,235 @@ async def test_clear_chat_resets_session(temp_db_path: Path): assert app.session_state.session_id != original_session_id +@pytest.mark.asyncio +async def test_handle_stream_event_extracts_citations_from_state_snapshot( + temp_db_path: Path, +): + """Test that _handle_stream_event extracts citations from STATE_SNAPSHOT events.""" + from ag_ui.core import EventType, StateSnapshotEvent + from pydantic_ai import FunctionToolResultEvent + from pydantic_ai.messages import ToolReturnPart + + from haiku.rag.agents.chat.state import AGUI_STATE_KEY + from haiku.rag.chat.app import ChatApp + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test(): + # Create a STATE_SNAPSHOT event with citations + snapshot_event = StateSnapshotEvent( + type=EventType.STATE_SNAPSHOT, + snapshot={ + AGUI_STATE_KEY: { + "session_id": "test", + "citations": [ + { + "index": 1, + "document_id": "doc1", + "chunk_id": "chunk1", + "document_uri": "test.pdf", + "document_title": "Test Doc", + "content": "Test content", + } + ], + "qa_history": [], + "citation_registry": {"chunk1": 1}, + } + }, + ) + + # Create a tool result with the snapshot in metadata + tool_return = ToolReturnPart( + tool_name="search", + content="Found results", + tool_call_id="test-call-1", + metadata=[snapshot_event], + ) + + event = FunctionToolResultEvent(result=tool_return) + + # Handle the event + await app._handle_stream_event(event) + + # Citations should be extracted + assert len(app._last_citations) == 1 + assert app._last_citations[0].chunk_id == "chunk1" + + +@pytest.mark.asyncio +async def test_handle_stream_event_extracts_citations_from_state_delta( + temp_db_path: Path, +): + """Test that _handle_stream_event extracts citations from STATE_DELTA events. + + This test demonstrates that state deltas need to be applied to extract citations. + After the first STATE_SNAPSHOT, subsequent tool calls emit STATE_DELTA events + containing JSON Patch operations. + """ + from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent + from pydantic_ai import FunctionToolResultEvent + from pydantic_ai.messages import ToolReturnPart + + from haiku.rag.agents.chat.state import AGUI_STATE_KEY + from haiku.rag.chat.app import ChatApp + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test(): + # First, handle a STATE_SNAPSHOT to establish initial state + initial_snapshot = StateSnapshotEvent( + type=EventType.STATE_SNAPSHOT, + snapshot={ + AGUI_STATE_KEY: { + "session_id": "test", + "citations": [], + "qa_history": [], + "citation_registry": {}, + } + }, + ) + tool_return1 = ToolReturnPart( + tool_name="search", + content="Initial search", + tool_call_id="test-call-1", + metadata=[initial_snapshot], + ) + event1 = FunctionToolResultEvent(result=tool_return1) + await app._handle_stream_event(event1) + assert len(app._last_citations) == 0 + + # Now handle a STATE_DELTA event that adds citations + delta_event = StateDeltaEvent( + type=EventType.STATE_DELTA, + delta=[ + { + "op": "replace", + "path": f"/{AGUI_STATE_KEY}/citations", + "value": [ + { + "index": 1, + "document_id": "doc1", + "chunk_id": "chunk1", + "document_uri": "test.pdf", + "document_title": "Test Doc", + "content": "Test content from delta", + } + ], + }, + { + "op": "add", + "path": f"/{AGUI_STATE_KEY}/citation_registry/chunk1", + "value": 1, + }, + ], + ) + tool_return2 = ToolReturnPart( + tool_name="ask", + content="Answer with citations", + tool_call_id="test-call-2", + metadata=[delta_event], + ) + event2 = FunctionToolResultEvent(result=tool_return2) + + # Handle the delta event + await app._handle_stream_event(event2) + + # Citations should be extracted from the delta + assert len(app._last_citations) == 1 + assert app._last_citations[0].chunk_id == "chunk1" + assert app._last_citations[0].content == "Test content from delta" + + +@pytest.mark.asyncio +async def test_handle_stream_event_delta_with_preinitialized_state( + temp_db_path: Path, +): + """Test delta handling when _agui_state_snapshot is pre-initialized. + + This is the actual TUI scenario: session_state exists from the start, + so the agent emits deltas (not snapshots) even on the first tool call. + The TUI pre-initializes _agui_state_snapshot from session_state. + """ + from ag_ui.core import EventType, StateDeltaEvent + from pydantic_ai import FunctionToolResultEvent + from pydantic_ai.messages import ToolReturnPart + + from haiku.rag.agents.chat.state import AGUI_STATE_KEY + from haiku.rag.chat.app import ChatApp + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test(): + # Pre-initialize _agui_state_snapshot (simulating what _run_agent does) + app._agui_state_snapshot = { + AGUI_STATE_KEY: { + "session_id": "test", + "citations": [], + "qa_history": [], + "citation_registry": {}, + "document_filter": [], + "initial_context": None, + "session_context": None, + } + } + + # Now handle a STATE_DELTA event directly (no prior snapshot event) + delta_event = StateDeltaEvent( + type=EventType.STATE_DELTA, + delta=[ + { + "op": "replace", + "path": f"/{AGUI_STATE_KEY}/citations", + "value": [ + { + "index": 1, + "document_id": "doc1", + "chunk_id": "chunk1", + "document_uri": "test.pdf", + "document_title": "Test Doc", + "content": "Content from first delta", + } + ], + }, + { + "op": "add", + "path": f"/{AGUI_STATE_KEY}/citation_registry/chunk1", + "value": 1, + }, + ], + ) + tool_return = ToolReturnPart( + tool_name="ask", + content="Answer with citations", + tool_call_id="test-call-1", + metadata=[delta_event], + ) + event = FunctionToolResultEvent(result=tool_return) + + # Handle the delta event + await app._handle_stream_event(event) + + # Citations should be extracted from the delta + assert len(app._last_citations) == 1 + assert app._last_citations[0].chunk_id == "chunk1" + assert app._last_citations[0].content == "Content from first delta" + + @pytest.mark.asyncio async def test_citation_expand_collapse_with_enter(temp_db_path: Path): """Test that pressing Enter on a focused citation toggles expand/collapse."""