From 026825ce02c86c7e1ed445be426efc5748ceb772 Mon Sep 17 00:00:00 2001 From: Ken Date: Tue, 26 May 2026 18:29:26 +0000 Subject: [PATCH] feat: AG-UI adapter translating amplifierd SSE events to AG-UI protocol (with MCP App metadata passthrough) --- backend/agui_adapter.py | 179 ++++++++++++++++++++++++++++++ tests/test_agui_adapter.py | 220 +++++++++++++++++++++++++++++++++++++ 2 files changed, 399 insertions(+) create mode 100644 backend/agui_adapter.py create mode 100644 tests/test_agui_adapter.py diff --git a/backend/agui_adapter.py b/backend/agui_adapter.py new file mode 100644 index 0000000..dfb0471 --- /dev/null +++ b/backend/agui_adapter.py @@ -0,0 +1,179 @@ +"""AG-UI adapter: translates amplifierd SSE events to AG-UI protocol events. + +Handles content_block:start/delta/end, tool:result, and orchestrator:complete +events from amplifierd and converts them into the AG-UI event types that +CopilotKit expects. Passes through MCP App metadata (_meta and structuredContent) +from tool results for iframe rendering. +""" + +import uuid +from typing import Any + +from ag_ui.core import ( + RunFinishedEvent, + TextMessageContentEvent, + TextMessageEndEvent, + TextMessageStartEvent, + ToolCallArgsEvent, + ToolCallEndEvent, + ToolCallResultEvent, + ToolCallStartEvent, +) + + +class AmplifierdEventTranslator: + """Translates amplifierd SSE events to AG-UI protocol events. + + Args: + run_id: The AG-UI run ID for this conversation run. + thread_id: The AG-UI thread ID for the current thread. + """ + + def __init__(self, run_id: str, thread_id: str) -> None: + self._run_id = run_id + self._thread_id = thread_id + # Tracks active blocks by index: {index: {type, id?, name?, message_id?}} + self._blocks: dict[int, dict[str, Any]] = {} + self._current_message_id: str | None = None + + def translate(self, event: dict[str, Any]) -> list[Any]: + """Translate a single amplifierd SSE event into zero or more AG-UI events. + + Args: + event: Dict with 'event_type' and 'data' keys. + + Returns: + List of ag_ui.core event objects (may be empty for unknown events). + """ + event_type = event.get("event_type", "") + data = event.get("data", {}) + + if event_type == "content_block:start": + return self._handle_block_start(data) + elif event_type == "content_block:delta": + return self._handle_block_delta(data) + elif event_type == "content_block:end": + return self._handle_block_end(data) + elif event_type == "tool:result": + return self._handle_tool_result(data) + elif event_type == "orchestrator:complete": + return self._handle_complete() + else: + return [] + + def _handle_block_start(self, data: dict[str, Any]) -> list[Any]: + """Handle content_block:start — emit TEXT_MESSAGE_START or TOOL_CALL_START.""" + index = data.get("index", 0) + block_type = data.get("type", "text") + + if block_type == "tool_use": + tool_id = data.get("id", "") + tool_name = data.get("name", "") + self._blocks[index] = { + "type": "tool_use", + "tool_id": tool_id, + "tool_name": tool_name, + } + return [ + ToolCallStartEvent( + tool_call_id=tool_id, + tool_call_name=tool_name, + parent_message_id=self._current_message_id, + ) + ] + else: + # Text block + message_id = str(uuid.uuid4()) + self._blocks[index] = {"type": "text", "message_id": message_id} + self._current_message_id = message_id + return [ + TextMessageStartEvent( + message_id=message_id, + role="assistant", + ) + ] + + def _handle_block_delta(self, data: dict[str, Any]) -> list[Any]: + """Handle content_block:delta — emit TEXT_MESSAGE_CONTENT or TOOL_CALL_ARGS.""" + index = data.get("index", 0) + delta = data.get("delta", {}) + delta_type = delta.get("type", "") + + block_info = self._blocks.get(index) + + if block_info is None: + # Graceful handling: no matching start — treat as tool args if input_json_delta + if delta_type == "input_json_delta": + return [ + ToolCallArgsEvent( + tool_call_id="", + delta=delta.get("partial_json", ""), + ) + ] + return [] + + if block_info["type"] == "tool_use": + return [ + ToolCallArgsEvent( + tool_call_id=block_info["tool_id"], + delta=delta.get("partial_json", ""), + ) + ] + else: + # Text block + token = delta.get("text", "") + if not token: + return [] + return [ + TextMessageContentEvent( + message_id=block_info["message_id"], + delta=token, + ) + ] + + def _handle_block_end(self, data: dict[str, Any]) -> list[Any]: + """Handle content_block:end — emit TEXT_MESSAGE_END or TOOL_CALL_END.""" + index = data.get("index", 0) + block_info = self._blocks.pop(index, None) + + if block_info is None: + return [] + + if block_info["type"] == "tool_use": + return [ToolCallEndEvent(tool_call_id=block_info["tool_id"])] + else: + return [TextMessageEndEvent(message_id=block_info["message_id"])] + + def _handle_tool_result(self, data: dict[str, Any]) -> list[Any]: + """Handle tool:result — emit TOOL_CALL_RESULT with optional MCP App metadata.""" + tool_call_id = data.get("tool_use_id", "") + content = data.get("content", "") + message_id = str(uuid.uuid4()) + + event = ToolCallResultEvent( + message_id=message_id, + tool_call_id=tool_call_id, + content=content if isinstance(content, str) else str(content), + role="tool", + ) + + # Conditionally pass through MCP App metadata for iframe rendering + extra: dict[str, Any] = {} + if "_meta" in data: + extra["_meta"] = data["_meta"] + if "structuredContent" in data: + extra["structuredContent"] = data["structuredContent"] + + if extra: + event.__pydantic_extra__ = extra + + return [event] + + def _handle_complete(self) -> list[Any]: + """Handle orchestrator:complete — emit RUN_FINISHED.""" + return [ + RunFinishedEvent( + thread_id=self._thread_id, + run_id=self._run_id, + ) + ] diff --git a/tests/test_agui_adapter.py b/tests/test_agui_adapter.py new file mode 100644 index 0000000..4614685 --- /dev/null +++ b/tests/test_agui_adapter.py @@ -0,0 +1,220 @@ +"""Tests for agui_adapter module — AmplifierdEventTranslator.""" + +import sys +import os + +# Add backend to sys.path so we can import agui_adapter +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "backend")) + +from agui_adapter import AmplifierdEventTranslator + + +class TestTextBlocks: + """Tests for text content_block events → AG-UI text message events.""" + + def setup_method(self): + self.translator = AmplifierdEventTranslator( + run_id="run_1", thread_id="thread_1" + ) + + def test_text_block_start(self): + """text_block_start → TEXT_MESSAGE_START with role='assistant' and messageId.""" + event = { + "event_type": "content_block:start", + "data": {"index": 0, "type": "text"}, + } + result = self.translator.translate(event) + assert len(result) == 1 + ev = result[0] + assert ev.type.value == "TEXT_MESSAGE_START" + assert ev.role == "assistant" + assert ev.message_id # must have a non-empty message_id + + def test_text_delta(self): + """text_delta → TEXT_MESSAGE_CONTENT with delta='Hello '.""" + # Start the text block first + start_event = { + "event_type": "content_block:start", + "data": {"index": 0, "type": "text"}, + } + self.translator.translate(start_event) + + delta_event = { + "event_type": "content_block:delta", + "data": { + "index": 0, + "delta": {"type": "text_delta", "text": "Hello "}, + }, + } + result = self.translator.translate(delta_event) + assert len(result) == 1 + ev = result[0] + assert ev.type.value == "TEXT_MESSAGE_CONTENT" + assert ev.delta == "Hello " + + def test_text_block_end(self): + """text_block_end → TEXT_MESSAGE_END.""" + # Start the block first + start_event = { + "event_type": "content_block:start", + "data": {"index": 0, "type": "text"}, + } + self.translator.translate(start_event) + + end_event = { + "event_type": "content_block:end", + "data": {"index": 0}, + } + result = self.translator.translate(end_event) + assert len(result) == 1 + ev = result[0] + assert ev.type.value == "TEXT_MESSAGE_END" + + +class TestToolBlocks: + """Tests for tool_use content_block events and tool results → AG-UI tool call events.""" + + def setup_method(self): + self.translator = AmplifierdEventTranslator( + run_id="run_1", thread_id="thread_1" + ) + + def test_tool_use_start(self): + """tool_use_start → TOOL_CALL_START with toolCallName='bash', toolCallId='tool_123'.""" + event = { + "event_type": "content_block:start", + "data": { + "index": 0, + "type": "tool_use", + "id": "tool_123", + "name": "bash", + }, + } + result = self.translator.translate(event) + assert len(result) == 1 + ev = result[0] + assert ev.type.value == "TOOL_CALL_START" + assert ev.tool_call_name == "bash" + assert ev.tool_call_id == "tool_123" + + def test_tool_use_delta(self): + """tool_use_delta → TOOL_CALL_ARGS with delta='{"command": "ls"}'.""" + # Start the tool block first + start_event = { + "event_type": "content_block:start", + "data": { + "index": 0, + "type": "tool_use", + "id": "tool_123", + "name": "bash", + }, + } + self.translator.translate(start_event) + + delta_event = { + "event_type": "content_block:delta", + "data": { + "index": 0, + "delta": { + "type": "input_json_delta", + "partial_json": '{"command": "ls"}', + }, + }, + } + result = self.translator.translate(delta_event) + assert len(result) == 1 + ev = result[0] + assert ev.type.value == "TOOL_CALL_ARGS" + assert ev.delta == '{"command": "ls"}' + + def test_tool_use_end(self): + """tool_use_end → TOOL_CALL_END.""" + # Start first + start_event = { + "event_type": "content_block:start", + "data": { + "index": 0, + "type": "tool_use", + "id": "tool_123", + "name": "bash", + }, + } + self.translator.translate(start_event) + + end_event = { + "event_type": "content_block:end", + "data": {"index": 0}, + } + result = self.translator.translate(end_event) + assert len(result) == 1 + ev = result[0] + assert ev.type.value == "TOOL_CALL_END" + + def test_tool_result(self): + """tool_result → TOOL_CALL_RESULT with toolCallId.""" + event = { + "event_type": "tool:result", + "data": {"tool_use_id": "tool_123", "content": "output"}, + } + result = self.translator.translate(event) + assert len(result) == 1 + ev = result[0] + assert ev.type.value == "TOOL_CALL_RESULT" + assert ev.tool_call_id == "tool_123" + + def test_tool_result_with_mcp_app_metadata(self): + """tool_result_with_mcp_app_metadata → TOOL_CALL_RESULT preserves _meta.ui.resourceUri='app://chart' and structuredContent.chartType='bar'.""" + event = { + "event_type": "tool:result", + "data": { + "tool_use_id": "tool_123", + "content": "output", + "_meta": {"ui": {"resourceUri": "app://chart"}}, + "structuredContent": {"chartType": "bar"}, + }, + } + result = self.translator.translate(event) + assert len(result) == 1 + ev = result[0] + assert ev.type.value == "TOOL_CALL_RESULT" + d = ev.model_dump() + assert d.get("_meta", {}).get("ui", {}).get("resourceUri") == "app://chart" + assert d.get("structuredContent", {}).get("chartType") == "bar" + + def test_tool_result_without_mcp_app_metadata(self): + """tool_result_without_mcp_app_metadata → no _meta or structuredContent keys.""" + event = { + "event_type": "tool:result", + "data": {"tool_use_id": "tool_123", "content": "output"}, + } + result = self.translator.translate(event) + assert len(result) == 1 + ev = result[0] + d = ev.model_dump() + assert "_meta" not in d + assert "structuredContent" not in d + + +class TestLifecycle: + """Tests for lifecycle events → AG-UI run lifecycle events.""" + + def setup_method(self): + self.translator = AmplifierdEventTranslator( + run_id="run_1", thread_id="thread_1" + ) + + def test_orchestrator_complete(self): + """orchestrator_complete → RUN_FINISHED with threadId and runId.""" + event = {"event_type": "orchestrator:complete", "data": {}} + result = self.translator.translate(event) + assert len(result) == 1 + ev = result[0] + assert ev.type.value == "RUN_FINISHED" + assert ev.thread_id == "thread_1" + assert ev.run_id == "run_1" + + def test_unknown_event(self): + """unknown event → empty list.""" + event = {"event_type": "some:unknown_event", "data": {}} + result = self.translator.translate(event) + assert result == []