Merge pull request #237 from ggozad/feat/ag-ui-features

Add AGUI_STATE_KEY for namespaced AG-UI state emission

CHANGELOG.md

@@ -1,6 +1,12 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- **AGUI_STATE_KEY Constant**: Exported `AGUI_STATE_KEY` (`"haiku.rag.chat"`) from `haiku.rag.agents.chat` for namespaced AG-UI state emission
+  - Enables integrators to use a consistent key when combining haiku.rag with other agents
+  - Backend, TUI, and frontend now use this key for state emission and extraction
+
 ## [0.26.3] - 2026-01-15
 
 ### Added

app/backend/main.py

@@ -13,6 +13,7 @@ from starlette.responses import JSONResponse, Response, StreamingResponse
 from starlette.routing import Route
 
 from haiku.rag.agents.chat import (
+    AGUI_STATE_KEY,
     ChatDeps,
     ChatSessionState,
     QAResponse,
@@ -79,11 +80,15 @@ async def stream_chat(request: Request) -> Response:
     accept = request.headers.get("accept", SSE_CONTENT_TYPE)
     run_input = AGUIAdapter.build_run_input(body)
 
-    # Restore qa_history from incoming state
+    # Restore qa_history from incoming state (look under namespaced key)
     initial_qa_history: list[QAResponse] = []
     state = getattr(run_input, "state", None)
-    if state and "qa_history" in state:
-        initial_qa_history = [QAResponse(**qa) for qa in state.get("qa_history", [])]
+    if state:
+        chat_state = state.get(AGUI_STATE_KEY, state)
+        if "qa_history" in chat_state:
+            initial_qa_history = [
+                QAResponse(**qa) for qa in chat_state.get("qa_history", [])
+            ]
 
     # Build deps with session state
     thread_id = getattr(run_input, "thread_id", None)
@@ -94,6 +99,7 @@ async def stream_chat(request: Request) -> Response:
             session_id=thread_id or "",
             qa_history=initial_qa_history,
         ),
+        state_key=AGUI_STATE_KEY,
     )
 
     # Use AGUIAdapter for streaming

app/frontend/components/Chat.tsx

@@ -11,6 +11,9 @@ import "@copilotkit/react-ui/styles.css";
 import CitationBlock from "./CitationBlock";
 import DbInfo from "./DbInfo";
 
+// Must match AGUI_STATE_KEY from haiku.rag.agents.chat
+const AGUI_STATE_KEY = "haiku.rag.chat";
+
 interface Citation {
 	index: number;
 	document_id: string;
@@ -35,6 +38,11 @@ interface ChatSessionState {
 	qa_history: QAResponse[];
 }
 
+// AG-UI state is namespaced under AGUI_STATE_KEY
+interface AgentState {
+	[AGUI_STATE_KEY]?: ChatSessionState;
+}
+
 function SpinnerIcon() {
 	return (
 		<svg
@@ -329,20 +337,23 @@ function ToolCallIndicator({
 }
 
 function ChatContent() {
-	useCoAgent<ChatSessionState>({
+	useCoAgent<AgentState>({
 		name: "chat_agent",
 		initialState: {
-			session_id: "",
-			citations: [],
-			qa_history: [],
+			[AGUI_STATE_KEY]: {
+				session_id: "",
+				citations: [],
+				qa_history: [],
+			},
 		},
 	});
 
-	useCoAgentStateRender<ChatSessionState>({
+	useCoAgentStateRender<AgentState>({
 		name: "chat_agent",
 		render: ({ state }) => {
-			if (state.citations && state.citations.length > 0) {
-				return <CitationBlock citations={state.citations} />;
+			const chatState = state[AGUI_STATE_KEY];
+			if (chatState?.citations.length) {
+				return <CitationBlock citations={chatState.citations} />;
 			}
 			return null;
 		},

docs/agents.md

@@ -111,6 +111,37 @@ Q/A history is used to:
 2. Avoid repeating previous answers
 3. Enable semantic ranking of relevant past answers
 
+### AG-UI Integration
+
+When using the chat agent with AG-UI streaming, state is emitted under a namespaced key to avoid conflicts with other agents:
+
+```python
+from haiku.rag.agents.chat import AGUI_STATE_KEY, ChatDeps, ChatSessionState
+
+# AGUI_STATE_KEY = "haiku.rag.chat"
+
+deps = ChatDeps(
+    client=client,
+    config=config,
+    session_state=ChatSessionState(),
+    state_key=AGUI_STATE_KEY,  # Enables namespaced state emission
+)
+```
+
+The emitted state structure:
+
+```json
+{
+  "haiku.rag.chat": {
+    "session_id": "",
+    "citations": [...],
+    "qa_history": [...]
+  }
+}
+```
+
+Frontend clients should extract state from under this key. See the [Conversational RAG App](apps.md#conversational-rag-app) for a complete implementation example.
+
 ## Research Graph
 
 The research workflow is implemented as a typed pydantic-graph. It plans, searches (in parallel batches), evaluates, and synthesizes into a final report.

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

@@ -1,6 +1,7 @@
 from haiku.rag.agents.chat.agent import create_chat_agent
 from haiku.rag.agents.chat.search import SearchAgent
 from haiku.rag.agents.chat.state import (
+    AGUI_STATE_KEY,
     ChatDeps,
     ChatSessionState,
     CitationInfo,
@@ -11,6 +12,7 @@ from haiku.rag.agents.chat.state import (
 )
 
 __all__ = [
+    "AGUI_STATE_KEY",
     "create_chat_agent",
     "SearchAgent",
     "ChatDeps",

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

@@ -16,6 +16,8 @@ if TYPE_CHECKING:
 
 MAX_QA_HISTORY = 50
 
+AGUI_STATE_KEY = "haiku.rag.chat"
+
 _embedding_cache: dict[str, list[float]] = {}
 
 

haiku_rag_slim/haiku/rag/chat/app.py

@@ -17,7 +17,12 @@ from pydantic_ai import (
 from pydantic_ai.messages import ModelMessage
 
 from haiku.rag.agents.chat.agent import create_chat_agent
-from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState, CitationInfo
+from haiku.rag.agents.chat.state import (
+    AGUI_STATE_KEY,
+    ChatDeps,
+    ChatSessionState,
+    CitationInfo,
+)
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import get_config
 
@@ -153,10 +158,10 @@ class ChatApp(App):  # type: ignore[misc]
                         and meta_event.type == EventType.STATE_SNAPSHOT
                     ):
                         snapshot = getattr(meta_event, "snapshot", {})
-                        if "citations" in snapshot:
-                            self._last_citations = [
-                                CitationInfo(**c) for c in snapshot["citations"]
-                            ]
+                        chat_state = snapshot.get(AGUI_STATE_KEY, snapshot)
+                        self._last_citations = [
+                            CitationInfo(**c) for c in chat_state["citations"]
+                        ]
 
     async def _event_stream_handler(
         self,
@@ -212,6 +217,7 @@ class ChatApp(App):  # type: ignore[misc]
                 client=self.client,
                 config=self.config,
                 session_state=self.session_state,
+                state_key=AGUI_STATE_KEY,
             )
 
             async with self.agent.run_stream(

tests/agents/chat/test_chat_agent.py

@@ -3,6 +3,7 @@ from pathlib import Path
 import pytest
 
 from haiku.rag.agents.chat import (
+    AGUI_STATE_KEY,
     ChatDeps,
     ChatSessionState,
     CitationInfo,
@@ -40,6 +41,11 @@ def test_chat_deps_initialization(temp_db_path):
     client.close()
 
 
+def test_agui_state_key_constant():
+    """Test AGUI_STATE_KEY is exported with correct value."""
+    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)
@@ -309,7 +315,7 @@ async def test_chat_agent_search_with_state_key(allow_model_requests, temp_db_pa
             client=client,
             config=Config,
             session_state=session_state,
-            state_key="haiku_rag",
+            state_key=AGUI_STATE_KEY,
         )
 
         result = await agent.run(
@@ -554,7 +560,7 @@ async def test_chat_agent_ask_with_state_key(allow_model_requests, temp_db_path)
             client=client,
             config=Config,
             session_state=session_state,
-            state_key="haiku_rag",
+            state_key=AGUI_STATE_KEY,
         )
 
         result = await agent.run(