diff --git a/app/frontend/components/Chat.tsx b/app/frontend/components/Chat.tsx index c6012fe9..d6212ab3 100644 --- a/app/frontend/components/Chat.tsx +++ b/app/frontend/components/Chat.tsx @@ -411,6 +411,29 @@ function ChatContentInner() { }, ); + const normalizeChatState = ( + state: ChatSessionState | undefined, + ): ChatSessionState => ({ + session_id: state?.session_id ?? "", + initial_context: state?.initial_context ?? null, + citations: state?.citations ?? [], + qa_history: state?.qa_history ?? [], + session_context: state?.session_context ?? null, + document_filter: state?.document_filter ?? [], + citation_registry: state?.citation_registry ?? {}, + }); + + const mergeChatState = (partial: Partial) => { + const current = normalizeChatState(agentState?.[AGUI_STATE_KEY]); + setAgentState({ + ...agentState, + [AGUI_STATE_KEY]: { + ...current, + ...partial, + }, + }); + }; + // Extract session context, document filter, and initial context from agent state const sessionContext = agentState?.[AGUI_STATE_KEY]?.session_context ?? null; const documentFilter = agentState?.[AGUI_STATE_KEY]?.document_filter ?? []; @@ -421,38 +444,12 @@ function ChatContentInner() { (agentState?.[AGUI_STATE_KEY]?.qa_history?.length ?? 0) > 0; const handleFilterApply = (selected: string[]) => { - setAgentState({ - ...agentState, - [AGUI_STATE_KEY]: { - ...agentState?.[AGUI_STATE_KEY], - session_id: agentState?.[AGUI_STATE_KEY]?.session_id ?? "", - initial_context: agentState?.[AGUI_STATE_KEY]?.initial_context ?? null, - citations: agentState?.[AGUI_STATE_KEY]?.citations ?? [], - qa_history: agentState?.[AGUI_STATE_KEY]?.qa_history ?? [], - session_context: agentState?.[AGUI_STATE_KEY]?.session_context ?? null, - document_filter: selected, - citation_registry: - agentState?.[AGUI_STATE_KEY]?.citation_registry ?? {}, - }, - }); + mergeChatState({ document_filter: selected }); }; const handleInitialContextChange = (value: string) => { if (isContextLocked) return; - setAgentState({ - ...agentState, - [AGUI_STATE_KEY]: { - ...agentState?.[AGUI_STATE_KEY], - session_id: agentState?.[AGUI_STATE_KEY]?.session_id ?? "", - initial_context: value || null, - citations: agentState?.[AGUI_STATE_KEY]?.citations ?? [], - qa_history: agentState?.[AGUI_STATE_KEY]?.qa_history ?? [], - session_context: agentState?.[AGUI_STATE_KEY]?.session_context ?? null, - document_filter: agentState?.[AGUI_STATE_KEY]?.document_filter ?? [], - citation_registry: - agentState?.[AGUI_STATE_KEY]?.citation_registry ?? {}, - }, - }); + mergeChatState({ initial_context: value || null }); }; useCoAgentStateRender({ diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index 30ea2dd8..6fbc59c4 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -15,6 +15,7 @@ from haiku.rag.agents.chat.state import ( AGUI_STATE_KEY, ChatSessionState, SessionContext, + build_chat_state_snapshot, ) from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig @@ -56,29 +57,13 @@ class ChatDeps: Combines SessionState and QASessionState into a single state dict matching the ChatSessionState schema expected by AG-UI clients. """ - snapshot: dict[str, Any] = {"session_id": self.session_id} - - # Add SessionState fields session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) - if session_state is not None: - snapshot["document_filter"] = session_state.document_filter - snapshot["citation_registry"] = session_state.citation_registry - snapshot["citations"] = [c.model_dump() for c in session_state.citations] - - # Add QASessionState fields qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) - if qa_session_state is not None: - snapshot["qa_history"] = [ - qa.model_dump() for qa in qa_session_state.qa_history - ] - # Convert string to SessionContext model for frontend - if qa_session_state.session_context: - snapshot["session_context"] = SessionContext( - summary=qa_session_state.session_context - ).model_dump(mode="json") - else: - snapshot["session_context"] = None - + snapshot = build_chat_state_snapshot( + session_state, + qa_session_state, + incoming=False, + ) if self.state_key: return {self.state_key: snapshot} return snapshot @@ -151,7 +136,7 @@ class ChatDeps: qa_session_state.session_context = None # Check cache for fresher session_context from background summarization - # Cache is authoritative - always use it if available + # Cache is authoritative so background summaries show up on next request if self.session_id: cached = get_cached_session_context(self.session_id) if cached and cached.summary: diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index 6c240f4b..84d22c28 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/state.py @@ -1,12 +1,15 @@ from datetime import datetime -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from pydantic import BaseModel from haiku.rag.agents.research.models import Citation if TYPE_CHECKING: - from haiku.rag.tools.qa import QAHistoryEntry + from ag_ui.core import StateDeltaEvent + + from haiku.rag.tools.qa import QAHistoryEntry, QASessionState + from haiku.rag.tools.session import SessionState AGUI_STATE_KEY = "haiku.rag.chat" @@ -38,3 +41,66 @@ def _rebuild_models(qa_history_entry_cls: type) -> None: ChatSessionState.model_rebuild( _types_namespace={"QAHistoryEntry": qa_history_entry_cls} ) + + +def build_chat_state_snapshot( + session_state: "SessionState | None", + qa_state: "QASessionState | None", + *, + incoming: bool = False, +) -> dict[str, Any]: + """Build a combined AG-UI chat state snapshot. + + Args: + session_state: SessionState from ToolContext. + qa_state: QASessionState from ToolContext. + incoming: If True, use client-sent values where applicable. + + Returns: + Snapshot dict, optionally wrapped by state_key. + """ + snapshot: dict[str, Any] = {"session_id": ""} + + if session_state is not None: + snapshot.update( + { + "session_id": ( + session_state.incoming_session_id + if incoming + else session_state.session_id + ), + "document_filter": session_state.document_filter.copy(), + "citation_registry": session_state.citation_registry.copy(), + "citations": [c.model_dump() for c in session_state.citations], + } + ) + + if qa_state is not None: + snapshot["qa_history"] = [qa.model_dump() for qa in qa_state.qa_history] + if incoming: + if qa_state.incoming_session_context is not None: + snapshot["session_context"] = ( + qa_state.incoming_session_context.model_dump(mode="json") + ) + else: + snapshot["session_context"] = None + else: + if qa_state.session_context: + snapshot["session_context"] = SessionContext( + summary=qa_state.session_context + ).model_dump(mode="json") + else: + snapshot["session_context"] = None + + return snapshot + + +def build_chat_state_delta( + old_snapshot: dict[str, Any], + new_snapshot: dict[str, Any], + state_key: str | None, +) -> "StateDeltaEvent | None": + """Compute a delta patch between two combined snapshots.""" + from haiku.rag.tools.session import compute_combined_state_delta + + return compute_combined_state_delta(old_snapshot, new_snapshot, state_key) diff --git a/haiku_rag_slim/haiku/rag/tools/qa.py b/haiku_rag_slim/haiku/rag/tools/qa.py index f2c33215..e51fd08c 100644 --- a/haiku_rag_slim/haiku/rag/tools/qa.py +++ b/haiku_rag_slim/haiku/rag/tools/qa.py @@ -8,7 +8,11 @@ from haiku.rag.agents.chat.context import ( get_cached_embedding, trigger_background_summarization, ) -from haiku.rag.agents.chat.state import SessionContext +from haiku.rag.agents.chat.state import ( + SessionContext, + build_chat_state_delta, + build_chat_state_snapshot, +) from haiku.rag.agents.research.dependencies import ResearchContext from haiku.rag.agents.research.graph import build_research_graph from haiku.rag.agents.research.models import Citation, SearchAnswer @@ -26,7 +30,6 @@ from haiku.rag.tools.models import QAResult from haiku.rag.tools.session import ( SESSION_NAMESPACE, SessionState, - compute_combined_state_delta, ) PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7 @@ -144,25 +147,11 @@ def create_qa_toolset( # Capture combined state snapshot before changes # Use incoming values (what client sent) so delta shows server-side updates if session_state is not None: - old_state_snapshot = { - "session_id": session_state.incoming_session_id, - "document_filter": session_state.document_filter.copy(), - "citation_registry": session_state.citation_registry.copy(), - "citations": [c.model_dump() for c in session_state.citations], - } - if qa_session_state is not None: - old_state_snapshot["qa_history"] = [ - qa.model_dump() for qa in qa_session_state.qa_history - ] - # Use incoming_session_context so delta shows what client sent - if qa_session_state.incoming_session_context is not None: - old_state_snapshot["session_context"] = ( - qa_session_state.incoming_session_context.model_dump( - mode="json" - ) - ) - else: - old_state_snapshot["session_context"] = None + old_state_snapshot = build_chat_state_snapshot( + session_state, + qa_session_state, + incoming=True, + ) # Build filter from session state, base_filter, and document_name doc_filter = build_document_filter(document_name) if document_name else None @@ -298,24 +287,13 @@ def create_qa_toolset( # Compute and return state delta if session state changed if session_state is not None and old_state_snapshot is not None: # Build new combined state snapshot - new_state_snapshot = { - "session_id": session_state.session_id, - "document_filter": session_state.document_filter, - "citation_registry": session_state.citation_registry, - "citations": [c.model_dump() for c in session_state.citations], - } - if qa_session_state is not None: - new_state_snapshot["qa_history"] = [ - qa.model_dump() for qa in qa_session_state.qa_history - ] - if qa_session_state.session_context: - new_state_snapshot["session_context"] = SessionContext( - summary=qa_session_state.session_context - ).model_dump(mode="json") - else: - new_state_snapshot["session_context"] = None + new_state_snapshot = build_chat_state_snapshot( + session_state, + qa_session_state, + incoming=False, + ) - state_event = compute_combined_state_delta( + state_event = build_chat_state_delta( old_state_snapshot, new_state_snapshot, state_key=session_state.state_key,