diff --git a/app/backend/main.py b/app/backend/main.py index 3dabf862..93ddf1f1 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -80,30 +80,30 @@ async def stream_chat(request: Request) -> Response: run_input = AGUIAdapter.build_run_input(body) # Restore session state from incoming AG-UI state (look under namespaced key) - initial_qa_history: list[QAResponse] = [] - session_id: str | None = None - document_filter: list[str] = [] - initial_context: str | None = None + session_state: ChatSessionState | None = None state = getattr(run_input, "state", None) - 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", []) - ] - session_id = chat_state.get("session_id") - document_filter = chat_state.get("document_filter", []) - initial_context = chat_state.get("initial_context") + if state and AGUI_STATE_KEY in state: + chat_state = state[AGUI_STATE_KEY] + if chat_state and chat_state.get("session_id"): + # Only restore state if client has a session_id (not first request) + # This ensures first request gets a full snapshot with generated UUID + # NOTE: We intentionally do NOT restore session_context from the client. + # The server maintains session_context via background summarization tasks, + # and the agent fetches it from the server-side cache (get_cached_session_context). + session_state = ChatSessionState( + session_id=chat_state["session_id"], + qa_history=[ + QAResponse(**qa) for qa in chat_state.get("qa_history", []) + ], + document_filter=chat_state.get("document_filter", []), + initial_context=chat_state.get("initial_context"), + citation_registry=chat_state.get("citation_registry", {}), + ) deps = ChatDeps( client=get_client(db_path), config=Config, - session_state=ChatSessionState( - qa_history=initial_qa_history, - document_filter=document_filter, - initial_context=initial_context, - **({"session_id": session_id} if session_id else {}), - ), + session_state=session_state, state_key=AGUI_STATE_KEY, ) @@ -111,25 +111,7 @@ async def stream_chat(request: Request) -> Response: adapter = AGUIAdapter(agent=chat_agent, run_input=run_input, accept=accept) event_stream = adapter.run_stream(deps=deps) - # Wrap to log state events - async def logged_event_stream(): - async for event in event_stream: - event_type = getattr(event, "type", None) - if event_type and "state" in str(event_type).lower(): - delta: list[dict[str, str]] | None = getattr(event, "delta", None) - snapshot: dict[str, object] | None = getattr(event, "snapshot", None) - if delta is not None: - logger.info(f"StateDeltaEvent: {len(delta)} ops") - for op in delta[:3]: # Log first 3 ops - logger.info(f" {op['op']} {op['path']}") - if len(delta) > 3: - logger.info(f" ... and {len(delta) - 3} more ops") - elif snapshot is not None: - snapshot_keys = list(snapshot.keys()) - logger.info(f"StateSnapshotEvent: keys={snapshot_keys}") - yield event - - sse_event_stream = adapter.encode_stream(logged_event_stream()) + sse_event_stream = adapter.encode_stream(event_stream) return StreamingResponse( sse_event_stream, diff --git a/app/frontend/components/Chat.tsx b/app/frontend/components/Chat.tsx index 4a0e11d3..c6012fe9 100644 --- a/app/frontend/components/Chat.tsx +++ b/app/frontend/components/Chat.tsx @@ -47,6 +47,7 @@ interface ChatSessionState { qa_history: QAResponse[]; session_context: SessionContext | null; document_filter: string[]; + citation_registry: Record; } // AG-UI state is namespaced under AGUI_STATE_KEY @@ -404,6 +405,7 @@ function ChatContentInner() { qa_history: [], session_context: null, document_filter: [], + citation_registry: {}, }, }, }, @@ -429,6 +431,8 @@ function ChatContentInner() { 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 ?? {}, }, }); }; @@ -445,6 +449,8 @@ function ChatContentInner() { 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 ?? {}, }, }); }; diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index c77d702f..3a07bb6b 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -116,12 +116,19 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: if not results: return ToolReturn(return_value="No results found.") - # Build citation infos using stable registry indices + # Copy session state to work with (avoids mutating original for delta computation) + new_state = ( + ctx.deps.session_state.model_copy(deep=True) + if ctx.deps.session_state + else ChatSessionState() + ) + + # Build citation infos using the copy's registry citation_infos = [] for r in results: chunk_id = r.chunk_id or "" - if ctx.deps.session_state is not None and chunk_id: - index = ctx.deps.session_state.get_or_assign_index(chunk_id) + if chunk_id: + index = new_state.get_or_assign_index(chunk_id) else: index = len(citation_infos) + 1 citation_infos.append( @@ -137,26 +144,10 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: ) ) - # Build new state with citations and registry - session_id = ctx.deps.session_state.session_id if ctx.deps.session_state else "" - new_state = ChatSessionState( - session_id=session_id, - citations=citation_infos, - qa_history=( - ctx.deps.session_state.qa_history if ctx.deps.session_state else [] - ), - session_context=get_cached_session_context(session_id) - if session_id - else None, - document_filter=( - ctx.deps.session_state.document_filter if ctx.deps.session_state else [] - ), - citation_registry=( - ctx.deps.session_state.citation_registry - if ctx.deps.session_state - else {} - ), - ) + # Update new_state with citations and fresh session_context + new_state.citations = citation_infos + if new_state.session_id: + new_state.session_context = get_cached_session_context(new_state.session_id) # Return detailed results for the agent to present result_lines = [] @@ -276,14 +267,17 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: result = await graph.run(state=state, deps=deps) - # Build citation infos using stable registry indices + # Copy session state to work with (avoids mutating original for delta computation) + new_state = ( + ctx.deps.session_state.model_copy(deep=True) + if ctx.deps.session_state + else ChatSessionState() + ) + + # Build citation infos using the copy's registry citation_infos = [] for c in result.citations: - # Use registry for stable indices across calls - if ctx.deps.session_state is not None: - index = ctx.deps.session_state.get_or_assign_index(c.chunk_id) - else: - index = len(citation_infos) + 1 + index = new_state.get_or_assign_index(c.chunk_id) citation_infos.append( Citation( index=index, @@ -297,54 +291,37 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: ) ) - # Accumulate Q&A in session state with full citation metadata - if ctx.deps.session_state is not None: - qa_response = QAResponse( - question=question, - answer=result.answer, - confidence=result.confidence, - citations=citation_infos, - ) - ctx.deps.session_state.qa_history.append(qa_response) - # Enforce FIFO limit - if len(ctx.deps.session_state.qa_history) > MAX_QA_HISTORY: - ctx.deps.session_state.qa_history = ctx.deps.session_state.qa_history[ - -MAX_QA_HISTORY: - ] - - # Spawn background task to update session context - # Cancel any previous summarization for this session - if session_id in _summarization_tasks: - _summarization_tasks[session_id].cancel() - - task = asyncio.create_task( - _update_context_background( - qa_history=list(ctx.deps.session_state.qa_history), - config=ctx.deps.config, - session_state=ctx.deps.session_state, - ) - ) - _summarization_tasks[session_id] = task - task.add_done_callback(lambda t: _summarization_tasks.pop(session_id, None)) - - # Build new state with citations, qa_history, and registry - new_state = ChatSessionState( - session_id=session_id, + # Add Q&A to the copy's history + qa_response = QAResponse( + question=question, + answer=result.answer, + confidence=result.confidence, citations=citation_infos, - qa_history=( - ctx.deps.session_state.qa_history if ctx.deps.session_state else [] - ), - session_context=get_cached_session_context(session_id) - if session_id - else None, - document_filter=( - ctx.deps.session_state.document_filter if ctx.deps.session_state else [] - ), - citation_registry=( - ctx.deps.session_state.citation_registry - if ctx.deps.session_state - else {} - ), + ) + new_state.qa_history.append(qa_response) + # Enforce FIFO limit + if len(new_state.qa_history) > MAX_QA_HISTORY: + new_state.qa_history = new_state.qa_history[-MAX_QA_HISTORY:] + + # Update citations and session_context + new_state.citations = citation_infos + if new_state.session_id: + new_state.session_context = get_cached_session_context(new_state.session_id) + + # Spawn background task to update session context + if new_state.session_id in _summarization_tasks: + _summarization_tasks[new_state.session_id].cancel() + + task = asyncio.create_task( + _update_context_background( + qa_history=list(new_state.qa_history), + config=ctx.deps.config, + session_state=new_state, + ) + ) + _summarization_tasks[new_state.session_id] = task + task.add_done_callback( + lambda t, sid=new_state.session_id: _summarization_tasks.pop(sid, None) ) # Format answer with citation references using stable indices