Merge pull request #233 from ggozad/feat/agui-state-key

Add state_key parameter for keyed AG-UI state emission

CHANGELOG.md

@@ -3,6 +3,9 @@
 
 ### Added
 
+- **Keyed State Emission for Chat Agent**: New `state_key` parameter in `ChatDeps` for namespaced AG-UI state snapshots
+  - When set, tools emit `{state_key: snapshot}` instead of bare state, enabling state merging when multiple agents share state
+  - Default `None` preserves backwards compatibility (bare state emission)
 - **Page Image Generation Control**: New `generate_page_images` option in `ConversionOptions` to control PDF page image extraction
   - `generate_page_images: bool = True` - Enable/disable rendered page images (used by `visualize_chunk()`)
   - Works with both `docling-local` and `docling-serve` converters

haiku_rag_slim/haiku/rag/agents/chat/agent.py

@@ -102,13 +102,17 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             line += f"\n    {snippet}"
             result_lines.append(line)
 
+        snapshot = new_state.model_dump()
+        if ctx.deps.state_key:
+            snapshot = {ctx.deps.state_key: snapshot}
+
         return ToolReturn(
             return_value=f"Found {len(results)} results:\n\n"
             + "\n\n".join(result_lines),
             metadata=[
                 StateSnapshotEvent(
                     type=EventType.STATE_SNAPSHOT,
-                    snapshot=new_state.model_dump(),
+                    snapshot=snapshot,
                 )
             ],
         )
@@ -241,12 +245,16 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             citation_refs = " ".join(f"[{i + 1}]" for i in range(len(citation_infos)))
             answer_text = f"{answer_text}\n\nSources: {citation_refs}"
 
+        snapshot = new_state.model_dump()
+        if ctx.deps.state_key:
+            snapshot = {ctx.deps.state_key: snapshot}
+
         return ToolReturn(
             return_value=answer_text,
             metadata=[
                 StateSnapshotEvent(
                     type=EventType.STATE_SNAPSHOT,
-                    snapshot=new_state.model_dump(),
+                    snapshot=snapshot,
                 )
             ],
         )

haiku_rag_slim/haiku/rag/agents/chat/state.py

@@ -159,6 +159,7 @@ class ChatDeps:
     config: AppConfig
     search_results: list[SearchResult] | None = None
     session_state: ChatSessionState | None = None
+    state_key: str | None = None
 
 
 @dataclass

tests/agents/chat/test_chat_agent.py

@@ -40,6 +40,28 @@ def test_chat_deps_initialization(temp_db_path):
     client.close()
 
 
+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)
+    deps = ChatDeps(client=client, config=Config, state_key="my_state")
+
+    assert deps.client is client
+    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)
+    deps = ChatDeps(client=client, config=Config)
+
+    assert deps.state_key is None
+
+    client.close()
+
+
 def test_chat_session_state():
     """Test ChatSessionState model."""
     state = ChatSessionState(session_id="test-session")
@@ -265,6 +287,40 @@ async def test_chat_agent_search_tool(allow_model_requests, temp_db_path):
         assert len(result.output) > 0
 
 
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_chat_agent_search_with_state_key(allow_model_requests, temp_db_path):
+    """Test search tool emits keyed state when state_key is set."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+        await client.create_document(
+            content=DOCLAYNET_ANNOTATION,
+            uri="doclaynet-annotation",
+            title="DocLayNet Annotation",
+        )
+
+        agent = create_chat_agent(Config)
+        session_state = ChatSessionState(session_id="test-search")
+        deps = ChatDeps(
+            client=client,
+            config=Config,
+            session_state=session_state,
+            state_key="haiku_rag",
+        )
+
+        result = await agent.run(
+            "Search for documents about class labels",
+            deps=deps,
+        )
+
+        assert result.output is not None
+        assert len(result.output) > 0
+
+
 @pytest.mark.asyncio
 @pytest.mark.vcr()
 async def test_chat_agent_search_tool_with_filter(allow_model_requests, temp_db_path):
@@ -481,6 +537,35 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path)
         assert len(session_state.qa_history) >= 1
 
 
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_chat_agent_ask_with_state_key(allow_model_requests, temp_db_path):
+    """Test ask tool emits keyed state when state_key is set."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+
+        agent = create_chat_agent(Config)
+        session_state = ChatSessionState(session_id="test-ask-keyed")
+        deps = ChatDeps(
+            client=client,
+            config=Config,
+            session_state=session_state,
+            state_key="haiku_rag",
+        )
+
+        result = await agent.run(
+            "What is the highest count class in the DocLayNet dataset?",
+            deps=deps,
+        )
+
+        assert result.output is not None
+        assert len(session_state.qa_history) >= 1
+
+
 def test_fifo_limit_enforcement():
     """Test that FIFO limit enforcement logic works correctly.