diff --git a/CHANGELOG.md b/CHANGELOG.md index acb50d7a..22551695 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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` diff --git a/haiku_rag_slim/haiku/rag/graph/agui/__init__.py b/haiku_rag_slim/haiku/rag/graph/agui/__init__.py index 7259a6b0..ede844a6 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/__init__.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/__init__.py @@ -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", ] diff --git a/haiku_rag_slim/haiku/rag/graph/agui/cli_renderer.py b/haiku_rag_slim/haiku/rag/graph/agui/cli_renderer.py index 3ff21abc..59a88e6d 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/cli_renderer.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/cli_renderer.py @@ -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: diff --git a/haiku_rag_slim/haiku/rag/graph/agui/emitter.py b/haiku_rag_slim/haiku/rag/graph/agui/emitter.py index 7df847f5..7e9cc7f9 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/emitter.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/emitter.py @@ -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, + ) + ) diff --git a/haiku_rag_slim/haiku/rag/graph/agui/events.py b/haiku_rag_slim/haiku/rag/graph/agui/events.py deleted file mode 100644 index 5cb1728b..00000000 --- a/haiku_rag_slim/haiku/rag/graph/agui/events.py +++ /dev/null @@ -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, - } diff --git a/haiku_rag_slim/haiku/rag/graph/agui/server.py b/haiku_rag_slim/haiku/rag/graph/agui/server.py index 26b463e6..41c9d1f0 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/server.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/server.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/graph/agui/stream.py b/haiku_rag_slim/haiku/rag/graph/agui/stream.py index acde9298..9612d1a6 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/stream.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/stream.py @@ -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") diff --git a/haiku_rag_slim/haiku/rag/graph/research/graph.py b/haiku_rag_slim/haiku/rag/graph/research/graph.py index e3c8befa..f81849b6 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/research/graph.py @@ -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, diff --git a/tests/graph/agui/test_cli_renderer.py b/tests/graph/agui/test_cli_renderer.py index 8c31f3d3..9d1ef82a 100644 --- a/tests/graph/agui/test_cli_renderer.py +++ b/tests/graph/agui/test_cli_renderer.py @@ -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, diff --git a/tests/graph/agui/test_events.py b/tests/graph/agui/test_events.py deleted file mode 100644 index 5c8795d2..00000000 --- a/tests/graph/agui/test_events.py +++ /dev/null @@ -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