research-workbench/backend/agui_adapter.py

"""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,
            )
        ]