Replace custom event classes with ag_ui.core types
This commit is contained in:
Yiorgis Gozadinos
parent
f6618e6037
commit
a1f5d944ba
10 changed files with 219 additions and 545 deletions
|
|
@ -27,6 +27,9 @@
|
|||
|
||||
### Changed
|
||||
|
||||
- **AG-UI Events**: Replaced custom event classes with `ag_ui.core` types
|
||||
- Removed `haiku.rag.graph.agui.events` module
|
||||
- Event factory functions (`emit_*`) now wrap official `ag_ui.core` event classes
|
||||
- **Chunker Sets Order**: Chunkers now set `chunk.order` directly
|
||||
- **Unified Research Graph**: Simplified and unified research and deep QA into a single configurable graph
|
||||
- Removed `analyze_insights` node - graph now flows directly from `collect_answers` to `decide`
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
"""Generic AG-UI protocol support for haiku.rag graphs."""
|
||||
|
||||
from haiku.rag.graph.agui.cli_renderer import AGUIConsoleRenderer
|
||||
from haiku.rag.graph.agui.emitter import AGUIEmitter
|
||||
from haiku.rag.graph.agui.events import (
|
||||
from haiku.rag.graph.agui.emitter import (
|
||||
AGUIEmitter,
|
||||
AGUIEvent,
|
||||
emit_activity,
|
||||
emit_activity_delta,
|
||||
|
|
@ -17,6 +17,9 @@ from haiku.rag.graph.agui.events import (
|
|||
emit_text_message_content,
|
||||
emit_text_message_end,
|
||||
emit_text_message_start,
|
||||
emit_tool_call_args,
|
||||
emit_tool_call_end,
|
||||
emit_tool_call_start,
|
||||
)
|
||||
from haiku.rag.graph.agui.server import (
|
||||
RunAgentInput,
|
||||
|
|
@ -48,6 +51,9 @@ __all__ = [
|
|||
"emit_text_message_content",
|
||||
"emit_text_message_end",
|
||||
"emit_text_message_start",
|
||||
"emit_tool_call_args",
|
||||
"emit_tool_call_end",
|
||||
"emit_tool_call_start",
|
||||
"format_sse_event",
|
||||
"stream_graph",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from typing import Any
|
|||
|
||||
from rich.console import Console
|
||||
|
||||
from haiku.rag.graph.agui.events import AGUIEvent
|
||||
from haiku.rag.graph.agui.emitter import AGUIEvent
|
||||
|
||||
|
||||
class AGUIConsoleRenderer:
|
||||
|
|
|
|||
|
|
@ -2,24 +2,39 @@
|
|||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
from ag_ui.core import (
|
||||
ActivitySnapshotEvent,
|
||||
BaseEvent,
|
||||
RunErrorEvent,
|
||||
RunFinishedEvent,
|
||||
RunStartedEvent,
|
||||
StateDeltaEvent,
|
||||
StateSnapshotEvent,
|
||||
StepFinishedEvent,
|
||||
StepStartedEvent,
|
||||
TextMessageChunkEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
TextMessageStartEvent,
|
||||
ToolCallArgsEvent,
|
||||
ToolCallEndEvent,
|
||||
ToolCallStartEvent,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
|
||||
from haiku.rag.graph.agui.events import (
|
||||
AGUIEvent,
|
||||
emit_activity,
|
||||
emit_run_error,
|
||||
emit_run_finished,
|
||||
emit_run_started,
|
||||
emit_state_delta,
|
||||
emit_state_snapshot,
|
||||
emit_step_finished,
|
||||
emit_step_started,
|
||||
emit_text_message,
|
||||
)
|
||||
from haiku.rag.graph.agui.state import compute_state_delta
|
||||
|
||||
AGUIEvent = dict[str, Any]
|
||||
|
||||
|
||||
def _serialize_event(event: BaseEvent) -> AGUIEvent:
|
||||
"""Serialize an ag_ui event to a dict with camelCase keys."""
|
||||
return event.model_dump(mode="json", by_alias=True, exclude_none=True)
|
||||
|
||||
|
||||
class AGUIEmitter[StateT: BaseModel, ResultT]:
|
||||
|
|
@ -80,8 +95,14 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
|
|||
self._thread_id = self._generate_thread_id(state_json)
|
||||
|
||||
# RunStarted (state snapshot follows immediately with full state)
|
||||
self.emit(emit_run_started(self._thread_id, self._run_id))
|
||||
self.emit(emit_state_snapshot(initial_state))
|
||||
self.emit(
|
||||
_serialize_event(
|
||||
RunStartedEvent(thread_id=self._thread_id, run_id=self._run_id)
|
||||
)
|
||||
)
|
||||
self.emit(
|
||||
_serialize_event(StateSnapshotEvent(snapshot=initial_state.model_dump()))
|
||||
)
|
||||
# Store a deep copy to detect future changes
|
||||
self._last_state = initial_state.model_copy(deep=True)
|
||||
|
||||
|
|
@ -92,12 +113,12 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
|
|||
step_name: Name of the step being started
|
||||
"""
|
||||
self._current_step = step_name
|
||||
self.emit(emit_step_started(step_name))
|
||||
self.emit(_serialize_event(StepStartedEvent(step_name=step_name)))
|
||||
|
||||
def finish_step(self) -> None:
|
||||
"""Emit StepFinished event for the current step."""
|
||||
if self._current_step:
|
||||
self.emit(emit_step_finished(self._current_step))
|
||||
self.emit(_serialize_event(StepFinishedEvent(step_name=self._current_step)))
|
||||
self._current_step = None
|
||||
|
||||
def log(self, message: str, role: str = "assistant") -> None:
|
||||
|
|
@ -107,7 +128,16 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
|
|||
message: The message content
|
||||
role: The role of the sender (default: assistant)
|
||||
"""
|
||||
self.emit(emit_text_message(message, role))
|
||||
message_id = str(uuid4())
|
||||
self.emit(
|
||||
_serialize_event(
|
||||
TextMessageChunkEvent(
|
||||
message_id=message_id,
|
||||
role=role, # type: ignore[arg-type]
|
||||
delta=message,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
def update_state(self, new_state: StateT) -> None:
|
||||
"""Emit StateDelta or StateSnapshot for state change.
|
||||
|
|
@ -117,10 +147,13 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
|
|||
"""
|
||||
if self._use_deltas and self._last_state is not None:
|
||||
# Emit delta for incremental updates
|
||||
self.emit(emit_state_delta(self._last_state, new_state))
|
||||
delta = compute_state_delta(self._last_state, new_state)
|
||||
self.emit(_serialize_event(StateDeltaEvent(delta=delta)))
|
||||
else:
|
||||
# Emit full snapshot for initial state or when deltas disabled
|
||||
self.emit(emit_state_snapshot(new_state))
|
||||
self.emit(
|
||||
_serialize_event(StateSnapshotEvent(snapshot=new_state.model_dump()))
|
||||
)
|
||||
# Store a deep copy to detect future changes
|
||||
self._last_state = new_state.model_copy(deep=True)
|
||||
|
||||
|
|
@ -139,7 +172,15 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
|
|||
"""
|
||||
if message_id is None:
|
||||
message_id = str(uuid4())
|
||||
self.emit(emit_activity(message_id, activity_type, content))
|
||||
self.emit(
|
||||
_serialize_event(
|
||||
ActivitySnapshotEvent(
|
||||
message_id=message_id,
|
||||
activity_type=activity_type,
|
||||
content=content,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
def finish_run(self, result: ResultT) -> None:
|
||||
"""Emit RunFinished event.
|
||||
|
|
@ -147,7 +188,18 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
|
|||
Args:
|
||||
result: The final result from the graph
|
||||
"""
|
||||
self.emit(emit_run_finished(self._thread_id, self._run_id, result))
|
||||
# Convert result to dict if it's a Pydantic model
|
||||
result_data: Any = result
|
||||
if hasattr(result, "model_dump"):
|
||||
result_data = result.model_dump() # type: ignore[union-attr]
|
||||
|
||||
self.emit(
|
||||
_serialize_event(
|
||||
RunFinishedEvent(
|
||||
thread_id=self._thread_id, run_id=self._run_id, result=result_data
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
def error(self, error: Exception, code: str | None = None) -> None:
|
||||
"""Emit RunError event.
|
||||
|
|
@ -156,7 +208,7 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
|
|||
error: The exception that occurred
|
||||
code: Optional error code
|
||||
"""
|
||||
self.emit(emit_run_error(str(error), code))
|
||||
self.emit(_serialize_event(RunErrorEvent(message=str(error), code=code)))
|
||||
|
||||
def emit(self, event: AGUIEvent) -> None:
|
||||
"""Put event in queue.
|
||||
|
|
@ -199,3 +251,132 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
|
|||
# Use hash of input for deterministic thread ID
|
||||
hash_obj = hashlib.sha256(input_data.encode("utf-8"))
|
||||
return hash_obj.hexdigest()[:16]
|
||||
|
||||
|
||||
def emit_text_message_start(message_id: str, role: str = "assistant") -> AGUIEvent:
|
||||
"""Create a TextMessageStart event."""
|
||||
return _serialize_event(
|
||||
TextMessageStartEvent(message_id=message_id, role=role) # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
|
||||
def emit_text_message_content(message_id: str, delta: str) -> AGUIEvent:
|
||||
"""Create a TextMessageContent event."""
|
||||
return _serialize_event(TextMessageContentEvent(message_id=message_id, delta=delta))
|
||||
|
||||
|
||||
def emit_text_message_end(message_id: str) -> AGUIEvent:
|
||||
"""Create a TextMessageEnd event."""
|
||||
return _serialize_event(TextMessageEndEvent(message_id=message_id))
|
||||
|
||||
|
||||
def emit_tool_call_start(
|
||||
tool_call_id: str,
|
||||
tool_name: str,
|
||||
parent_message_id: str | None = None,
|
||||
) -> AGUIEvent:
|
||||
"""Create a ToolCallStart event."""
|
||||
return _serialize_event(
|
||||
ToolCallStartEvent(
|
||||
tool_call_id=tool_call_id,
|
||||
tool_call_name=tool_name,
|
||||
parent_message_id=parent_message_id,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def emit_tool_call_args(tool_call_id: str, args: dict[str, Any]) -> AGUIEvent:
|
||||
"""Create a ToolCallArgs event."""
|
||||
return _serialize_event(
|
||||
ToolCallArgsEvent(tool_call_id=tool_call_id, delta=json.dumps(args))
|
||||
)
|
||||
|
||||
|
||||
def emit_tool_call_end(tool_call_id: str) -> AGUIEvent:
|
||||
"""Create a ToolCallEnd event."""
|
||||
return _serialize_event(ToolCallEndEvent(tool_call_id=tool_call_id))
|
||||
|
||||
|
||||
def emit_run_started(thread_id: str, run_id: str) -> AGUIEvent:
|
||||
"""Create a RunStarted event."""
|
||||
return _serialize_event(RunStartedEvent(thread_id=thread_id, run_id=run_id))
|
||||
|
||||
|
||||
def emit_run_finished(thread_id: str, run_id: str, result: Any) -> AGUIEvent:
|
||||
"""Create a RunFinished event."""
|
||||
# Convert result to dict if it's a Pydantic model
|
||||
if hasattr(result, "model_dump"):
|
||||
result = result.model_dump()
|
||||
return _serialize_event(
|
||||
RunFinishedEvent(thread_id=thread_id, run_id=run_id, result=result)
|
||||
)
|
||||
|
||||
|
||||
def emit_run_error(message: str, code: str | None = None) -> AGUIEvent:
|
||||
"""Create a RunError event."""
|
||||
return _serialize_event(RunErrorEvent(message=message, code=code))
|
||||
|
||||
|
||||
def emit_step_started(step_name: str) -> AGUIEvent:
|
||||
"""Create a StepStarted event."""
|
||||
return _serialize_event(StepStartedEvent(step_name=step_name))
|
||||
|
||||
|
||||
def emit_step_finished(step_name: str) -> AGUIEvent:
|
||||
"""Create a StepFinished event."""
|
||||
return _serialize_event(StepFinishedEvent(step_name=step_name))
|
||||
|
||||
|
||||
def emit_text_message(content: str, role: str = "assistant") -> AGUIEvent:
|
||||
"""Create a TextMessageChunk event (convenience wrapper)."""
|
||||
message_id = str(uuid4())
|
||||
return _serialize_event(
|
||||
TextMessageChunkEvent(
|
||||
message_id=message_id,
|
||||
role=role, # type: ignore[arg-type]
|
||||
delta=content,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def emit_state_snapshot(state: BaseModel) -> AGUIEvent:
|
||||
"""Create a StateSnapshot event."""
|
||||
return _serialize_event(StateSnapshotEvent(snapshot=state.model_dump()))
|
||||
|
||||
|
||||
def emit_state_delta(old_state: BaseModel, new_state: BaseModel) -> AGUIEvent:
|
||||
"""Create a StateDelta event with JSON Patch operations."""
|
||||
delta = compute_state_delta(old_state, new_state)
|
||||
return _serialize_event(StateDeltaEvent(delta=delta))
|
||||
|
||||
|
||||
def emit_activity(
|
||||
message_id: str,
|
||||
activity_type: str,
|
||||
content: dict[str, Any],
|
||||
) -> AGUIEvent:
|
||||
"""Create an ActivitySnapshot event."""
|
||||
return _serialize_event(
|
||||
ActivitySnapshotEvent(
|
||||
message_id=message_id,
|
||||
activity_type=activity_type,
|
||||
content=content,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def emit_activity_delta(
|
||||
message_id: str,
|
||||
activity_type: str,
|
||||
patch: list[dict[str, Any]],
|
||||
) -> AGUIEvent:
|
||||
"""Create an ActivityDelta event with JSON Patch operations."""
|
||||
from ag_ui.core import ActivityDeltaEvent
|
||||
|
||||
return _serialize_event(
|
||||
ActivityDeltaEvent(
|
||||
message_id=message_id,
|
||||
activity_type=activity_type,
|
||||
patch=patch,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,313 +0,0 @@
|
|||
"""Generic AG-UI event creation utilities for any graph."""
|
||||
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from haiku.rag.graph.agui.state import compute_state_delta
|
||||
|
||||
# Type aliases for AG-UI events (actual types from ag_ui.core will be used at runtime)
|
||||
AGUIEvent = dict[str, Any]
|
||||
|
||||
|
||||
def emit_run_started(
|
||||
thread_id: str, run_id: str, input_data: str | None = None
|
||||
) -> dict[str, Any]:
|
||||
"""Create a RunStarted event.
|
||||
|
||||
Args:
|
||||
thread_id: Unique identifier for the conversation thread
|
||||
run_id: Unique identifier for this run
|
||||
input_data: Optional input that started the run
|
||||
|
||||
Returns:
|
||||
RunStarted event dict
|
||||
"""
|
||||
event: dict[str, Any] = {
|
||||
"type": "RUN_STARTED",
|
||||
"threadId": thread_id,
|
||||
"runId": run_id,
|
||||
}
|
||||
if input_data:
|
||||
event["input"] = input_data
|
||||
return event
|
||||
|
||||
|
||||
def emit_run_finished(thread_id: str, run_id: str, result: Any) -> dict[str, Any]:
|
||||
"""Create a RunFinished event.
|
||||
|
||||
Args:
|
||||
thread_id: Unique identifier for the conversation thread
|
||||
run_id: Unique identifier for this run
|
||||
result: The final result of the run
|
||||
|
||||
Returns:
|
||||
RunFinished event dict
|
||||
"""
|
||||
# Convert result to dict if it's a Pydantic model
|
||||
if hasattr(result, "model_dump"):
|
||||
result = result.model_dump()
|
||||
|
||||
return {
|
||||
"type": "RUN_FINISHED",
|
||||
"threadId": thread_id,
|
||||
"runId": run_id,
|
||||
"result": result,
|
||||
}
|
||||
|
||||
|
||||
def emit_run_error(message: str, code: str | None = None) -> dict[str, Any]:
|
||||
"""Create a RunError event.
|
||||
|
||||
Args:
|
||||
message: Error message
|
||||
code: Optional error code
|
||||
|
||||
Returns:
|
||||
RunError event dict
|
||||
"""
|
||||
event: dict[str, Any] = {
|
||||
"type": "RUN_ERROR",
|
||||
"message": message,
|
||||
}
|
||||
if code:
|
||||
event["code"] = code
|
||||
return event
|
||||
|
||||
|
||||
def emit_step_started(step_name: str) -> dict[str, Any]:
|
||||
"""Create a StepStarted event.
|
||||
|
||||
Args:
|
||||
step_name: Name of the step being started
|
||||
|
||||
Returns:
|
||||
StepStarted event dict
|
||||
"""
|
||||
return {
|
||||
"type": "STEP_STARTED",
|
||||
"stepName": step_name,
|
||||
}
|
||||
|
||||
|
||||
def emit_step_finished(step_name: str) -> dict[str, Any]:
|
||||
"""Create a StepFinished event.
|
||||
|
||||
Args:
|
||||
step_name: Name of the step that finished
|
||||
|
||||
Returns:
|
||||
StepFinished event dict
|
||||
"""
|
||||
return {
|
||||
"type": "STEP_FINISHED",
|
||||
"stepName": step_name,
|
||||
}
|
||||
|
||||
|
||||
def emit_text_message(content: str, role: str = "assistant") -> dict[str, Any]:
|
||||
"""Create a TextMessageChunk event (convenience wrapper).
|
||||
|
||||
This creates a complete text message in one event.
|
||||
|
||||
Args:
|
||||
content: The message content
|
||||
role: The role of the sender (default: assistant)
|
||||
|
||||
Returns:
|
||||
TextMessageChunk event dict
|
||||
"""
|
||||
message_id = str(uuid4())
|
||||
return {
|
||||
"type": "TEXT_MESSAGE_CHUNK",
|
||||
"messageId": message_id,
|
||||
"role": role,
|
||||
"delta": content,
|
||||
}
|
||||
|
||||
|
||||
def emit_text_message_start(message_id: str, role: str = "assistant") -> dict[str, Any]:
|
||||
"""Create a TextMessageStart event.
|
||||
|
||||
Args:
|
||||
message_id: Unique identifier for this message
|
||||
role: The role of the sender
|
||||
|
||||
Returns:
|
||||
TextMessageStart event dict
|
||||
"""
|
||||
return {
|
||||
"type": "TEXT_MESSAGE_START",
|
||||
"messageId": message_id,
|
||||
"role": role,
|
||||
}
|
||||
|
||||
|
||||
def emit_text_message_content(message_id: str, delta: str) -> dict[str, Any]:
|
||||
"""Create a TextMessageContent event.
|
||||
|
||||
Args:
|
||||
message_id: Identifier for the message being streamed
|
||||
delta: Content chunk to append
|
||||
|
||||
Returns:
|
||||
TextMessageContent event dict
|
||||
"""
|
||||
return {
|
||||
"type": "TEXT_MESSAGE_CONTENT",
|
||||
"messageId": message_id,
|
||||
"delta": delta,
|
||||
}
|
||||
|
||||
|
||||
def emit_text_message_end(message_id: str) -> dict[str, Any]:
|
||||
"""Create a TextMessageEnd event.
|
||||
|
||||
Args:
|
||||
message_id: Identifier for the message being completed
|
||||
|
||||
Returns:
|
||||
TextMessageEnd event dict
|
||||
"""
|
||||
return {
|
||||
"type": "TEXT_MESSAGE_END",
|
||||
"messageId": message_id,
|
||||
}
|
||||
|
||||
|
||||
def emit_state_snapshot(state: BaseModel) -> dict[str, Any]:
|
||||
"""Create a StateSnapshot event.
|
||||
|
||||
Args:
|
||||
state: The complete state to snapshot (any Pydantic BaseModel)
|
||||
|
||||
Returns:
|
||||
StateSnapshot event dict
|
||||
"""
|
||||
return {
|
||||
"type": "STATE_SNAPSHOT",
|
||||
"snapshot": state.model_dump(),
|
||||
}
|
||||
|
||||
|
||||
def emit_state_delta(old_state: BaseModel, new_state: BaseModel) -> dict[str, Any]:
|
||||
"""Create a StateDelta event with JSON Patch operations.
|
||||
|
||||
Args:
|
||||
old_state: Previous state (any Pydantic BaseModel)
|
||||
new_state: Current state (same type as old_state)
|
||||
|
||||
Returns:
|
||||
StateDelta event dict
|
||||
"""
|
||||
delta = compute_state_delta(old_state, new_state)
|
||||
return {
|
||||
"type": "STATE_DELTA",
|
||||
"delta": delta,
|
||||
}
|
||||
|
||||
|
||||
def emit_activity(
|
||||
message_id: str,
|
||||
activity_type: 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: Structured payload representing the activity state
|
||||
|
||||
Returns:
|
||||
ActivitySnapshot event dict
|
||||
"""
|
||||
return {
|
||||
"type": "ACTIVITY_SNAPSHOT",
|
||||
"messageId": message_id,
|
||||
"activityType": activity_type,
|
||||
"content": content,
|
||||
}
|
||||
|
||||
|
||||
def emit_activity_delta(
|
||||
message_id: str,
|
||||
activity_type: str,
|
||||
patch: list[dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
"""Create an ActivityDelta event with JSON Patch operations.
|
||||
|
||||
Args:
|
||||
message_id: Message ID of the activity being updated
|
||||
activity_type: Type of activity being updated
|
||||
patch: JSON Patch operations to apply
|
||||
|
||||
Returns:
|
||||
ActivityDelta event dict
|
||||
"""
|
||||
return {
|
||||
"type": "ACTIVITY_DELTA",
|
||||
"messageId": message_id,
|
||||
"activityType": activity_type,
|
||||
"patch": patch,
|
||||
}
|
||||
|
||||
|
||||
def emit_tool_call_start(
|
||||
tool_call_id: str,
|
||||
tool_name: str,
|
||||
parent_message_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Create a ToolCallStart event.
|
||||
|
||||
Args:
|
||||
tool_call_id: Unique identifier for this tool call
|
||||
tool_name: Name of the tool being called
|
||||
parent_message_id: Optional parent message ID
|
||||
|
||||
Returns:
|
||||
ToolCallStart event dict
|
||||
"""
|
||||
event: dict[str, Any] = {
|
||||
"type": "TOOL_CALL_START",
|
||||
"toolCallId": tool_call_id,
|
||||
"toolCallName": tool_name,
|
||||
}
|
||||
if parent_message_id:
|
||||
event["parentMessageId"] = parent_message_id
|
||||
return event
|
||||
|
||||
|
||||
def emit_tool_call_args(tool_call_id: str, args: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Create a ToolCallArgs event.
|
||||
|
||||
Args:
|
||||
tool_call_id: Identifier for the tool call
|
||||
args: Tool arguments
|
||||
|
||||
Returns:
|
||||
ToolCallArgs event dict
|
||||
"""
|
||||
import json
|
||||
|
||||
return {
|
||||
"type": "TOOL_CALL_ARGS",
|
||||
"toolCallId": tool_call_id,
|
||||
"delta": json.dumps(args),
|
||||
}
|
||||
|
||||
|
||||
def emit_tool_call_end(tool_call_id: str) -> dict[str, Any]:
|
||||
"""Create a ToolCallEnd event.
|
||||
|
||||
Args:
|
||||
tool_call_id: Identifier for the tool call being completed
|
||||
|
||||
Returns:
|
||||
ToolCallEnd event dict
|
||||
"""
|
||||
return {
|
||||
"type": "TOOL_CALL_END",
|
||||
"toolCallId": tool_call_id,
|
||||
}
|
||||
|
|
@ -18,8 +18,7 @@ from starlette.responses import JSONResponse, StreamingResponse
|
|||
from starlette.routing import Route
|
||||
|
||||
from haiku.rag.config.models import AGUIConfig
|
||||
from haiku.rag.graph.agui.emitter import AGUIEmitter
|
||||
from haiku.rag.graph.agui.events import AGUIEvent
|
||||
from haiku.rag.graph.agui.emitter import AGUIEmitter, AGUIEvent
|
||||
from haiku.rag.graph.agui.stream import stream_graph
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,8 +8,7 @@ from typing import Protocol, TypeVar
|
|||
from pydantic import BaseModel
|
||||
from pydantic_graph.beta import Graph
|
||||
|
||||
from haiku.rag.graph.agui.emitter import AGUIEmitter
|
||||
from haiku.rag.graph.agui.events import AGUIEvent
|
||||
from haiku.rag.graph.agui.emitter import AGUIEmitter, AGUIEvent
|
||||
|
||||
StateT = TypeVar("StateT", bound=BaseModel)
|
||||
ResultT = TypeVar("ResultT")
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from pydantic_graph.beta.join import reduce_list_append
|
|||
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.graph.agui.events import (
|
||||
from haiku.rag.graph.agui.emitter import (
|
||||
emit_text_message_end,
|
||||
emit_text_message_start,
|
||||
emit_tool_call_args,
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from pydantic import BaseModel
|
|||
from rich.console import Console
|
||||
|
||||
from haiku.rag.graph.agui.cli_renderer import AGUIConsoleRenderer
|
||||
from haiku.rag.graph.agui.events import (
|
||||
from haiku.rag.graph.agui.emitter import (
|
||||
emit_activity,
|
||||
emit_run_error,
|
||||
emit_run_finished,
|
||||
|
|
|
|||
|
|
@ -1,201 +0,0 @@
|
|||
"""Tests for AG-UI event creation utilities."""
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from haiku.rag.graph.agui.events import (
|
||||
emit_activity,
|
||||
emit_run_error,
|
||||
emit_run_finished,
|
||||
emit_run_started,
|
||||
emit_state_snapshot,
|
||||
emit_step_finished,
|
||||
emit_step_started,
|
||||
emit_text_message,
|
||||
emit_tool_call_args,
|
||||
emit_tool_call_end,
|
||||
emit_tool_call_start,
|
||||
)
|
||||
|
||||
|
||||
class TestState(BaseModel):
|
||||
"""Test state model."""
|
||||
|
||||
value: int
|
||||
|
||||
|
||||
class TestResult(BaseModel):
|
||||
"""Test result model."""
|
||||
|
||||
status: str
|
||||
|
||||
|
||||
def test_emit_run_started():
|
||||
"""Test RUN_STARTED event creation."""
|
||||
event = emit_run_started("thread-1", "run-1")
|
||||
|
||||
assert event["type"] == "RUN_STARTED"
|
||||
assert event["threadId"] == "thread-1"
|
||||
assert event["runId"] == "run-1"
|
||||
assert "input" not in event
|
||||
|
||||
|
||||
def test_emit_run_started_with_input():
|
||||
"""Test RUN_STARTED event with input data."""
|
||||
event = emit_run_started("thread-1", "run-1", input_data="test input")
|
||||
|
||||
assert event["type"] == "RUN_STARTED"
|
||||
assert event["input"] == "test input"
|
||||
|
||||
|
||||
def test_emit_run_finished():
|
||||
"""Test RUN_FINISHED event creation."""
|
||||
result = TestResult(status="complete")
|
||||
event = emit_run_finished("thread-1", "run-1", result)
|
||||
|
||||
assert event["type"] == "RUN_FINISHED"
|
||||
assert event["threadId"] == "thread-1"
|
||||
assert event["runId"] == "run-1"
|
||||
assert event["result"] == {"status": "complete"}
|
||||
|
||||
|
||||
def test_emit_run_finished_with_dict():
|
||||
"""Test RUN_FINISHED event with dict result."""
|
||||
result = {"status": "complete", "count": 42}
|
||||
event = emit_run_finished("thread-1", "run-1", result)
|
||||
|
||||
assert event["type"] == "RUN_FINISHED"
|
||||
assert event["result"] == result
|
||||
|
||||
|
||||
def test_emit_run_error():
|
||||
"""Test RUN_ERROR event creation."""
|
||||
event = emit_run_error("Something went wrong")
|
||||
|
||||
assert event["type"] == "RUN_ERROR"
|
||||
assert event["message"] == "Something went wrong"
|
||||
assert "code" not in event
|
||||
|
||||
|
||||
def test_emit_run_error_with_code():
|
||||
"""Test RUN_ERROR event with error code."""
|
||||
event = emit_run_error("Something went wrong", code="ERR_001")
|
||||
|
||||
assert event["type"] == "RUN_ERROR"
|
||||
assert event["message"] == "Something went wrong"
|
||||
assert event["code"] == "ERR_001"
|
||||
|
||||
|
||||
def test_emit_step_started():
|
||||
"""Test STEP_STARTED event creation."""
|
||||
event = emit_step_started("plan")
|
||||
|
||||
assert event["type"] == "STEP_STARTED"
|
||||
assert event["stepName"] == "plan"
|
||||
|
||||
|
||||
def test_emit_step_finished():
|
||||
"""Test STEP_FINISHED event creation."""
|
||||
event = emit_step_finished("plan")
|
||||
|
||||
assert event["type"] == "STEP_FINISHED"
|
||||
assert event["stepName"] == "plan"
|
||||
|
||||
|
||||
def test_emit_text_message():
|
||||
"""Test TEXT_MESSAGE_CHUNK event creation."""
|
||||
event = emit_text_message("Hello world")
|
||||
|
||||
assert event["type"] == "TEXT_MESSAGE_CHUNK"
|
||||
assert event["delta"] == "Hello world"
|
||||
assert event["role"] == "assistant"
|
||||
assert "messageId" in event
|
||||
|
||||
|
||||
def test_emit_text_message_with_role():
|
||||
"""Test TEXT_MESSAGE_CHUNK event with custom role."""
|
||||
event = emit_text_message("Hello", role="user")
|
||||
|
||||
assert event["type"] == "TEXT_MESSAGE_CHUNK"
|
||||
assert event["role"] == "user"
|
||||
|
||||
|
||||
def test_emit_state_snapshot():
|
||||
"""Test STATE_SNAPSHOT event creation."""
|
||||
state = TestState(value=42)
|
||||
event = emit_state_snapshot(state)
|
||||
|
||||
assert event["type"] == "STATE_SNAPSHOT"
|
||||
assert event["snapshot"] == {"value": 42}
|
||||
|
||||
|
||||
def test_emit_activity():
|
||||
"""Test ACTIVITY_SNAPSHOT event creation."""
|
||||
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"] == {"message": "Working on task"}
|
||||
|
||||
|
||||
def test_emit_tool_call_start():
|
||||
"""Test TOOL_CALL_START event creation."""
|
||||
event = emit_tool_call_start("call-1", "search_documents")
|
||||
|
||||
assert event["type"] == "TOOL_CALL_START"
|
||||
assert event["toolCallId"] == "call-1"
|
||||
assert event["toolCallName"] == "search_documents"
|
||||
assert "parentMessageId" not in event
|
||||
|
||||
|
||||
def test_emit_tool_call_start_with_parent():
|
||||
"""Test TOOL_CALL_START event with parent message ID."""
|
||||
event = emit_tool_call_start("call-1", "search", parent_message_id="msg-1")
|
||||
|
||||
assert event["type"] == "TOOL_CALL_START"
|
||||
assert event["toolCallId"] == "call-1"
|
||||
assert event["toolCallName"] == "search"
|
||||
assert event["parentMessageId"] == "msg-1"
|
||||
|
||||
|
||||
def test_emit_tool_call_args():
|
||||
"""Test TOOL_CALL_ARGS event creation."""
|
||||
import json
|
||||
|
||||
args = {"query": "test query", "limit": 10}
|
||||
event = emit_tool_call_args("call-1", args)
|
||||
|
||||
assert event["type"] == "TOOL_CALL_ARGS"
|
||||
assert event["toolCallId"] == "call-1"
|
||||
assert event["delta"] == json.dumps(args)
|
||||
|
||||
|
||||
def test_emit_tool_call_end():
|
||||
"""Test TOOL_CALL_END event creation."""
|
||||
event = emit_tool_call_end("call-1")
|
||||
|
||||
assert event["type"] == "TOOL_CALL_END"
|
||||
assert event["toolCallId"] == "call-1"
|
||||
|
||||
|
||||
def test_event_structure_consistency():
|
||||
"""Test that all events have consistent structure."""
|
||||
events = [
|
||||
emit_run_started("t1", "r1"),
|
||||
emit_run_finished("t1", "r1", {"result": "done"}),
|
||||
emit_run_error("error"),
|
||||
emit_step_started("step1"),
|
||||
emit_step_finished("step1"),
|
||||
emit_text_message("text"),
|
||||
emit_state_snapshot(TestState(value=1)),
|
||||
emit_activity("m1", "type", {"content": "value"}),
|
||||
emit_tool_call_start("c1", "tool"),
|
||||
emit_tool_call_args("c1", {"arg": "value"}),
|
||||
emit_tool_call_end("c1"),
|
||||
]
|
||||
|
||||
for event in events:
|
||||
assert isinstance(event, dict)
|
||||
assert "type" in event
|
||||
assert isinstance(event["type"], str)
|
||||
assert event["type"].isupper() # Event types are uppercase
|
||||
Loading…
Reference in a new issue