From 5f6d9e9812731642af3924068480b7886e767841 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 13 Feb 2026 10:23:12 +0200 Subject: [PATCH] handle state_key duplication between deps and ToolContext --- app/backend/main.py | 2 - docs/tools.md | 15 +++--- haiku_rag_slim/haiku/rag/agents/chat/agent.py | 36 +++---------- haiku_rag_slim/haiku/rag/tools/deps.py | 22 ++++---- tests/agents/chat/test_chat_agent.py | 54 ++++++------------- tests/agents/chat/test_features.py | 16 ++++-- tests/tools/test_deps.py | 14 +++-- 7 files changed, 66 insertions(+), 93 deletions(-) diff --git a/app/backend/main.py b/app/backend/main.py index b49beeca..99482f11 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -13,7 +13,6 @@ from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route from haiku.rag.agents.chat import ( - AGUI_STATE_KEY, ChatDeps, create_chat_agent, prepare_chat_context, @@ -92,7 +91,6 @@ async def stream_chat(request: Request) -> Response: config=Config, client=get_client(), tool_context=context, - state_key=AGUI_STATE_KEY, ) adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept) diff --git a/docs/tools.md b/docs/tools.md index f2650503..6468e825 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -235,10 +235,12 @@ async with HaikuRAG("path/to/db.lancedb") as client: print(f"Total search results: {len(search_state.results)}") ``` -`AgentDeps` satisfies the `RAGDeps` protocol and implements the AG-UI state protocol (`state` getter/setter). For AG-UI streaming, pass a `state_key`: +`AgentDeps` satisfies the `RAGDeps` protocol and implements the AG-UI state protocol (`state` getter/setter). For AG-UI streaming, set `state_key` on the `ToolContext` (via `prepare_context`): ```python -deps = AgentDeps(client=client, tool_context=context, state_key="my_app") +context = ToolContext() +prepare_context(context, features=["search", "qa"], state_key="my_app") +deps = AgentDeps(client=client, tool_context=context) ``` Tool functions access `client` and `tool_context` via pydantic-ai's `RunContext.deps`, so toolsets can be created once and reused across requests. @@ -247,7 +249,7 @@ All toolsets respect session-level document filters when a `SessionState` is reg ## AG-UI State Management -Both `AgentDeps` and `ChatDeps` implement the AG-UI `StateHandler` protocol. State is emitted under a namespaced key via `state_key`. +Both `AgentDeps` and `ChatDeps` implement the AG-UI `StateHandler` protocol. `ChatDeps` extends `AgentDeps` with chat-specific config and state handling. State is emitted under a namespaced key via `state_key` on the `ToolContext` — set it once via `prepare_context()`. **Custom agents** use `AgentDeps` + `prepare_context`: @@ -256,14 +258,14 @@ from haiku.rag.tools import AgentDeps, ToolContext, ToolContextCache, prepare_co context = ToolContext() prepare_context(context, features=["search", "qa"], state_key="my_app") -deps = AgentDeps(client=client, tool_context=context, state_key="my_app") +deps = AgentDeps(client=client, tool_context=context) ``` **Chat agent** uses `ChatDeps` + `prepare_chat_context` (adds chat-specific overrides like background summarization and initial context handling): ```python from haiku.rag.agents.chat import ( - AGUI_STATE_KEY, ChatDeps, create_chat_agent, prepare_chat_context, + ChatDeps, create_chat_agent, prepare_chat_context, ) from haiku.rag.tools import ToolContext, ToolContextCache @@ -272,13 +274,12 @@ agent = create_chat_agent(config) # For multi-session apps, cache ToolContext per thread cache = ToolContextCache() context, _is_new = cache.get_or_create(thread_id) -prepare_chat_context(context) # idempotent namespace registration +prepare_chat_context(context) # idempotent; sets state_key="haiku.rag.chat" deps = ChatDeps( config=config, client=client, tool_context=context, - state_key=AGUI_STATE_KEY, # "haiku.rag.chat" ) ``` diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index 36400a41..852afd11 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any from pydantic_ai import Agent @@ -8,9 +8,9 @@ from haiku.rag.agents.chat.context import ( ) from haiku.rag.agents.chat.prompts import build_chat_prompt from haiku.rag.agents.chat.state import AGUI_STATE_KEY -from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig from haiku.rag.tools.context import ToolContext, prepare_context +from haiku.rag.tools.deps import AgentDeps from haiku.rag.tools.document import create_document_toolset from haiku.rag.tools.qa import ( QA_SESSION_NAMESPACE, @@ -34,41 +34,21 @@ def _on_qa_complete(qa_session_state: QASessionState, config: AppConfig) -> None @dataclass -class ChatDeps: +class ChatDeps(AgentDeps): """Dependencies for chat agent. - Implements RAGDeps protocol and StateHandler protocol for AG-UI state management. + Extends AgentDeps with chat-specific config and state handling. """ - config: AppConfig - client: HaikuRAG - tool_context: ToolContext - state_key: str | None = None + config: AppConfig = field(default_factory=AppConfig) - @property - def state(self) -> dict[str, Any]: - """Get current state for AG-UI protocol. - - Combines all registered namespace states into a single flat dict, - matching the ChatSessionState schema expected by AG-UI clients. - """ - snapshot = self.tool_context.build_state_snapshot() - if self.state_key: - return {self.state_key: snapshot} - return snapshot - - @state.setter + @AgentDeps.state.setter def state(self, value: dict[str, Any] | None) -> None: - """Set state from AG-UI protocol.""" + """Set state from AG-UI protocol with chat-specific overrides.""" 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 + state_data = self._extract_state_data(value) # Preserve server's session_context before restore overwrites it qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) diff --git a/haiku_rag_slim/haiku/rag/tools/deps.py b/haiku_rag_slim/haiku/rag/tools/deps.py index b247cb8a..d9ed5123 100644 --- a/haiku_rag_slim/haiku/rag/tools/deps.py +++ b/haiku_rag_slim/haiku/rag/tools/deps.py @@ -14,14 +14,14 @@ class AgentDeps: client: HaikuRAG tool_context: ToolContext - state_key: str | None = None @property def state(self) -> dict[str, Any]: """Get current state for AG-UI protocol.""" snapshot = self.tool_context.build_state_snapshot() - if self.state_key: - return {self.state_key: snapshot} + state_key = self.tool_context.state_key + if state_key: + return {state_key: snapshot} return snapshot @state.setter @@ -29,10 +29,14 @@ class AgentDeps: """Set state from AG-UI protocol.""" if value is None: return - - data: dict[str, Any] = value - if self.state_key and self.state_key in value: - nested = value[self.state_key] - if isinstance(nested, dict): - data = nested + data = self._extract_state_data(value) self.tool_context.restore_state_snapshot(data) + + def _extract_state_data(self, value: dict[str, Any]) -> dict[str, Any]: + """Extract flat state dict, unwrapping state_key if present.""" + state_key = self.tool_context.state_key + if state_key and state_key in value: + nested = value[state_key] + if isinstance(nested, dict): + return nested + return value diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index 771418a5..eaed1556 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/tests/agents/chat/test_chat_agent.py @@ -64,7 +64,17 @@ def test_chat_deps_initialization(temp_db_path): assert deps.config is Config assert deps.client is client assert deps.tool_context is context - assert deps.state_key is None + client.close() + + +def test_chat_deps_is_agent_deps(temp_db_path): + """Test ChatDeps is a subclass of AgentDeps.""" + from haiku.rag.tools.deps import AgentDeps + + client = HaikuRAG(temp_db_path, create=True) + context = ToolContext() + deps = ChatDeps(config=Config, client=client, tool_context=context) + assert isinstance(deps, AgentDeps) client.close() @@ -73,42 +83,18 @@ def test_agui_state_key_constant(): assert AGUI_STATE_KEY == "haiku.rag.chat" -def test_chat_deps_with_state_key(temp_db_path): - """Test ChatDeps can be initialized with state_key for keyed state emission.""" - client = HaikuRAG(temp_db_path, create=True) - context = ToolContext() - deps = ChatDeps( - config=Config, client=client, tool_context=context, state_key="my_state" - ) - - assert deps.config is Config - assert deps.state_key == "my_state" - client.close() - - -def test_chat_deps_state_key_default_none(temp_db_path): - """Test ChatDeps state_key defaults to None.""" - client = HaikuRAG(temp_db_path, create=True) - context = ToolContext() - deps = ChatDeps(config=Config, client=client, tool_context=context) - - assert deps.state_key is None - client.close() - - def test_chat_deps_state_setter_handles_initial_context(temp_db_path): """Test ChatDeps.state setter transfers initial_context to qa_session_state.""" from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState client = HaikuRAG(temp_db_path, create=True) context = ToolContext() + context.state_key = AGUI_STATE_KEY # Register QASessionState (normally done by prepare_chat_context) context.register(QA_SESSION_NAMESPACE, QASessionState()) context.register(SESSION_NAMESPACE, SessionState()) - deps = ChatDeps( - config=Config, client=client, tool_context=context, state_key=AGUI_STATE_KEY - ) + deps = ChatDeps(config=Config, client=client, tool_context=context) # Client sends initial_context with no session_context incoming_state = { @@ -140,12 +126,11 @@ def test_chat_deps_state_setter_parses_session_context_dict(temp_db_path): client = HaikuRAG(temp_db_path, create=True) context = ToolContext() + context.state_key = AGUI_STATE_KEY context.register(QA_SESSION_NAMESPACE, QASessionState()) context.register(SESSION_NAMESPACE, SessionState()) - deps = ChatDeps( - config=Config, client=client, tool_context=context, state_key=AGUI_STATE_KEY - ) + deps = ChatDeps(config=Config, client=client, tool_context=context) # Client sends session_context as a dict (as it comes from JSON) incoming_state = { @@ -178,6 +163,7 @@ def test_chat_deps_state_setter_preserves_server_session_context(temp_db_path): client = HaikuRAG(temp_db_path, create=True) context = ToolContext() + context.state_key = AGUI_STATE_KEY qa_state = QASessionState() qa_state.session_context = SessionContext( summary="Fresh summary from background summarizer" @@ -185,9 +171,7 @@ def test_chat_deps_state_setter_preserves_server_session_context(temp_db_path): context.register(QA_SESSION_NAMESPACE, qa_state) context.register(SESSION_NAMESPACE, SessionState()) - deps = ChatDeps( - config=Config, client=client, tool_context=context, state_key=AGUI_STATE_KEY - ) + deps = ChatDeps(config=Config, client=client, tool_context=context) # Client sends stale session_context incoming_state = { @@ -500,7 +484,6 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path) config=Config, client=client, tool_context=context, - state_key=AGUI_STATE_KEY, ) # Ask a question that should use the ask tool @@ -557,7 +540,6 @@ async def test_chat_agent_ask_triggers_background_summarization( config=Config, client=client, tool_context=context, - state_key=AGUI_STATE_KEY, ) # Patch internal trigger to avoid concurrent HTTP calls during VCR @@ -623,7 +605,6 @@ async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_ config=Config, client=client, tool_context=context, - state_key=AGUI_STATE_KEY, ) # Set initial state with initial_context (mimicking AG-UI client) @@ -725,7 +706,6 @@ async def test_chat_agent_ask_with_prior_answer_retrieval( config=Config, client=client, tool_context=context, - state_key=AGUI_STATE_KEY, ) # First ask - establishes qa_history diff --git a/tests/agents/chat/test_features.py b/tests/agents/chat/test_features.py index fbda785d..21c89134 100644 --- a/tests/agents/chat/test_features.py +++ b/tests/agents/chat/test_features.py @@ -95,6 +95,8 @@ def test_no_qa_skips_qa_session_state(temp_db_path): def test_chat_deps_state_without_qa(temp_db_path): """ChatDeps.state getter omits qa_history/session_context when QASessionState absent.""" + from haiku.rag.agents.chat.state import AGUI_STATE_KEY + client = HaikuRAG(temp_db_path, create=True) context = ToolContext() prepare_chat_context(context, features=[FEATURE_SEARCH]) @@ -102,13 +104,17 @@ def test_chat_deps_state_without_qa(temp_db_path): deps = ChatDeps(config=Config, client=client, tool_context=context) state = deps.state + # State is wrapped under the AGUI state key + assert AGUI_STATE_KEY in state + inner = state[AGUI_STATE_KEY] + # SessionState fields should be present - assert "document_filter" in state - assert "citation_registry" in state - assert "citations" in state + assert "document_filter" in inner + assert "citation_registry" in inner + assert "citations" in inner # QA fields should NOT be present - assert "qa_history" not in state - assert "session_context" not in state + assert "qa_history" not in inner + assert "session_context" not in inner client.close() diff --git a/tests/tools/test_deps.py b/tests/tools/test_deps.py index 2608ac30..61a9a580 100644 --- a/tests/tools/test_deps.py +++ b/tests/tools/test_deps.py @@ -31,10 +31,11 @@ def test_agent_deps_state_getter_with_session(mock_client): def test_agent_deps_state_getter_with_state_key(mock_client): - """state wraps snapshot under state_key when set.""" + """state wraps snapshot under state_key when set on context.""" ctx = ToolContext() + ctx.state_key = "my_app" ctx.register(SESSION_NAMESPACE, SessionState()) - deps = AgentDeps(client=mock_client, tool_context=ctx, state_key="my_app") + deps = AgentDeps(client=mock_client, tool_context=ctx) snapshot = deps.state assert "my_app" in snapshot assert "citations" in snapshot["my_app"] @@ -54,8 +55,9 @@ def test_agent_deps_state_setter_restores(mock_client): def test_agent_deps_state_setter_with_state_key(mock_client): """state setter extracts data from namespaced key.""" ctx = ToolContext() + ctx.state_key = "my_app" ctx.register(SESSION_NAMESPACE, SessionState()) - deps = AgentDeps(client=mock_client, tool_context=ctx, state_key="my_app") + deps = AgentDeps(client=mock_client, tool_context=ctx) deps.state = {"my_app": {"document_filter": ["doc1"]}} session = ctx.get(SESSION_NAMESPACE, SessionState) assert session is not None @@ -76,16 +78,18 @@ def test_agent_deps_state_setter_ignores_none(mock_client): def test_agent_deps_state_roundtrip(mock_client): """Build snapshot then restore produces equivalent state.""" ctx = ToolContext() + ctx.state_key = "app" ctx.register(SESSION_NAMESPACE, SessionState(document_filter=["doc1"])) ctx.register(QA_SESSION_NAMESPACE, QASessionState()) - deps = AgentDeps(client=mock_client, tool_context=ctx, state_key="app") + deps = AgentDeps(client=mock_client, tool_context=ctx) snapshot = deps.state ctx2 = ToolContext() + ctx2.state_key = "app" ctx2.register(SESSION_NAMESPACE, SessionState()) ctx2.register(QA_SESSION_NAMESPACE, QASessionState()) - deps2 = AgentDeps(client=mock_client, tool_context=ctx2, state_key="app") + deps2 = AgentDeps(client=mock_client, tool_context=ctx2) deps2.state = snapshot session = ctx2.get(SESSION_NAMESPACE, SessionState)