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

Merge pull request #152 from ggozad/fix/agui-activity

Browse source 
Fix ACTIVITY_SNAPSHOT payloads.
This commit is contained in:
Yiorgis Gozadinos 2025-11-21 12:35:28 +02:00 committed by GitHub
commit d8f1f3fe92
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
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