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

Enhanced ActivitySnapshot events with richer structured data

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-11-24 12:46:29 +02:00
parent 3bdbd10516
commit 13d3a366e8
No known key found for this signature in database
7 changed files with 108 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
CHANGELOG.md
View file

@ -19,6 +19,14 @@
- Dataset specifications now declare their retrieval evaluator (MRR for RepliQA, MAP for Wix)
- Replaced Recall@K and Success@K with industry-standard MRR and MAP metrics
- Unified evaluation framework for both retrieval and QA benchmarks
- **AG-UI Events**: Enhanced ActivitySnapshot events with richer structured data
- Added `stepName` field to identify which graph node emitted each activity
- Added structured fields to activity content while preserving backward-compatible `message` field:
- **Planning**: `sub_questions` - list of sub-question strings
- **Searching**: `query` - the search query, `confidence` - answer confidence (on success), `error` - error message (on failure)
- **Analyzing** (research): `insights` - list of insight objects, `gaps` - list of gap objects, `resolved_gaps` - list of resolved gap strings
- **Evaluating** (research): `confidence` - confidence score, `is_sufficient` - sufficiency flag
- **Evaluating** (deep QA): `is_sufficient` - sufficiency flag, `iterations` - iteration count
### Changed

13
docs/agents.md
View file

@ -271,7 +271,18 @@ async with HaikuRAG(path_to_db) as client:
if event["type"] == "STEP_STARTED":
print(f"Starting step: {event['stepName']}")
elif event["type"] == "ACTIVITY_SNAPSHOT":
print(f" {event['content']}")
# Activity events include structured data alongside messages
content = event['content']
print(f" {content['message']}")
# Different activity types have different structured fields
if 'confidence' in content:
print(f" Confidence: {content['confidence']:.0%}")
if 'sub_questions' in content:
for q in content['sub_questions']:
print(f" - {q}")
if 'insights' in content:
print(f" New insights: {len(content['insights'])}")
elif event["type"] == "RUN_FINISHED":
print("\nResearch complete!\n")
result = event["result"]

22
docs/server.md
View file

@ -175,9 +175,10 @@ The `--no-buffer` flag ensures curl displays events as they arrive instead of bu
**Response:** Server-Sent Events stream with AG-UI protocol events:
- `RUN_STARTED` - Graph execution started
- `STATE_SNAPSHOT` - Current state snapshot
- `STATE_DELTA` - Incremental state changes (JSON Patch format)
- `STEP_STARTED` - Node execution started
- `STEP_FINISHED` - Node execution completed
- `ACTIVITY_SNAPSHOT` - Progress update
- `ACTIVITY_SNAPSHOT` - Progress update with structured data
- `RUN_FINISHED` - Graph execution completed with result
- `RUN_ERROR` - Error during execution
@ -189,13 +190,30 @@ data: {"type":"STATE_SNAPSHOT","snapshot":{"context":{"original_question":"What
data: {"type":"STEP_STARTED","stepName":"plan"}
data: {"type":"ACTIVITY_SNAPSHOT","messageId":"msg-1","activityType":"planning","content":"Creating research plan"}
data: {"type":"ACTIVITY_SNAPSHOT","messageId":"msg-1","activityType":"planning","stepName":"plan","content":{"message":"Created plan with 3 sub-questions","sub_questions":["What is X?","How does Y work?","Why is Z important?"]}}
data: {"type":"STEP_FINISHED","stepName":"plan"}
data: {"type":"STATE_DELTA","delta":[{"op":"replace","path":"/iterations","value":1}]}
data: {"type":"ACTIVITY_SNAPSHOT","messageId":"msg-2","activityType":"evaluating","stepName":"decide","content":{"message":"Confidence: 85%, Sufficient: Yes","confidence":0.85,"is_sufficient":true}}
data: {"type":"RUN_FINISHED","threadId":"abc123","runId":"xyz789","result":{"title":"Research Report","executive_summary":"..."}}
```
**Activity Event Structure:**
`ACTIVITY_SNAPSHOT` events include a `content` object with:
- `message` - Human-readable progress message (always present)
- `stepName` - The graph step emitting this activity (when available)
- Additional structured fields depending on the activity type:
- **Planning**: `sub_questions` (list of strings)
- **Searching**: `query` (string), `confidence` (float, on completion), `error` (string, on failure)
- **Analyzing**: `insights` (list of insight objects), `gaps` (list of gap objects), `resolved_gaps` (list of strings)
- **Evaluating**: `confidence` (float), `is_sufficient` (boolean) for research; `is_sufficient` (boolean), `iterations` (int) for deep QA
The `message` field is always present for simple rendering, while structured fields enable richer UI features like displaying lists, charts, and detailed status information.
The endpoint follows the [AG-UI protocol](https://docs.ag-ui.com/concepts/events) for event streaming.
### Running Multiple Services

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

@ -79,7 +79,9 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
if deps.agui_emitter:
deps.agui_emitter.start_step("plan")
deps.agui_emitter.update_activity("planning", {"message": activity_message})
deps.agui_emitter.update_activity(
"planning", {"stepName": "plan", "message": activity_message}
)
try:
# Build agent configuration
@ -124,7 +126,12 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
deps.agui_emitter.update_state(state)
count = len(state.context.sub_questions)
deps.agui_emitter.update_activity(
"planning", {"message": f"Created plan with {count} sub-questions"}
"planning",
{
"stepName": "plan",
"message": f"Created plan with {count} sub-questions",
"sub_questions": list(state.context.sub_questions),
},
)
finally:
if deps.agui_emitter:
@ -207,7 +214,12 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
"""Internal search implementation."""
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching", {"message": f"Searching: {sub_q}"}
"searching",
{
"stepName": "search_one",
"message": f"Searching: {sub_q}",
"query": sub_q,
},
)
agent = Agent(
@ -258,14 +270,28 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
)
else:
message = success_message_format.format(sub_q=sub_q)
deps.agui_emitter.update_activity("searching", {"message": message})
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": message,
"query": sub_q,
"confidence": answer.confidence,
},
)
return answer
except Exception as e:
if handle_exceptions:
# Narrate the error
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching", {"message": f"Search failed: {e}"}
"searching",
{
"stepName": "search_one",
"message": f"Search failed: {e}",
"query": sub_q,
"error": str(e),
},
)
failure_answer = SearchAnswer(
query=sub_q,

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

@ -136,7 +136,10 @@ def build_deep_qa_graph(
deps.agui_emitter.update_activity(
"evaluating",
{
"message": f"Information {status} after {state.iterations} iteration(s)"
"stepName": "decide",
"message": f"Information {status} after {state.iterations} iteration(s)",
"is_sufficient": evaluation.is_sufficient,
"iterations": state.iterations,
},
)

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

@ -126,9 +126,9 @@ def build_research_graph(
# State updated with insights/gaps - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
highlights = len(analysis.highlights) if analysis.highlights else 0
gaps = len(analysis.gap_assessments) if analysis.gap_assessments else 0
resolved = len(analysis.resolved_gaps) if analysis.resolved_gaps else 0
highlights = len(analysis.highlights)
gaps = len(analysis.gap_assessments)
resolved = len(analysis.resolved_gaps)
parts = []
if highlights:
parts.append(f"{highlights} insights")
@ -138,7 +138,18 @@ def build_research_graph(
parts.append(f"{resolved} resolved")
summary = ", ".join(parts) if parts else "No updates"
deps.agui_emitter.update_activity(
"analyzing", {"message": f"Analysis: {summary}"}
"analyzing",
{
"stepName": "analyze_insights",
"message": f"Analysis: {summary}",
"insights": [
h.model_dump(mode="json") for h in analysis.highlights
],
"gaps": [
g.model_dump(mode="json") for g in analysis.gap_assessments
],
"resolved_gaps": list(analysis.resolved_gaps),
},
)
finally:
if deps.agui_emitter:
@ -204,7 +215,10 @@ def build_research_graph(
deps.agui_emitter.update_activity(
"evaluating",
{
"message": f"Confidence: {output.confidence_score:.0%}, Sufficient: {sufficient}"
"stepName": "decide",
"message": f"Confidence: {output.confidence_score:.0%}, Sufficient: {sufficient}",
"confidence": output.confidence_score,
"is_sufficient": output.is_sufficient,
},
)

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

@ -124,8 +124,16 @@ async def test_emitter_activity_events():
"""Test activity events."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
# Activity without a step
emitter.update_activity("processing", {"message": "Processing data"})
emitter.update_activity("done", {"message": "Completed"}, message_id="msg-1")
# Activity within a step (stepName explicitly included in content)
emitter.start_step("analyze")
emitter.update_activity(
"done", {"stepName": "analyze", "message": "Completed"}, message_id="msg-1"
)
emitter.finish_step()
await emitter.close()
events = []
@ -135,9 +143,13 @@ 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"] == {"message": "Processing data"}
assert activity_events[0]["content"]["message"] == "Processing data"
assert "stepName" not in activity_events[0]["content"] # No step context
assert activity_events[1]["messageId"] == "msg-1"
assert activity_events[1]["activityType"] == "done"
assert activity_events[1]["content"]["message"] == "Completed"
assert activity_events[1]["content"]["stepName"] == "analyze" # Has step context
@pytest.mark.asyncio