diff --git a/app/backend/main.py b/app/backend/main.py index 506f15d3..3dabf862 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -110,7 +110,26 @@ async def stream_chat(request: Request) -> Response: # Use AGUIAdapter for streaming adapter = AGUIAdapter(agent=chat_agent, run_input=run_input, accept=accept) event_stream = adapter.run_stream(deps=deps) - sse_event_stream = adapter.encode_stream(event_stream) + + # 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()) return StreamingResponse( sse_event_stream, diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index 85134763..c77d702f 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -1,7 +1,6 @@ import asyncio import math -from ag_ui.core import EventType, StateSnapshotEvent from pydantic_ai import Agent, RunContext, ToolReturn from haiku.rag.agents.chat.context import ( @@ -20,6 +19,7 @@ from haiku.rag.agents.chat.state import ( build_document_filter, build_multi_document_filter, combine_filters, + emit_state_event, ) from haiku.rag.agents.research.dependencies import ResearchContext from haiku.rag.agents.research.graph import build_conversational_graph @@ -173,19 +173,14 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: line += f"\n {snippet}" result_lines.append(line) - snapshot = new_state.model_dump(mode="json") - if ctx.deps.state_key: - snapshot = {ctx.deps.state_key: snapshot} + state_event = emit_state_event( + ctx.deps.session_state, new_state, ctx.deps.state_key + ) return ToolReturn( return_value=f"Found {len(results)} results:\n\n" + "\n\n".join(result_lines), - metadata=[ - StateSnapshotEvent( - type=EventType.STATE_SNAPSHOT, - snapshot=snapshot, - ) - ], + metadata=[state_event] if state_event else None, ) @agent.tool @@ -358,18 +353,13 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: citation_refs = " ".join(f"[{c.index}]" for c in citation_infos) answer_text = f"{answer_text}\n\nSources: {citation_refs}" - snapshot = new_state.model_dump(mode="json") - if ctx.deps.state_key: - snapshot = {ctx.deps.state_key: snapshot} + state_event = emit_state_event( + ctx.deps.session_state, new_state, ctx.deps.state_key + ) return ToolReturn( return_value=answer_text, - metadata=[ - StateSnapshotEvent( - type=EventType.STATE_SNAPSHOT, - snapshot=snapshot, - ) - ], + metadata=[state_event] if state_event else None, ) @agent.tool diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index 5e7ccfed..56114b3f 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/state.py @@ -3,6 +3,8 @@ from dataclasses import dataclass, field from datetime import datetime from typing import Any +import jsonpatch +from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent from pydantic import BaseModel, Field from haiku.rag.agents.research.models import Citation, SearchAnswer @@ -194,3 +196,32 @@ def combine_filters(filter1: str | None, filter2: str | None) -> str | None: if len(filters) == 1: return filters[0] return f"({filters[0]}) AND ({filters[1]})" + + +def emit_state_event( + current_state: ChatSessionState | None, + new_state: ChatSessionState, + state_key: str | None = None, +) -> StateSnapshotEvent | StateDeltaEvent | None: + """Emit state delta against current state, or full snapshot if no current state.""" + new_snapshot = new_state.model_dump(mode="json") + wrapped_new = {state_key: new_snapshot} if state_key else new_snapshot + + if current_state is None: + return StateSnapshotEvent( + type=EventType.STATE_SNAPSHOT, + snapshot=wrapped_new, + ) + + current_snapshot = current_state.model_dump(mode="json") + wrapped_current = {state_key: current_snapshot} if state_key else current_snapshot + + patch = jsonpatch.make_patch(wrapped_current, wrapped_new) + + if not patch.patch: + return None + + return StateDeltaEvent( + type=EventType.STATE_DELTA, + delta=patch.patch, + ) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 1158f73e..27c69475 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -24,6 +24,7 @@ classifiers = [ dependencies = [ "docling-core==2.60.1", "httpx>=0.28.1", + "jsonpatch>=1.33", "lancedb==0.27.0", "pathspec>=1.0.3", "pydantic>=2.12.5", diff --git a/tests/agents/chat/test_state.py b/tests/agents/chat/test_state.py index 2ed78e7a..f3ab9127 100644 --- a/tests/agents/chat/test_state.py +++ b/tests/agents/chat/test_state.py @@ -641,3 +641,118 @@ def test_chat_session_state_model_dump_json_serializes_datetime(): # datetime should be serialized as ISO string, not datetime object assert isinstance(snapshot["session_context"]["last_updated"], str) assert snapshot["session_context"]["last_updated"] == "2025-01-27T12:00:00" + + +def test_emit_state_event_returns_snapshot_when_no_current_state(): + """emit_state_event returns StateSnapshotEvent when current_state is None.""" + from ag_ui.core import EventType, StateSnapshotEvent + + from haiku.rag.agents.chat.state import emit_state_event + + new_state = ChatSessionState(session_id="test-123") + + event = emit_state_event(None, new_state) + + assert isinstance(event, StateSnapshotEvent) + assert event.type == EventType.STATE_SNAPSHOT + assert event.snapshot["session_id"] == "test-123" + + +def test_emit_state_event_returns_snapshot_with_state_key(): + """emit_state_event wraps snapshot in state_key namespace.""" + from ag_ui.core import EventType, StateSnapshotEvent + + from haiku.rag.agents.chat.state import AGUI_STATE_KEY, emit_state_event + + new_state = ChatSessionState(session_id="test-123") + + event = emit_state_event(None, new_state, state_key=AGUI_STATE_KEY) + + assert isinstance(event, StateSnapshotEvent) + assert event.type == EventType.STATE_SNAPSHOT + assert AGUI_STATE_KEY in event.snapshot + assert event.snapshot[AGUI_STATE_KEY]["session_id"] == "test-123" + + +def test_emit_state_event_returns_none_when_no_changes(): + """emit_state_event returns None when states are identical.""" + from haiku.rag.agents.chat.state import emit_state_event + + state = ChatSessionState(session_id="test-123", qa_history=[], citations=[]) + + event = emit_state_event(state, state) + + assert event is None + + +def test_emit_state_event_returns_delta_with_changes(): + """emit_state_event returns StateDeltaEvent with JSON Patch ops for changes.""" + from ag_ui.core import EventType, StateDeltaEvent + + from haiku.rag.agents.chat.state import emit_state_event + + current_state = ChatSessionState(session_id="test-123", qa_history=[], citations=[]) + new_state = ChatSessionState( + session_id="test-123", + qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)], + citations=[], + ) + + event = emit_state_event(current_state, new_state) + + assert isinstance(event, StateDeltaEvent) + assert event.type == EventType.STATE_DELTA + assert len(event.delta) > 0 + # Delta should contain an "add" operation for the new qa_history entry + ops = event.delta + qa_history_op = next((op for op in ops if "/qa_history" in op["path"]), None) + assert qa_history_op is not None + + +def test_emit_state_event_delta_with_state_key(): + """emit_state_event wraps delta paths with state_key namespace.""" + from ag_ui.core import StateDeltaEvent + + from haiku.rag.agents.chat.state import AGUI_STATE_KEY, emit_state_event + + current_state = ChatSessionState(session_id="test-123", qa_history=[]) + new_state = ChatSessionState( + session_id="test-123", + qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)], + ) + + event = emit_state_event(current_state, new_state, state_key=AGUI_STATE_KEY) + + assert isinstance(event, StateDeltaEvent) + # Paths should be namespaced under state_key + for op in event.delta: + assert op["path"].startswith(f"/{AGUI_STATE_KEY}") + + +def test_emit_state_event_delta_produces_valid_patch(): + """emit_state_event delta can be applied to reproduce new state.""" + import jsonpatch + + from haiku.rag.agents.chat.state import emit_state_event + + current_state = ChatSessionState( + session_id="test-123", + qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)], + citations=[], + ) + new_state = ChatSessionState( + session_id="test-123", + qa_history=[ + QAResponse(question="Q1", answer="A1", confidence=0.9), + QAResponse(question="Q2", answer="A2", confidence=0.8), + ], + citations=[], + ) + + event = emit_state_event(current_state, new_state) + + # Apply patch to current state and verify it produces new state + current_snapshot = current_state.model_dump(mode="json") + patched = jsonpatch.apply_patch(current_snapshot, event.delta) + new_snapshot = new_state.model_dump(mode="json") + assert patched == new_snapshot diff --git a/uv.lock b/uv.lock index f430bc01..943fab80 100644 --- a/uv.lock +++ b/uv.lock @@ -1368,6 +1368,7 @@ source = { editable = "haiku_rag_slim" } dependencies = [ { name = "docling-core" }, { name = "httpx" }, + { name = "jsonpatch" }, { name = "lancedb" }, { name = "pathspec" }, { name = "pydantic" }, @@ -1429,6 +1430,7 @@ requires-dist = [ { name = "docling", marker = "extra == 'docling'", specifier = "==2.69.1" }, { name = "docling-core", specifier = "==2.60.1" }, { name = "httpx", specifier = ">=0.28.1" }, + { name = "jsonpatch", specifier = ">=1.33" }, { name = "lancedb", specifier = "==0.27.0" }, { name = "mxbai-rerank", marker = "extra == 'mxbai'", specifier = ">=0.1.6" }, { name = "opencv-python-headless", marker = "extra == 'docling'", specifier = ">=4.13.0.90" },