From 205445014229c9ed8fbb68bf97587d4046dd0c5a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 16 Jan 2026 13:55:26 +0200 Subject: [PATCH] ChatDeps implements StateHandler protocol for proper AG-UI state management --- haiku_rag_slim/haiku/rag/agents/chat/state.py | 45 +++- tests/agents/chat/test_state.py | 203 ++++++++++++++++++ tests/test_cli.py | 4 + 3 files changed, 250 insertions(+), 2 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index f60169df..61c23499 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/state.py @@ -1,6 +1,6 @@ import hashlib from dataclasses import dataclass, field -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import numpy as np from numpy.typing import NDArray @@ -156,7 +156,10 @@ async def rank_qa_history_by_similarity( @dataclass class ChatDeps: - """Dependencies for chat agent.""" + """Dependencies for chat agent. + + Implements StateHandler protocol for AG-UI state management. + """ client: HaikuRAG config: AppConfig @@ -164,6 +167,44 @@ class ChatDeps: session_state: ChatSessionState | None = None state_key: str | None = None + @property + def state(self) -> dict[str, Any] | None: + """Get current state for AG-UI protocol.""" + if self.session_state is None: + return None + snapshot = self.session_state.model_dump() + if self.state_key: + return {self.state_key: snapshot} + return snapshot + + @state.setter + def state(self, value: dict[str, Any] | None) -> None: + """Set state from AG-UI protocol.""" + if value is None: + return + # Extract from namespaced key if present + state_data: dict[str, Any] = value + if self.state_key and self.state_key in value: + nested = value[self.state_key] + if isinstance(nested, dict): + state_data = nested + # Update session_state from incoming state + if self.session_state is not None: + if "qa_history" in state_data: + self.session_state.qa_history = [ + QAResponse(**qa) if isinstance(qa, dict) else qa + for qa in state_data.get("qa_history", []) + ] + if "citations" in state_data: + self.session_state.citations = [ + CitationInfo(**c) if isinstance(c, dict) else c + for c in state_data.get("citations", []) + ] + if "initial_context" in state_data: + self.session_state.initial_context = state_data.get("initial_context") + if "session_id" in state_data: + self.session_state.session_id = state_data.get("session_id", "") + @dataclass class SearchDeps: diff --git a/tests/agents/chat/test_state.py b/tests/agents/chat/test_state.py index 60f59855..5e933648 100644 --- a/tests/agents/chat/test_state.py +++ b/tests/agents/chat/test_state.py @@ -267,3 +267,206 @@ def test_chat_session_state_initial_context_defaults_to_none(): state = ChatSessionState(session_id="test-session") assert state.initial_context is None + + +def test_chat_deps_state_getter_returns_namespaced_state(): + """Test ChatDeps.state getter returns state under namespaced key.""" + from unittest.mock import MagicMock + + from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState + + mock_client = MagicMock() + mock_config = MagicMock() + + session_state = ChatSessionState( + session_id="test-123", + qa_history=[ + QAResponse(question="Q1", answer="A1", confidence=0.9), + ], + initial_context="Background info", + ) + + deps = ChatDeps( + client=mock_client, + config=mock_config, + session_state=session_state, + state_key=AGUI_STATE_KEY, + ) + + state = deps.state + assert state is not None + assert AGUI_STATE_KEY in state + assert state[AGUI_STATE_KEY]["session_id"] == "test-123" + assert len(state[AGUI_STATE_KEY]["qa_history"]) == 1 + assert state[AGUI_STATE_KEY]["qa_history"][0]["question"] == "Q1" + assert state[AGUI_STATE_KEY]["initial_context"] == "Background info" + + +def test_chat_deps_state_getter_without_namespace(): + """Test ChatDeps.state getter returns flat state when no state_key.""" + from unittest.mock import MagicMock + + from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState + + mock_client = MagicMock() + mock_config = MagicMock() + + session_state = ChatSessionState(session_id="test-123") + deps = ChatDeps( + client=mock_client, + config=mock_config, + session_state=session_state, + state_key=None, + ) + + state = deps.state + assert state is not None + assert "session_id" in state + assert state["session_id"] == "test-123" + + +def test_chat_deps_state_getter_returns_none_without_session(): + """Test ChatDeps.state getter returns None when no session_state.""" + from unittest.mock import MagicMock + + from haiku.rag.agents.chat.state import ChatDeps + + mock_client = MagicMock() + mock_config = MagicMock() + + deps = ChatDeps( + client=mock_client, + config=mock_config, + session_state=None, + ) + + assert deps.state is None + + +def test_chat_deps_state_setter_updates_from_namespaced_state(): + """Test ChatDeps.state setter updates session_state from namespaced incoming state.""" + from unittest.mock import MagicMock + + from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState + + mock_client = MagicMock() + mock_config = MagicMock() + + session_state = ChatSessionState(session_id="initial") + deps = ChatDeps( + client=mock_client, + config=mock_config, + session_state=session_state, + state_key=AGUI_STATE_KEY, + ) + + # Simulate incoming AG-UI state with namespaced key + incoming_state = { + AGUI_STATE_KEY: { + "session_id": "updated-123", + "qa_history": [ + {"question": "Q1", "answer": "A1", "confidence": 0.9, "citations": []} + ], + "citations": [], + "initial_context": "New context", + } + } + + deps.state = incoming_state + + assert deps.session_state is not None + assert deps.session_state.session_id == "updated-123" + assert len(deps.session_state.qa_history) == 1 + assert deps.session_state.qa_history[0].question == "Q1" + assert deps.session_state.initial_context == "New context" + + +def test_chat_deps_state_setter_handles_none(): + """Test ChatDeps.state setter handles None gracefully.""" + from unittest.mock import MagicMock + + from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState + + mock_client = MagicMock() + mock_config = MagicMock() + + session_state = ChatSessionState(session_id="original") + deps = ChatDeps( + client=mock_client, + config=mock_config, + session_state=session_state, + ) + + # Setting None should not raise and should not change state + deps.state = None + + assert deps.session_state is not None + assert deps.session_state.session_id == "original" + + +def test_chat_deps_state_setter_without_session_state(): + """Test ChatDeps.state setter does nothing when session_state is None.""" + from unittest.mock import MagicMock + + from haiku.rag.agents.chat.state import ChatDeps + + mock_client = MagicMock() + mock_config = MagicMock() + + deps = ChatDeps( + client=mock_client, + config=mock_config, + session_state=None, + ) + + # Should not raise even with valid incoming state + deps.state = {"session_id": "test", "qa_history": [], "citations": []} + + assert deps.session_state is None + + +def test_chat_deps_state_setter_with_citation_dicts(): + """Test ChatDeps.state setter converts citation dicts to CitationInfo.""" + from unittest.mock import MagicMock + + from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState + + mock_client = MagicMock() + mock_config = MagicMock() + + session_state = ChatSessionState(session_id="test") + deps = ChatDeps( + client=mock_client, + config=mock_config, + session_state=session_state, + state_key=AGUI_STATE_KEY, + ) + + incoming_state = { + AGUI_STATE_KEY: { + "session_id": "test", + "qa_history": [], + "citations": [ + { + "index": 1, + "document_id": "doc-1", + "chunk_id": "chunk-1", + "document_uri": "test.md", + "document_title": "Test Doc", + "page_numbers": [1, 2], + "headings": ["Intro"], + "content": "Test content", + } + ], + "initial_context": None, + } + } + + deps.state = incoming_state + + assert deps.session_state is not None + assert len(deps.session_state.citations) == 1 + citation = deps.session_state.citations[0] + assert citation.document_id == "doc-1" + assert citation.chunk_id == "chunk-1" + assert citation.page_numbers == [1, 2] diff --git a/tests/test_cli.py b/tests/test_cli.py index 933d35e5..c92e6882 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -283,6 +283,7 @@ def test_ask(): cite=False, deep=False, filter=None, + initial_context=None, ) @@ -300,6 +301,7 @@ def test_ask_with_cite(): cite=True, deep=False, filter=None, + initial_context=None, ) @@ -317,6 +319,7 @@ def test_ask_with_deep(): cite=False, deep=True, filter=None, + initial_context=None, ) @@ -334,6 +337,7 @@ def test_ask_with_deep_and_cite(): cite=True, deep=True, filter=None, + initial_context=None, )