From 0434cd068b7cf6081e5ed491ce9619b9dc21b8fc Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 11 Feb 2026 15:10:50 +0200 Subject: [PATCH] Simplify TUI app --- haiku_rag_slim/haiku/rag/chat/app.py | 164 ++-------- .../haiku/rag/chat/widgets/context_modal.py | 27 +- tests/chat/test_chat_app.py | 305 +----------------- 3 files changed, 55 insertions(+), 441 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 3709bb67..aae5f7e0 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -6,8 +6,6 @@ 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, AgentStreamEvent, @@ -22,10 +20,6 @@ from haiku.rag.agents.chat.agent import ( create_chat_agent, trigger_background_summarization, ) -from haiku.rag.agents.chat.state import ( - AGUI_STATE_KEY, - ChatSessionState, -) from haiku.rag.client import HaikuRAG from haiku.rag.config import get_config from haiku.rag.tools.context import ToolContext @@ -104,13 +98,11 @@ class ChatApp(App): self.client: HaikuRAG | None = None self.config = get_config() self.agent: Agent[ChatDeps, str] | None = None - self.session_state = ChatSessionState() self._is_processing = False self._tool_call_widgets: dict[str, Any] = {} 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.""" @@ -162,17 +154,11 @@ class ChatApp(App): self.tool_context = ToolContext() self.agent = create_chat_agent(self.config, self.client, self.tool_context) - # Initialize session state in tool context + # Sync document filter to tool context session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) if session_state is not None: session_state.document_filter = self._document_filter - # Keep ChatSessionState for UI state sync (used by _sync_session_state) - self.session_state = ChatSessionState( - initial_context=self._initial_context, - document_filter=self._document_filter, - ) - # Focus the input field self.query_one(Input).focus() @@ -181,40 +167,6 @@ class ChatApp(App): if self.client: await self.client.__aexit__(None, None, None) - def _sync_session_state(self, chat_state: dict[str, Any]) -> None: - """Sync session_state from AG-UI state. - - Updates existing session_state with fields from incoming state, - preserving fields not present in the update (e.g., initial_context). - """ - from haiku.rag.agents.research.models import Citation - - # Update specific fields rather than replacing the entire state - if "document_filter" in chat_state: - self.session_state.document_filter = chat_state["document_filter"] - if "citation_registry" in chat_state: - self.session_state.citation_registry = chat_state["citation_registry"] - if "citations" in chat_state: - self.session_state.citations = [ - Citation(**c) if isinstance(c, dict) else c - for c in chat_state["citations"] - ] - if "qa_history" in chat_state: - from haiku.rag.tools.qa import QAHistoryEntry - - self.session_state.qa_history = [ - QAHistoryEntry(**qa) if isinstance(qa, dict) else qa - for qa in chat_state["qa_history"] - ] - if "session_context" in chat_state: - from haiku.rag.agents.chat.state import SessionContext - - ctx = chat_state["session_context"] - if ctx is not None: - self.session_state.session_context = ( - SessionContext(**ctx) if isinstance(ctx, dict) else ctx - ) - async def _handle_stream_event(self, event: AgentStreamEvent) -> None: """Handle streaming events from the agent.""" chat_history = self.query_one(ChatHistory) @@ -232,32 +184,6 @@ class ChatApp(App): widget = self._tool_call_widgets[tool_call_id] chat_history.mark_tool_complete(widget) - # 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 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._sync_session_state(chat_state) - - 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 - ) - self._sync_session_state(chat_state) - async def _event_stream_handler( self, _ctx: RunContext[ChatDeps], @@ -290,7 +216,9 @@ class ChatApp(App): # Clear for new query self._tool_call_widgets.clear() - self.session_state.citations.clear() + session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) + if session_state: + session_state.citations.clear() # Run agent in a worker to keep UI responsive self._is_processing = True @@ -310,39 +238,20 @@ 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") - } - - # Sync session state to tool context before running - session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) - if session_state is not None: - session_state.document_filter = self.session_state.document_filter - session_state.citation_registry = self.session_state.citation_registry - session_state.citations = list(self.session_state.citations) - - # Sync initial_context to QA session state + # Promote initial_context to QA session context on first run from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState qa_session_state = self.tool_context.get( QA_SESSION_NAMESPACE, QASessionState ) if qa_session_state is not None: - if ( - not qa_session_state.session_context - and self.session_state.initial_context - ): - qa_session_state.session_context = ( - self.session_state.initial_context - ) + if not qa_session_state.session_context and self._initial_context: + qa_session_state.session_context = self._initial_context deps = ChatDeps( config=self.config, tool_context=self.tool_context, is_new=False, - state_key=AGUI_STATE_KEY, ) async with self.agent.run_stream( @@ -367,27 +276,14 @@ class ChatApp(App): # Update message history with this conversation self._message_history = stream.all_messages() - # Add citations captured from tool metadata - if self.session_state.citations: - await chat_history.add_citations(self.session_state.citations) + # Add citations from ToolContext + session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) + if session_state and session_state.citations: + await chat_history.add_citations(session_state.citations) # Trigger background summarization trigger_background_summarization(deps) - # Sync session context from QASessionState to ChatSessionState for modal - qa_session_state = self.tool_context.get( - QA_SESSION_NAMESPACE, QASessionState - ) - if qa_session_state is not None and qa_session_state.session_context: - from datetime import datetime - - from haiku.rag.agents.chat.state import SessionContext - - self.session_state.session_context = SessionContext( - summary=qa_session_state.session_context, - last_updated=datetime.now(), - ) - except asyncio.CancelledError: chat_history.hide_thinking() await chat_history.add_message("assistant", "*Cancelled*") @@ -403,16 +299,18 @@ class ChatApp(App): async def action_clear_chat(self) -> None: """Clear the chat history and reset session.""" + from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState + chat_history = self.query_one(ChatHistory) await chat_history.clear_messages() 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( - initial_context=self._initial_context, - document_filter=self._document_filter, + # Re-register fresh states in ToolContext + self.tool_context.register( + SESSION_NAMESPACE, + SessionState(document_filter=self._document_filter), ) + self.tool_context.register(QA_SESSION_NAMESPACE, QASessionState()) def action_focus_input(self) -> None: """Focus the input field, or cancel if processing.""" @@ -462,29 +360,30 @@ class ChatApp(App): async def action_show_context(self) -> None: """Show context modal (edit initial context or view session context).""" + from haiku.rag.agents.chat.state import SessionContext from haiku.rag.chat.widgets.context_modal import ContextModal from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState - # Sync session context from QASessionState before showing modal + session_context = None qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) - if qa_session_state is not None and qa_session_state.session_context: - from datetime import datetime - - from haiku.rag.agents.chat.state import SessionContext - - self.session_state.session_context = SessionContext( + if qa_session_state and qa_session_state.session_context: + session_context = SessionContext( summary=qa_session_state.session_context, last_updated=datetime.now(), ) await self.push_screen( - ContextModal(self.session_state, is_locked=self._context_locked) + ContextModal( + initial_context=self._initial_context, + session_context=session_context, + is_locked=self._context_locked, + ) ) def on_context_modal_context_updated(self, event: Any) -> None: """Handle context updates from modal.""" - if self.session_state and not self._context_locked: - self.session_state.initial_context = event.context or None + if not self._context_locked: + self._initial_context = event.context or None def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None: """Handle citation selection.""" @@ -514,5 +413,6 @@ class ChatApp(App): def on_document_filter_modal_filter_changed(self, event: Any) -> None: """Handle document filter changes from modal.""" self._document_filter = event.selected - if self.session_state: - self.session_state.document_filter = self._document_filter + session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) + if session_state: + session_state.document_filter = self._document_filter diff --git a/haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py b/haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py index 065688d3..c8ec876d 100644 --- a/haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py +++ b/haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py @@ -8,7 +8,7 @@ from textual.screen import ModalScreen from textual.widgets import Button, Markdown, Static, TextArea if TYPE_CHECKING: - from haiku.rag.agents.chat.state import ChatSessionState + from haiku.rag.agents.chat.state import SessionContext class ContextModal(ModalScreen): # pragma: no cover @@ -80,20 +80,20 @@ class ContextModal(ModalScreen): # pragma: no cover self.context = context def __init__( - self, session_state: "ChatSessionState | None", is_locked: bool = False + self, + initial_context: str | None = None, + session_context: "SessionContext | None" = None, + is_locked: bool = False, ) -> None: super().__init__() - self.session_state = session_state + self._initial_context = initial_context + self._session_context = session_context self._is_locked = is_locked @property def _is_edit_mode(self) -> bool: """Edit mode when not locked and no session context yet.""" - has_session_context = ( - self.session_state - and self.session_state.session_context - and self.session_state.session_context.summary - ) + has_session_context = self._session_context and self._session_context.summary return not self._is_locked and not has_session_context def compose(self) -> ComposeResult: @@ -105,9 +105,7 @@ class ContextModal(ModalScreen): # pragma: no cover "This will be locked after you send your first message.", id="context-description", ) - initial_value = "" - if self.session_state and self.session_state.initial_context: - initial_value = self.session_state.initial_context + initial_value = self._initial_context or "" yield TextArea(initial_value, id="context-editor") with Horizontal(id="button-row"): yield Button("Cancel", id="cancel-btn", variant="default") @@ -124,13 +122,10 @@ class ContextModal(ModalScreen): # pragma: no cover yield Button("Close", id="cancel-btn", variant="primary") def _get_session_content(self) -> str: - if not self.session_state: - return "*No session state.*" - - if not self.session_state.session_context: + if not self._session_context: return "*No session context yet. Ask a question first.*" - ctx = self.session_state.session_context + ctx = self._session_context updated = ( ctx.last_updated.strftime("%Y-%m-%d %H:%M:%S") if ctx.last_updated diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 8d87bd00..48f78d44 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -5,6 +5,7 @@ import pytest from typer.testing import CliRunner from haiku.rag.cli import _cli as cli +from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState runner = CliRunner() @@ -273,302 +274,20 @@ async def test_clear_chat_resets_session(temp_db_path: Path): # Verify messages cleared assert len(chat_history.messages) == 0 - # Verify session state reset - assert app.session_state is not None - assert app.session_state.qa_history == [] - assert app.session_state.citations == [] + # Verify ToolContext states are reset + from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState + session_state = app.tool_context.get(SESSION_NAMESPACE, SessionState) + assert session_state is not None + assert session_state.citations == [] + assert session_state.citation_registry == {} -@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: { - "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) - - # Session state should be synced - assert len(app.session_state.citations) == 1 - assert app.session_state.citations[0].chunk_id == "chunk1" - assert app.session_state.citation_registry == {"chunk1": 1} - - -@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: { - "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.session_state.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) - - # Session state should be synced - assert len(app.session_state.citations) == 1 - assert app.session_state.citations[0].chunk_id == "chunk1" - assert app.session_state.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: { - "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) - - # Session state should be synced - assert len(app.session_state.citations) == 1 - assert app.session_state.citations[0].chunk_id == "chunk1" - assert app.session_state.citations[0].content == "Content from first delta" - - -@pytest.mark.asyncio -async def test_handle_stream_event_syncs_session_context(temp_db_path: Path): - """Test that _handle_stream_event syncs session_context to 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 state - app._agui_state_snapshot = { - AGUI_STATE_KEY: { - "citations": [], - "qa_history": [], - "citation_registry": {}, - "document_filter": [], - "initial_context": None, - "session_context": None, - } - } - - # Verify session_context starts as None - assert app.session_state.session_context is None - - # Handle a delta that adds session_context - delta_event = StateDeltaEvent( - type=EventType.STATE_DELTA, - delta=[ - { - "op": "replace", - "path": f"/{AGUI_STATE_KEY}/session_context", - "value": { - "summary": "User asked about Python async patterns.", - "last_updated": "2025-01-15T10:30:00", - }, - }, - ], - ) - tool_return = ToolReturnPart( - tool_name="ask", - content="Answer about async", - tool_call_id="test-call-1", - metadata=[delta_event], - ) - event = FunctionToolResultEvent(result=tool_return) - - await app._handle_stream_event(event) - - # Session context should be synced to session_state - assert app.session_state.session_context is not None - assert ( - app.session_state.session_context.summary - == "User asked about Python async patterns." + qa_session_state = app.tool_context.get( + QA_SESSION_NAMESPACE, QASessionState ) + assert qa_session_state is not None + assert qa_session_state.qa_history == [] + assert qa_session_state.session_context is None @pytest.mark.asyncio