danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Extract shared snapshot helpers

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-10 17:34:30 +02:00
parent a592a04031
commit f85aaefed7
No known key found for this signature in database
4 changed files with 116 additions and 90 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
app/frontend/components/Chat.tsx
View file

@ -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<ChatSessionState>) => {
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 // Extract session context, document filter, and initial context from agent state
const sessionContext = agentState?.[AGUI_STATE_KEY]?.session_context ?? null; const sessionContext = agentState?.[AGUI_STATE_KEY]?.session_context ?? null;
const documentFilter = agentState?.[AGUI_STATE_KEY]?.document_filter ?? []; const documentFilter = agentState?.[AGUI_STATE_KEY]?.document_filter ?? [];
@ -421,38 +444,12 @@ function ChatContentInner() {
(agentState?.[AGUI_STATE_KEY]?.qa_history?.length ?? 0) > 0; (agentState?.[AGUI_STATE_KEY]?.qa_history?.length ?? 0) > 0;
const handleFilterApply = (selected: string[]) => { const handleFilterApply = (selected: string[]) => {
setAgentState({ mergeChatState({ document_filter: selected });
...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 ?? {},
},
});
}; };
const handleInitialContextChange = (value: string) => { const handleInitialContextChange = (value: string) => {
if (isContextLocked) return; if (isContextLocked) return;
setAgentState({ mergeChatState({ initial_context: value || null });
...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 ?? {},
},
});
}; };
useCoAgentStateRender<AgentState>({ useCoAgentStateRender<AgentState>({

29
haiku_rag_slim/haiku/rag/agents/chat/agent.py
View file

@ -15,6 +15,7 @@ from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY, AGUI_STATE_KEY,
ChatSessionState, ChatSessionState,
SessionContext, SessionContext,
build_chat_state_snapshot,
) )
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
@ -56,29 +57,13 @@ class ChatDeps:
Combines SessionState and QASessionState into a single state dict Combines SessionState and QASessionState into a single state dict
matching the ChatSessionState schema expected by AG-UI clients. 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) 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) qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState)
if qa_session_state is not None: snapshot = build_chat_state_snapshot(
snapshot["qa_history"] = [ session_state,
qa.model_dump() for qa in qa_session_state.qa_history qa_session_state,
] incoming=False,
# 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
if self.state_key: if self.state_key:
return {self.state_key: snapshot} return {self.state_key: snapshot}
return snapshot return snapshot
@ -151,7 +136,7 @@ class ChatDeps:
qa_session_state.session_context = None qa_session_state.session_context = None
# Check cache for fresher session_context from background summarization # 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: if self.session_id:
cached = get_cached_session_context(self.session_id) cached = get_cached_session_context(self.session_id)
if cached and cached.summary: if cached and cached.summary:

70
haiku_rag_slim/haiku/rag/agents/chat/state.py
View file

@ -1,12 +1,15 @@
from datetime import datetime from datetime import datetime
from typing import TYPE_CHECKING from typing import TYPE_CHECKING, Any
from pydantic import BaseModel from pydantic import BaseModel
from haiku.rag.agents.research.models import Citation from haiku.rag.agents.research.models import Citation
if TYPE_CHECKING: 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" AGUI_STATE_KEY = "haiku.rag.chat"
@ -38,3 +41,66 @@ def _rebuild_models(qa_history_entry_cls: type) -> None:
ChatSessionState.model_rebuild( ChatSessionState.model_rebuild(
_types_namespace={"QAHistoryEntry": qa_history_entry_cls} _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)

54
haiku_rag_slim/haiku/rag/tools/qa.py
View file

@ -8,7 +8,11 @@ from haiku.rag.agents.chat.context import (
get_cached_embedding, get_cached_embedding,
trigger_background_summarization, 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.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.models import Citation, SearchAnswer 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 ( from haiku.rag.tools.session import (
SESSION_NAMESPACE, SESSION_NAMESPACE,
SessionState, SessionState,
compute_combined_state_delta,
) )
PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7 PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7
@ -144,25 +147,11 @@ def create_qa_toolset(
# Capture combined state snapshot before changes # Capture combined state snapshot before changes
# Use incoming values (what client sent) so delta shows server-side updates # Use incoming values (what client sent) so delta shows server-side updates
if session_state is not None: if session_state is not None:
old_state_snapshot = { old_state_snapshot = build_chat_state_snapshot(
"session_id": session_state.incoming_session_id, session_state,
"document_filter": session_state.document_filter.copy(), qa_session_state,
"citation_registry": session_state.citation_registry.copy(), incoming=True,
"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
# Build filter from session state, base_filter, and document_name # Build filter from session state, base_filter, and document_name
doc_filter = build_document_filter(document_name) if document_name else None 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 # Compute and return state delta if session state changed
if session_state is not None and old_state_snapshot is not None: if session_state is not None and old_state_snapshot is not None:
# Build new combined state snapshot # Build new combined state snapshot
new_state_snapshot = { new_state_snapshot = build_chat_state_snapshot(
"session_id": session_state.session_id, session_state,
"document_filter": session_state.document_filter, qa_session_state,
"citation_registry": session_state.citation_registry, incoming=False,
"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
state_event = compute_combined_state_delta( state_event = build_chat_state_delta(
old_state_snapshot, old_state_snapshot,
new_state_snapshot, new_state_snapshot,
state_key=session_state.state_key, state_key=session_state.state_key,