agui_protocol.py

"""AG-UI (Agent-User Interaction Protocol) 协议实现

AG-UI 是一种开源、轻量级、基于事件的协议，用于标准化 AI Agent 与前端应用之间的交互。
参考: https://docs.ag-ui.com/

本实现使用 ag-ui-protocol 包提供的事件类型和编码器，
将 AgentResult 事件转换为 AG-UI SSE 格式。
"""

from dataclasses import dataclass, field
from typing import (
    Any,
    AsyncIterator,
    Dict,
    Iterator,
    List,
    Optional,
    TYPE_CHECKING,
)
import uuid

if TYPE_CHECKING:
    from ag_ui.core import (
        Message as AguiMessage,
    )
    from ag_ui.encoder import EventEncoder

from fastapi import APIRouter, Request
from fastapi.responses import StreamingResponse
import pydash

from ..utils.helper import merge, MergeOptions
from .model import (
    AgentEvent,
    AgentRequest,
    EventType,
    Message,
    MessageRole,
    ServerConfig,
    Tool,
    ToolCall,
)
from .protocol import BaseProtocolHandler

if TYPE_CHECKING:
    from .invoker import AgentInvoker


# ============================================================================
# AG-UI 协议处理器
# ============================================================================

DEFAULT_PREFIX = "/ag-ui/agent"


@dataclass
class TextState:
    started: bool = False
    ended: bool = False
    message_id: str = field(default_factory=lambda: str(uuid.uuid4()))


@dataclass
class ToolCallState:
    name: str = ""
    started: bool = False
    ended: bool = False
    has_result: bool = False
    is_hitl: bool = False


@dataclass
class StreamStateMachine:
    text: TextState = field(default_factory=TextState)
    tool_call_states: Dict[str, ToolCallState] = field(default_factory=dict)
    tool_result_chunks: Dict[str, List[str]] = field(default_factory=dict)
    run_errored: bool = False

    def end_all_tools(
        self, encoder: "EventEncoder", exclude: Optional[str] = None
    ) -> Iterator[str]:
        from ag_ui.core import ToolCallEndEvent

        for tool_id, state in self.tool_call_states.items():
            if exclude and tool_id == exclude:
                continue
            if state.started and not state.ended:
                yield encoder.encode(ToolCallEndEvent(tool_call_id=tool_id))
                state.ended = True

    def ensure_text_started(self, encoder: "EventEncoder") -> Iterator[str]:
        from ag_ui.core import TextMessageStartEvent

        if not self.text.started or self.text.ended:
            if self.text.ended:
                self.text = TextState()
            yield encoder.encode(
                TextMessageStartEvent(
                    message_id=self.text.message_id,
                    role="assistant",
                )
            )
            self.text.started = True
            self.text.ended = False

    def end_text_if_open(self, encoder: "EventEncoder") -> Iterator[str]:
        from ag_ui.core import TextMessageEndEvent

        if self.text.started and not self.text.ended:
            yield encoder.encode(
                TextMessageEndEvent(message_id=self.text.message_id)
            )
            self.text.ended = True

    def cache_tool_result_chunk(self, tool_id: str, delta: str) -> None:
        if not tool_id or delta is None:
            return
        if delta:
            self.tool_result_chunks.setdefault(tool_id, []).append(delta)

    def pop_tool_result_chunks(self, tool_id: str) -> str:
        return "".join(self.tool_result_chunks.pop(tool_id, []))


class AGUIProtocolHandler(BaseProtocolHandler):
    """AG-UI 协议处理器

    实现 AG-UI (Agent-User Interaction Protocol) 兼容接口。
    参考: https://docs.ag-ui.com/

    使用 ag-ui-protocol 包提供的事件类型和编码器。

    特点:
    - 基于事件的流式通信
    - 完整支持所有 AG-UI 事件类型
    - 支持状态同步
    - 支持工具调用

    Example:
        >>> from agentrun.server import AgentRunServer, AGUIProtocolHandler
        >>>
        >>> server = AgentRunServer(
        ...     invoke_agent=my_agent,
        ...     protocols=[AGUIProtocolHandler()]
        ... )
        >>> server.start(port=8000)
        # 可访问: POST http://localhost:8000/ag-ui/agent
    """

    name = "ag-ui"

    def __init__(self, config: Optional[ServerConfig] = None):
        from ag_ui.encoder import EventEncoder

        self._config = config.agui if config else None
        self._encoder = EventEncoder()

    def get_prefix(self) -> str:
        """AG-UI 协议建议使用 /ag-ui/agent 前缀"""
        return pydash.get(self._config, "prefix", DEFAULT_PREFIX)

    def as_fastapi_router(self, agent_invoker: "AgentInvoker") -> APIRouter:
        """创建 AG-UI 协议的 FastAPI Router"""
        router = APIRouter()

        @router.post("")
        async def run_agent(request: Request):
            """AG-UI 运行 Agent 端点

            接收 AG-UI 格式的请求，返回 SSE 事件流。
            """
            sse_headers = {
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
            }

            try:
                request_data = await request.json()
                agent_request, context = await self.parse_request(
                    request, request_data
                )

                # 使用 invoke_stream 获取流式结果
                event_stream = self._format_stream(
                    agent_invoker.invoke_stream(agent_request),
                    context,
                )

                return StreamingResponse(
                    event_stream,
                    media_type=self._encoder.get_content_type(),
                    headers=sse_headers,
                )

            except ValueError as e:
                return StreamingResponse(
                    self._error_stream(str(e)),
                    media_type=self._encoder.get_content_type(),
                    headers=sse_headers,
                )
            except Exception as e:
                return StreamingResponse(
                    self._error_stream(f"Internal error: {str(e)}"),
                    media_type=self._encoder.get_content_type(),
                    headers=sse_headers,
                )

        @router.get("/health")
        async def health_check():
            """健康检查端点"""
            return {"status": "ok", "protocol": "ag-ui", "version": "1.0"}

        return router

    async def parse_request(
        self,
        request: Request,
        request_data: Dict[str, Any],
    ) -> tuple[AgentRequest, Dict[str, Any]]:
        """解析 AG-UI 格式的请求

        Args:
            request: FastAPI Request 对象
            request_data: HTTP 请求体 JSON 数据

        Returns:
            tuple: (AgentRequest, context)
        """
        # 创建上下文
        context = {
            "thread_id": request_data.get("threadId") or str(uuid.uuid4()),
            "run_id": request_data.get("runId") or str(uuid.uuid4()),
        }

        # 解析消息列表
        messages = self._parse_messages(request_data.get("messages", []))

        # 解析工具列表
        tools = self._parse_tools(request_data.get("tools"))

        # 构建 AgentRequest
        agent_request = AgentRequest(
            protocol="agui",  # 设置协议名称
            messages=messages,
            stream=True,  # AG-UI 总是流式
            tools=tools,
            raw_request=request,  # 保留原始请求对象
        )

        return agent_request, context

    def _parse_messages(
        self, raw_messages: List[Dict[str, Any]]
    ) -> List[Message]:
        """解析消息列表

        Args:
            raw_messages: 原始消息数据

        Returns:
            标准化的消息列表
        """
        messages = []

        for msg_data in raw_messages:
            if not isinstance(msg_data, dict):
                continue

            role_str = msg_data.get("role", "user")
            try:
                role = MessageRole(role_str)
            except ValueError:
                role = MessageRole.USER

            # 解析 tool_calls
            tool_calls = None
            if msg_data.get("toolCalls"):
                tool_calls = [
                    ToolCall(
                        id=tc.get("id", ""),
                        type=tc.get("type", "function"),
                        function=tc.get("function", {}),
                    )
                    for tc in msg_data["toolCalls"]
                ]

            messages.append(
                Message(
                    id=msg_data.get("id"),
                    role=role,
                    content=msg_data.get("content"),
                    name=msg_data.get("name"),
                    tool_calls=tool_calls,
                    tool_call_id=msg_data.get("toolCallId"),
                )
            )

        return messages

    def _parse_tools(
        self, raw_tools: Optional[List[Dict[str, Any]]]
    ) -> Optional[List[Tool]]:
        """解析工具列表

        Args:
            raw_tools: 原始工具数据

        Returns:
            标准化的工具列表
        """
        if not raw_tools:
            return None

        tools = []
        for tool_data in raw_tools:
            if not isinstance(tool_data, dict):
                continue

            tools.append(
                Tool(
                    type=tool_data.get("type", "function"),
                    function=tool_data.get("function", {}),
                )
            )

        return tools if tools else None

    async def _format_stream(
        self,
        event_stream: AsyncIterator[AgentEvent],
        context: Dict[str, Any],
    ) -> AsyncIterator[str]:
        """将 AgentEvent 流转换为 AG-UI SSE 格式

        自动生成边界事件：
        - RUN_STARTED / RUN_FINISHED（生命周期）
        - TEXT_MESSAGE_START / TEXT_MESSAGE_END（文本边界）
        - TOOL_CALL_START / TOOL_CALL_END（工具调用边界）

        注意：RUN_ERROR 之后不能再发送任何事件（包括 RUN_FINISHED）

        Args:
            event_stream: AgentEvent 流
            context: 上下文信息

        Yields:
            SSE 格式的字符串
        """
        from ag_ui.core import RunFinishedEvent, RunStartedEvent

        state = StreamStateMachine()

        # 发送 RUN_STARTED
        yield self._encoder.encode(
            RunStartedEvent(
                thread_id=context.get("thread_id"),
                run_id=context.get("run_id"),
            )
        )

        async for event in event_stream:
            # RUN_ERROR 后不再处理任何事件
            if state.run_errored:
                continue

            # 检查是否是错误事件
            if event.event == EventType.ERROR:
                state.run_errored = True

            # 处理边界事件注入
            for sse_data in self._process_event_with_boundaries(
                event,
                context,
                state,
            ):
                if sse_data:
                    yield sse_data

        # RUN_ERROR 后不发送任何清理事件
        if state.run_errored:
            return

        # 结束所有未结束的工具调用
        for sse_data in state.end_all_tools(self._encoder):
            yield sse_data

        # 发送 TEXT_MESSAGE_END（如果有文本消息且未结束）
        for sse_data in state.end_text_if_open(self._encoder):
            yield sse_data

        # 发送 RUN_FINISHED
        yield self._encoder.encode(
            RunFinishedEvent(
                thread_id=context.get("thread_id"),
                run_id=context.get("run_id"),
            )
        )

    def _process_event_with_boundaries(
        self,
        event: AgentEvent,
        context: Dict[str, Any],
        state: StreamStateMachine,
    ) -> Iterator[str]:
        """处理事件并注入边界事件"""
        import json

        from ag_ui.core import CustomEvent as AguiCustomEvent
        from ag_ui.core import (
            RunErrorEvent,
            StateDeltaEvent,
            StateSnapshotEvent,
            TextMessageContentEvent,
            ToolCallArgsEvent,
            ToolCallEndEvent,
            ToolCallResultEvent,
            ToolCallStartEvent,
        )

        # RAW 事件直接透传
        if event.event == EventType.RAW:
            raw_data = event.data.get("raw", "")
            if raw_data:
                if not raw_data.endswith("\n\n"):
                    raw_data = raw_data.rstrip("\n") + "\n\n"
                yield raw_data
            return

        # TEXT 事件：在首个 TEXT 前注入 TEXT_MESSAGE_START
        # AG-UI 协议要求：发送 TEXT_MESSAGE_START 前必须先结束所有未结束的 TOOL_CALL
        if event.event == EventType.TEXT:
            for sse_data in state.end_all_tools(self._encoder):
                yield sse_data

            for sse_data in state.ensure_text_started(self._encoder):
                yield sse_data

            agui_event = TextMessageContentEvent(
                message_id=state.text.message_id,
                delta=event.data.get("delta", ""),
            )
            if event.addition:
                event_dict = agui_event.model_dump(
                    by_alias=True, exclude_none=True
                )
                event_dict = self._apply_addition(
                    event_dict,
                    event.addition,
                    event.addition_merge_options,
                )
                json_str = json.dumps(event_dict, ensure_ascii=False)
                yield f"data: {json_str}\n\n"
            else:
                yield self._encoder.encode(agui_event)
            return

        # TOOL_CALL_CHUNK 事件：在首个 CHUNK 前注入 TOOL_CALL_START
        if event.event == EventType.TOOL_CALL_CHUNK:
            tool_id = event.data.get("id", "")
            tool_name = event.data.get("name", "")

            for sse_data in state.end_text_if_open(self._encoder):
                yield sse_data

            need_start = False
            current_state = state.tool_call_states.get(tool_id)
            if tool_id:
                if current_state is None or current_state.ended:
                    need_start = True

            if need_start:
                yield self._encoder.encode(
                    ToolCallStartEvent(
                        tool_call_id=tool_id,
                        tool_call_name=tool_name,
                    )
                )
                state.tool_call_states[tool_id] = ToolCallState(
                    name=tool_name,
                    started=True,
                    ended=False,
                )

            yield self._encoder.encode(
                ToolCallArgsEvent(
                    tool_call_id=tool_id,
                    delta=event.data.get("args_delta", ""),
                )
            )
            return

        # TOOL_CALL 事件：完整的工具调用事件
        if event.event == EventType.TOOL_CALL:
            tool_id = event.data.get("id", "")
            tool_name = event.data.get("name", "")
            tool_args = event.data.get("args", "")

            for sse_data in state.end_text_if_open(self._encoder):
                yield sse_data

            need_start = False
            current_state = state.tool_call_states.get(tool_id)
            if tool_id:
                if current_state is None or current_state.ended:
                    need_start = True

            if need_start:
                yield self._encoder.encode(
                    ToolCallStartEvent(
                        tool_call_id=tool_id,
                        tool_call_name=tool_name,
                    )
                )
                state.tool_call_states[tool_id] = ToolCallState(
                    name=tool_name,
                    started=True,
                    ended=False,
                )

            # 发送工具参数（如果存在）
            if tool_args:
                yield self._encoder.encode(
                    ToolCallArgsEvent(
                        tool_call_id=tool_id,
                        delta=tool_args,
                    )
                )
            return

        # TOOL_RESULT_CHUNK 事件：工具执行过程中的流式输出
        if event.event == EventType.TOOL_RESULT_CHUNK:
            tool_id = event.data.get("id", "")
            delta = event.data.get("delta", "")
            state.cache_tool_result_chunk(tool_id, delta)
            return

        # HITL 事件：请求人类介入
        if event.event == EventType.HITL:
            hitl_id = event.data.get("id", "")
            tool_call_id = event.data.get("tool_call_id", "")
            hitl_type = event.data.get("type", "confirmation")
            prompt = event.data.get("prompt", "")
            options = event.data.get("options")
            default = event.data.get("default")
            timeout = event.data.get("timeout")
            schema = event.data.get("schema")

            for sse_data in state.end_text_if_open(self._encoder):
                yield sse_data

            if tool_call_id and tool_call_id in state.tool_call_states:
                tool_state = state.tool_call_states[tool_call_id]
                if tool_state.started and not tool_state.ended:
                    yield self._encoder.encode(
                        ToolCallEndEvent(tool_call_id=tool_call_id)
                    )
                    tool_state.ended = True
                tool_state.is_hitl = True
                tool_state.has_result = False
                return

            import json as json_module

            args_dict: Dict[str, Any] = {
                "type": hitl_type,
                "prompt": prompt,
            }
            if options:
                args_dict["options"] = options
            if default is not None:
                args_dict["default"] = default
            if timeout is not None:
                args_dict["timeout"] = timeout
            if schema:
                args_dict["schema"] = schema

            args_json = json_module.dumps(args_dict, ensure_ascii=False)
            actual_id = tool_call_id or hitl_id

            yield self._encoder.encode(
                ToolCallStartEvent(
                    tool_call_id=actual_id,
                    tool_call_name=f"hitl_{hitl_type}",
                )
            )
            yield self._encoder.encode(
                ToolCallArgsEvent(
                    tool_call_id=actual_id,
                    delta=args_json,
                )
            )
            yield self._encoder.encode(ToolCallEndEvent(tool_call_id=actual_id))

            state.tool_call_states[actual_id] = ToolCallState(
                name=f"hitl_{hitl_type}",
                started=True,
                ended=True,
                has_result=False,
                is_hitl=True,
            )
            return

        # TOOL_RESULT 事件：确保当前工具调用已结束
        if event.event == EventType.TOOL_RESULT:
            tool_id = event.data.get("id", "")
            tool_name = event.data.get("name", "")

            for sse_data in state.end_text_if_open(self._encoder):
                yield sse_data

            tool_state = (
                state.tool_call_states.get(tool_id) if tool_id else None
            )
            if tool_id and tool_state is None:
                yield self._encoder.encode(
                    ToolCallStartEvent(
                        tool_call_id=tool_id,
                        tool_call_name=tool_name or "",
                    )
                )
                tool_state = ToolCallState(
                    name=tool_name, started=True, ended=False
                )
                state.tool_call_states[tool_id] = tool_state

            if tool_state and tool_state.started and not tool_state.ended:
                yield self._encoder.encode(
                    ToolCallEndEvent(tool_call_id=tool_id)
                )
                tool_state.ended = True

            final_result = event.data.get("content") or event.data.get(
                "result", ""
            )
            if tool_id:
                cached_chunks = state.pop_tool_result_chunks(tool_id)
                if cached_chunks:
                    final_result = cached_chunks + final_result

            yield self._encoder.encode(
                ToolCallResultEvent(
                    message_id=event.data.get(
                        "message_id", f"tool-result-{tool_id}"
                    ),
                    tool_call_id=tool_id,
                    content=final_result,
                    role="tool",
                )
            )
            return

        # ERROR 事件
        if event.event == EventType.ERROR:
            yield self._encoder.encode(
                RunErrorEvent(
                    message=event.data.get("message", ""),
                    code=event.data.get("code"),
                )
            )
            return

        # STATE 事件
        if event.event == EventType.STATE:
            if "snapshot" in event.data:
                yield self._encoder.encode(
                    StateSnapshotEvent(snapshot=event.data.get("snapshot", {}))
                )
            elif "delta" in event.data:
                yield self._encoder.encode(
                    StateDeltaEvent(delta=event.data.get("delta", []))
                )
            else:
                yield self._encoder.encode(
                    StateSnapshotEvent(snapshot=event.data)
                )
            return

        # CUSTOM 事件
        if event.event == EventType.CUSTOM:
            yield self._encoder.encode(
                AguiCustomEvent(
                    name=event.data.get("name", "custom"),
                    value=event.data.get("value"),
                )
            )
            return

        # 其他未知事件
        event_name = (
            event.event.value
            if hasattr(event.event, "value")
            else str(event.event)
        )
        yield self._encoder.encode(
            AguiCustomEvent(
                name=event_name,
                value=event.data,
            )
        )

    def _convert_messages_for_snapshot(
        self, messages: List[Dict[str, Any]]
    ) -> List["AguiMessage"]:
        """将消息列表转换为 ag-ui-protocol 格式

        Args:
            messages: 消息字典列表

        Returns:
            ag-ui-protocol 消息列表
        """
        from ag_ui.core import AssistantMessage, SystemMessage
        from ag_ui.core import ToolMessage as AguiToolMessage
        from ag_ui.core import UserMessage

        result = []
        for msg in messages:
            if not isinstance(msg, dict):
                continue

            role = msg.get("role", "user")
            content = msg.get("content", "")
            msg_id = msg.get("id", str(uuid.uuid4()))

            if role == "user":
                result.append(
                    UserMessage(id=msg_id, role="user", content=content)
                )
            elif role == "assistant":
                result.append(
                    AssistantMessage(
                        id=msg_id,
                        role="assistant",
                        content=content,
                    )
                )
            elif role == "system":
                result.append(
                    SystemMessage(id=msg_id, role="system", content=content)
                )
            elif role == "tool":
                result.append(
                    AguiToolMessage(
                        id=msg_id,
                        role="tool",
                        content=content,
                        tool_call_id=msg.get("tool_call_id", ""),
                    )
                )

        return result

    def _apply_addition(
        self,
        event_data: Dict[str, Any],
        addition: Optional[Dict[str, Any]],
        merge_options: Optional[MergeOptions] = None,
    ) -> Dict[str, Any]:
        """应用 addition 字段

        Args:
            event_data: 原始事件数据
            addition: 附加字段
            merge_options: 合并选项，透传给 utils.helper.merge

        Returns:
            合并后的事件数据
        """
        if not addition:
            return event_data

        return merge(event_data, addition, **(merge_options or {}))

    async def _error_stream(self, message: str) -> AsyncIterator[str]:
        """生成错误事件流

        Args:
            message: 错误消息

        Yields:
            SSE 格式的错误事件
        """
        from ag_ui.core import RunErrorEvent, RunStartedEvent

        thread_id = str(uuid.uuid4())
        run_id = str(uuid.uuid4())

        # 生命周期开始
        yield self._encoder.encode(
            RunStartedEvent(thread_id=thread_id, run_id=run_id)
        )

        # 错误事件
        yield self._encoder.encode(
            RunErrorEvent(message=message, code="REQUEST_ERROR")
        )