diff --git a/CHANGELOG.md b/CHANGELOG.md index b400759a..5ac01892 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/app/backend/main.py b/app/backend/main.py index 6695be55..42a4fa5d 100644 --- a/app/backend/main.py +++ b/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 diff --git a/app/frontend/components/Chat.tsx b/app/frontend/components/Chat.tsx index c13dde5d..dafca4a9 100644 --- a/app/frontend/components/Chat.tsx +++ b/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 ( ({ + useCoAgent({ name: "chat_agent", initialState: { - session_id: "", - citations: [], - qa_history: [], + [AGUI_STATE_KEY]: { + session_id: "", + citations: [], + qa_history: [], + }, }, }); - useCoAgentStateRender({ + useCoAgentStateRender({ name: "chat_agent", render: ({ state }) => { - if (state.citations && state.citations.length > 0) { - return ; + const chatState = state[AGUI_STATE_KEY]; + if (chatState?.citations.length) { + return ; } return null; }, diff --git a/docs/agents.md b/docs/agents.md index 0785f226..70c436f1 100644 --- a/docs/agents.md +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/agents/chat/__init__.py b/haiku_rag_slim/haiku/rag/agents/chat/__init__.py index 8947cca9..b5cea851 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/__init__.py +++ b/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", diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index 44d2c4e8..7708df89 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/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]] = {} diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 7536ac96..94b4f68a 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/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( diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index 323cee6f..7bb2ac43 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/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(