diff --git a/app/backend/main.py b/app/backend/main.py index 5c998014..f527ec10 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -79,13 +79,12 @@ async def stream_chat(request: Request) -> Response: run_input = AGUIAdapter.build_run_input(body) thread_id = getattr(run_input, "thread_id", None) or "default" - context, is_new = context_cache.get_or_create(thread_id) + context, _is_new = context_cache.get_or_create(thread_id) agent = create_chat_agent(Config, get_client(), context) deps = ChatDeps( config=Config, tool_context=context, - is_new=is_new, state_key=AGUI_STATE_KEY, ) diff --git a/docs/agents/index.md b/docs/agents/index.md index 4fa00b3f..e8d00caf 100644 --- a/docs/agents/index.md +++ b/docs/agents/index.md @@ -127,13 +127,27 @@ The system prompt is automatically composed to match the selected features. See ### Session State -The `ChatSessionState` maintains: +Session state is managed through `ToolContext` — a namespace-based state container shared across all toolsets. The chat agent uses two namespaces: + +**`SessionState`** (session management): -- `session_id` — Unique identifier for the session -- `qa_history` — List of previous Q/A pairs -- `session_context` — Automatically maintained session context summary - `document_filter` — List of document titles/URIs to restrict searches - `citation_registry` — Stable mapping of chunk IDs to citation indices +- `citations` — Citations from the current query + +**`QASessionState`** (QA history and context): + +- `qa_history` — List of previous Q/A pairs with embeddings +- `session_context` — Automatically maintained session context summary + +For multi-session applications (e.g., web backends), use `ToolContextCache` to cache `ToolContext` instances by external session/thread ID: + +```python +from haiku.rag.tools import ToolContext, ToolContextCache + +cache = ToolContextCache() # TTL-based, defaults to 1 hour +context, _is_new = cache.get_or_create(thread_id) +``` **Citation Registry**: Citation indices persist across tool calls within a session. The same `chunk_id` always returns the same citation index (first-occurrence-wins). This ensures consistent citation numbering in multi-turn conversations — `[1]` always refers to the same source. diff --git a/docs/tools.md b/docs/tools.md index fc6d0ec8..a12fb550 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -227,9 +227,12 @@ When using the chat agent with [AG-UI](https://docs.ag-ui.com) streaming, `ChatD ```python from haiku.rag.agents.chat import AGUI_STATE_KEY, ChatDeps, create_chat_agent -from haiku.rag.tools import ToolContext +from haiku.rag.tools import ToolContext, ToolContextCache + +# For multi-session apps, cache ToolContext per thread +cache = ToolContextCache() +context, _is_new = cache.get_or_create(thread_id) -context = ToolContext() agent = create_chat_agent(config, client, context) deps = ChatDeps( config=config, @@ -243,7 +246,6 @@ The emitted state structure: ```json { "haiku.rag.chat": { - "session_id": "uuid", "citations": [], "qa_history": [], "session_context": null, @@ -253,7 +255,7 @@ The emitted state structure: } ``` -State flows bidirectionally — the frontend sends its current state on each request, and the agent emits deltas (JSON Patch) reflecting server-side updates (new citations, QA history entries, session context). See the [Web Application](apps.md#web-application) for a complete implementation. +State flows bidirectionally — the frontend sends its current state on each request, and the agent emits deltas (JSON Patch) reflecting server-side updates (new citations, QA history entries, session context). The server always prefers its own `session_context` over the client's value, since background summarization may have updated it between requests. See the [Web Application](apps.md#web-application) for a complete implementation. ## Filter Helpers diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index f429eb1d..45606318 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -43,7 +43,6 @@ class ChatDeps: config: AppConfig tool_context: ToolContext - is_new: bool = True state_key: str | None = None @property @@ -74,57 +73,43 @@ class ChatDeps: state_data = nested session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) + if session_state is not None: + if "document_filter" in state_data: + session_state.document_filter = state_data.get("document_filter", []) + if "citation_registry" in state_data: + session_state.citation_registry = state_data["citation_registry"] + if "citations" in state_data: + from haiku.rag.agents.research.models import Citation - if self.is_new: - # First request for this context: fully populate from client state - if session_state is not None: - if "document_filter" in state_data: - session_state.document_filter = state_data.get( - "document_filter", [] - ) - if "citation_registry" in state_data: - session_state.citation_registry = state_data["citation_registry"] - if "citations" in state_data: - from haiku.rag.agents.research.models import Citation + session_state.citations = [ + Citation(**c) if isinstance(c, dict) else c + for c in state_data.get("citations", []) + ] - session_state.citations = [ - Citation(**c) if isinstance(c, dict) else c - for c in state_data.get("citations", []) - ] + qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) + if qa_session_state is not None: + if "qa_history" in state_data: + from haiku.rag.tools.qa import QAHistoryEntry - qa_session_state = self.tool_context.get( - QA_SESSION_NAMESPACE, QASessionState - ) - if qa_session_state is not None: - if "qa_history" in state_data: - from haiku.rag.tools.qa import QAHistoryEntry + qa_session_state.qa_history = [ + QAHistoryEntry(**qa) if isinstance(qa, dict) else qa + for qa in state_data.get("qa_history", []) + ] - qa_session_state.qa_history = [ - QAHistoryEntry(**qa) if isinstance(qa, dict) else qa - for qa in state_data.get("qa_history", []) - ] - - # Restore session_context from client + # Prefer server's session_context (background summarizer may + # have updated it since the client's last snapshot). + if not qa_session_state.session_context: session_context = state_data.get("session_context") if isinstance(session_context, dict): qa_session_state.session_context = SessionContext( **session_context ).summary - elif session_context is None: - qa_session_state.session_context = None # Handle initial_context -> session_context for first message if "initial_context" in state_data: initial = state_data.get("initial_context") if initial and not qa_session_state.session_context: qa_session_state.session_context = initial - else: - # Returning request: only merge client-controlled fields - if session_state is not None: - if "document_filter" in state_data: - session_state.document_filter = state_data.get( - "document_filter", [] - ) def create_chat_agent( diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index aae5f7e0..9b49c12a 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -251,7 +251,6 @@ class ChatApp(App): deps = ChatDeps( config=self.config, tool_context=self.tool_context, - is_new=False, ) async with self.agent.run_stream( diff --git a/haiku_rag_slim/haiku/rag/tools/__init__.py b/haiku_rag_slim/haiku/rag/tools/__init__.py index aa2a80cb..58e24fe4 100644 --- a/haiku_rag_slim/haiku/rag/tools/__init__.py +++ b/haiku_rag_slim/haiku/rag/tools/__init__.py @@ -48,6 +48,7 @@ __all__ = [ "QASessionState", "QAHistoryEntry", "create_qa_toolset", + "run_qa_core", "create_analysis_toolset", "SESSION_NAMESPACE", "SessionState", diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index 1446f633..b7c4d644 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/tests/agents/chat/test_chat_agent.py @@ -64,7 +64,6 @@ def test_chat_deps_initialization(temp_db_path): assert deps.config is Config assert deps.tool_context is context - assert deps.is_new is True assert deps.state_key is None @@ -153,6 +152,42 @@ def test_chat_deps_state_setter_parses_session_context_dict(): assert qa_session_state.session_context == "Previous conversation summary" +def test_chat_deps_state_setter_preserves_server_session_context(): + """Test that server's session_context is preferred over client's stale value.""" + from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState + + context = ToolContext() + qa_state = QASessionState() + qa_state.session_context = "Fresh summary from background summarizer" + context.register(QA_SESSION_NAMESPACE, qa_state) + context.register(SESSION_NAMESPACE, SessionState()) + + deps = ChatDeps(config=Config, tool_context=context, state_key=AGUI_STATE_KEY) + + # Client sends stale session_context + incoming_state = { + AGUI_STATE_KEY: { + "session_context": { + "summary": "Stale summary from client", + "last_updated": "2025-01-27T12:00:00", + }, + "qa_history": [], + "citations": [], + "document_filter": [], + "citation_registry": {}, + } + } + + deps.state = incoming_state + + # Server's fresher session_context should be preserved + qa_session_state = context.get(QA_SESSION_NAMESPACE) + assert isinstance(qa_session_state, QASessionState) + assert ( + qa_session_state.session_context == "Fresh summary from background summarizer" + ) + + def test_chat_session_state(): """Test ChatSessionState model.""" state = ChatSessionState()