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

Fix ACTIVITY_SNAPSHOT payloads. Closes #150

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-11-21 10:58:10 +02:00
parent 09900f8e8b
commit 5c4a32f1c8
No known key found for this signature in database
11 changed files with 59 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -1,6 +1,10 @@
# Changelog
## [Unreleased]
### Fixed
- **AG-UI Activity Events**: Activity events now correctly use structured dict content instead of strings
## [0.17.2] - 2025-11-19
### Added

8
haiku_rag_slim/haiku/rag/graph/agui/emitter.py
View file

@ -3,6 +3,7 @@
import asyncio
import hashlib
from collections.abc import AsyncIterator
from typing import Any
from uuid import uuid4
from pydantic import BaseModel
@ -124,13 +125,16 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
self._last_state = new_state.model_copy(deep=True)
def update_activity(
self, activity_type: str, content: str, message_id: str | None = None
self,
activity_type: str,
content: dict[str, Any],
message_id: str | None = None,
) -> None:
"""Emit ActivitySnapshot event.
Args:
activity_type: Type of activity (e.g., "planning", "searching")
content: Description of the activity
content: Structured payload representing the activity state
message_id: Optional message ID to associate activity with (auto-generated if None)
"""
if message_id is None:

4
haiku_rag_slim/haiku/rag/graph/agui/events.py
View file

@ -211,14 +211,14 @@ def emit_state_delta(old_state: BaseModel, new_state: BaseModel) -> dict[str, An
def emit_activity(
message_id: str,
activity_type: str,
content: str,
content: dict[str, Any],
) -> dict[str, Any]:
"""Create an ActivitySnapshot event.
Args:
message_id: Message ID to associate activity with (required)
activity_type: Type of activity (e.g., "planning", "searching")
content: Description of the activity
content: Structured payload representing the activity state
Returns:
ActivitySnapshot event dict

14
haiku_rag_slim/haiku/rag/graph/common/nodes.py
View file

@ -75,7 +75,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
if deps.agui_emitter:
deps.agui_emitter.start_step("plan")
deps.agui_emitter.update_activity("planning", activity_message)
deps.agui_emitter.update_activity("planning", {"message": activity_message})
try:
# Build agent configuration
@ -120,7 +120,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
deps.agui_emitter.update_state(state)
count = len(state.context.sub_questions)
deps.agui_emitter.update_activity(
"planning", f"Created plan with {count} sub-questions"
"planning", {"message": f"Created plan with {count} sub-questions"}
)
finally:
if deps.agui_emitter:
@ -198,7 +198,9 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
) -> SearchAnswer:
"""Internal search implementation."""
if deps.agui_emitter:
deps.agui_emitter.update_activity("searching", f"Searching: {sub_q}")
deps.agui_emitter.update_activity(
"searching", {"message": f"Searching: {sub_q}"}
)
agent = Agent(
model=get_model(provider, model),
@ -248,13 +250,15 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
)
else:
message = success_message_format.format(sub_q=sub_q)
deps.agui_emitter.update_activity("searching", message)
deps.agui_emitter.update_activity("searching", {"message": message})
return answer
except Exception as e:
if handle_exceptions:
# Narrate the error
if deps.agui_emitter:
deps.agui_emitter.update_activity("searching", f"Search failed: {e}")
deps.agui_emitter.update_activity(
"searching", {"message": f"Search failed: {e}"}
)
failure_answer = SearchAnswer(
query=sub_q,
answer=f"Search failed after retries: {str(e)}",

12
haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py
View file

@ -85,7 +85,7 @@ def build_deep_qa_graph(
if deps.agui_emitter:
deps.agui_emitter.start_step("decide")
deps.agui_emitter.update_activity(
"evaluating", "Evaluating information sufficiency"
"evaluating", {"message": "Evaluating information sufficiency"}
)
try:
@ -133,7 +133,9 @@ def build_deep_qa_graph(
status = "sufficient" if evaluation.is_sufficient else "insufficient"
deps.agui_emitter.update_activity(
"evaluating",
f"Information {status} after {state.iterations} iteration(s)",
{
"message": f"Information {status} after {state.iterations} iteration(s)"
},
)
should_continue = (
@ -155,7 +157,7 @@ def build_deep_qa_graph(
if deps.agui_emitter:
deps.agui_emitter.start_step("synthesize")
deps.agui_emitter.update_activity(
"synthesizing", "Synthesizing final answer"
"synthesizing", {"message": "Synthesizing final answer"}
)
try:
@ -195,7 +197,9 @@ def build_deep_qa_graph(
result = await agent.run(prompt, deps=agent_deps)
if deps.agui_emitter:
deps.agui_emitter.update_activity("synthesizing", "Answer complete")
deps.agui_emitter.update_activity(
"synthesizing", {"message": "Answer complete"}
)
return result.output
finally:

14
haiku_rag_slim/haiku/rag/graph/research/graph.py
View file

@ -92,7 +92,7 @@ def build_research_graph(
if deps.agui_emitter:
deps.agui_emitter.start_step("analyze_insights")
deps.agui_emitter.update_activity(
"analyzing", "Synthesizing insights and gaps"
"analyzing", {"message": "Synthesizing insights and gaps"}
)
try:
@ -135,7 +135,9 @@ def build_research_graph(
if resolved:
parts.append(f"{resolved} resolved")
summary = ", ".join(parts) if parts else "No updates"
deps.agui_emitter.update_activity("analyzing", f"Analysis: {summary}")
deps.agui_emitter.update_activity(
"analyzing", {"message": f"Analysis: {summary}"}
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
@ -148,7 +150,7 @@ def build_research_graph(
if deps.agui_emitter:
deps.agui_emitter.start_step("decide")
deps.agui_emitter.update_activity(
"evaluating", "Evaluating research sufficiency"
"evaluating", {"message": "Evaluating research sufficiency"}
)
try:
@ -199,7 +201,9 @@ def build_research_graph(
sufficient = "Yes" if output.is_sufficient else "No"
deps.agui_emitter.update_activity(
"evaluating",
f"Confidence: {output.confidence_score:.0%}, Sufficient: {sufficient}",
{
"message": f"Confidence: {output.confidence_score:.0%}, Sufficient: {sufficient}"
},
)
should_continue = (
@ -222,7 +226,7 @@ def build_research_graph(
if deps.agui_emitter:
deps.agui_emitter.start_step("synthesize")
deps.agui_emitter.update_activity(
"synthesizing", "Generating final research report"
"synthesizing", {"message": "Generating final research report"}
)
try:

4
tests/graph/agui/test_cli_renderer.py
View file

@ -41,7 +41,7 @@ async def test_renderer_basic_flow():
emit_state_snapshot(SimpleState(value=1)),
emit_step_started("plan"),
emit_step_finished("plan"),
emit_activity("m1", "planning", "Planning research"),
emit_activity("m1", "planning", {"message": "Planning research"}),
emit_run_finished("t1", "r1", {"status": "complete"}),
]
@ -76,7 +76,7 @@ async def test_renderer_handles_all_event_types():
emit_step_started("step1"),
emit_step_finished("step1"),
emit_text_message("message"),
emit_activity("m1", "work", "Working"),
emit_activity("m1", "work", {"message": "Working"}),
emit_run_error("error occurred"),
emit_run_finished("t1", "r1", {"result": "done"}),
]

6
tests/graph/agui/test_emitter.py
View file

@ -124,8 +124,8 @@ async def test_emitter_activity_events():
"""Test activity events."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
emitter.update_activity("processing", "Processing data")
emitter.update_activity("done", "Completed", message_id="msg-1")
emitter.update_activity("processing", {"message": "Processing data"})
emitter.update_activity("done", {"message": "Completed"}, message_id="msg-1")
await emitter.close()
events = []
@ -135,7 +135,7 @@ async def test_emitter_activity_events():
activity_events = [e for e in events if e["type"] == "ACTIVITY_SNAPSHOT"]
assert len(activity_events) == 2
assert activity_events[0]["activityType"] == "processing"
assert activity_events[0]["content"] == "Processing data"
assert activity_events[0]["content"] == {"message": "Processing data"}
assert activity_events[1]["messageId"] == "msg-1"
assert activity_events[1]["activityType"] == "done"

6
tests/graph/agui/test_events.py
View file

@ -127,12 +127,12 @@ def test_emit_state_snapshot():
def test_emit_activity():
"""Test ACTIVITY_SNAPSHOT event creation."""
event = emit_activity("msg-1", "processing", "Working on task")
event = emit_activity("msg-1", "processing", {"message": "Working on task"})
assert event["type"] == "ACTIVITY_SNAPSHOT"
assert event["messageId"] == "msg-1"
assert event["activityType"] == "processing"
assert event["content"] == "Working on task"
assert event["content"] == {"message": "Working on task"}
def test_event_structure_consistency():
@ -145,7 +145,7 @@ def test_event_structure_consistency():
emit_step_finished("step1"),
emit_text_message("text"),
emit_state_snapshot(TestState(value=1)),
emit_activity("m1", "type", "content"),
emit_activity("m1", "type", {"content": "value"}),
]
for event in events:

8
tests/graph/agui/test_server.py
View file

@ -93,7 +93,7 @@ def test_create_agui_app_basic():
app = create_agui_app(
graph_factory=graph_factory, # type: ignore[arg-type]
state_factory=state_factory,
deps_factory=deps_factory,
deps_factory=deps_factory, # type: ignore[arg-type]
config=config,
)
@ -124,7 +124,7 @@ def test_server_health_endpoint():
app = create_agui_app(
graph_factory=graph_factory, # type: ignore[arg-type]
state_factory=state_factory,
deps_factory=deps_factory,
deps_factory=deps_factory, # type: ignore[arg-type]
config=config,
)
@ -159,7 +159,7 @@ async def test_server_stream_endpoint():
app = create_agui_app(
graph_factory=graph_factory, # type: ignore[arg-type]
state_factory=state_factory,
deps_factory=deps_factory,
deps_factory=deps_factory, # type: ignore[arg-type]
config=config,
)
@ -222,7 +222,7 @@ def test_server_cors_headers():
app = create_agui_app(
graph_factory=graph_factory, # type: ignore[arg-type]
state_factory=state_factory,
deps_factory=deps_factory,
deps_factory=deps_factory, # type: ignore[arg-type]
config=config,
)

18
tests/graph/agui/test_stream.py
View file

@ -36,7 +36,7 @@ class MockGraph:
# Emit some events through the emitter
if deps.agui_emitter:
deps.agui_emitter.start_step("mock_step")
deps.agui_emitter.update_activity("working", "Doing work")
deps.agui_emitter.update_activity("working", {"message": "Doing work"})
deps.agui_emitter.finish_step()
return self.result
@ -50,7 +50,7 @@ async def test_stream_graph_basic():
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps):
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Should have collected events
@ -73,7 +73,7 @@ async def test_stream_graph_emits_initial_state():
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps):
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Should have initial state snapshot
@ -91,7 +91,7 @@ async def test_stream_graph_emits_step_events():
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps):
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Should have step events from MockGraph
@ -108,13 +108,13 @@ async def test_stream_graph_emits_activity():
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps):
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Should have activity events from MockGraph
activities = [e for e in events if e["type"] == "ACTIVITY_SNAPSHOT"]
assert len(activities) > 0
assert activities[0]["content"] == "Doing work"
assert activities[0]["content"] == {"message": "Doing work"}
@pytest.mark.asyncio
@ -130,7 +130,7 @@ async def test_stream_graph_handles_error():
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps):
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Should have error event
@ -154,7 +154,7 @@ async def test_stream_graph_closes_emitter():
events = []
try:
async for event in stream_graph(graph, state, deps):
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
except RuntimeError:
# Expected - graph didn't return a result
@ -190,7 +190,7 @@ async def test_stream_graph_result_in_finish_event():
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps):
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Find RUN_FINISHED event