ai-video-fullstack/backend/services/client_tools.py

"""Conversation-scoped bridge for tools implemented by the connected client."""

from __future__ import annotations

import asyncio
from collections import deque
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Any, Literal, Protocol
from uuid import uuid4

from loguru import logger
from pipecat.frames.frames import (
    CancelFrame,
    EndFrame,
    InputTransportMessageFrame,
    OutputTransportMessageUrgentFrame,
    StopFrame,
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.utils.time import time_now_iso8601


class ClientToolError(RuntimeError):
    """Raised when a client tool cannot be delivered or completed."""


ClientToolResponseWaitMode = Literal["timeout", "session"]


@dataclass(frozen=True)
class _PendingClientToolCall:
    future: asyncio.Future[dict[str, Any]]
    function_name: str
    interrupt_on_result: bool = False


@dataclass(frozen=True)
class _DeferredClientToolResult:
    """A successful result waiting for the media interruption barrier."""

    future: asyncio.Future[dict[str, Any]]
    response: dict[str, Any]


class ClientToolPort(Protocol):
    def supports(self, function_name: str) -> bool: ...

    async def call(
        self,
        function_name: str,
        arguments: dict[str, Any],
        *,
        timeout_seconds: float,
        wait_for_response: bool = True,
        response_wait_mode: ClientToolResponseWaitMode = "timeout",
        interrupt_on_result: bool = False,
    ) -> dict[str, Any]: ...


class ClientToolBroker(FrameProcessor):
    """Send client tool calls and correlate their app-message results."""

    def __init__(self) -> None:
        super().__init__()
        self._pending: dict[str, _PendingClientToolCall] = {}
        self._deferred_results: deque[_DeferredClientToolResult] = deque()
        self._interrupt_handler: Callable[[], Awaitable[None]] | None = None
        self._closed_message: str | None = None
        self._capabilities: frozenset[str] | None = None

    def set_capabilities(self, names: set[str] | None) -> None:
        self._capabilities = None if names is None else frozenset(names)

    def supports(self, function_name: str) -> bool:
        return self._capabilities is None or function_name in self._capabilities

    def set_interrupt_handler(
        self,
        handler: Callable[[], Awaitable[None]],
    ) -> None:
        """Use a processor downstream of muted user input as interrupt source."""
        self._interrupt_handler = handler

    async def call(
        self,
        function_name: str,
        arguments: dict[str, Any],
        *,
        timeout_seconds: float,
        wait_for_response: bool = True,
        response_wait_mode: ClientToolResponseWaitMode = "timeout",
        interrupt_on_result: bool = False,
    ) -> dict[str, Any]:
        if self._closed_message is not None:
            raise ClientToolError(self._closed_message)
        if not self.supports(function_name):
            raise ClientToolError(f"客户端未注册工具: {function_name}")

        tool_call_id = f"client_{uuid4().hex}"
        message = {
            "type": "client-tool-call",
            "tool_call_id": tool_call_id,
            "function_name": function_name,
            "arguments": dict(arguments),
            "wait_for_response": wait_for_response,
        }
        if not wait_for_response:
            try:
                await self.push_frame(
                    OutputTransportMessageUrgentFrame(message=message)
                )
            except Exception as exc:
                await self._emit_trace(
                    "client_tool_failed",
                    toolCallId=tool_call_id,
                    functionName=function_name,
                    error="客户端工具调用发送失败",
                )
                raise ClientToolError(
                    f"客户端工具调用发送失败: {function_name}"
                ) from exc
            await self._emit_trace(
                "client_tool_completed",
                toolCallId=tool_call_id,
                functionName=function_name,
                status="dispatched",
            )
            return {
                "status": "ok",
                "data": {"dispatched": True},
            }

        loop = asyncio.get_running_loop()
        future: asyncio.Future[dict[str, Any]] = loop.create_future()
        self._pending[tool_call_id] = _PendingClientToolCall(
            future=future,
            function_name=function_name,
            interrupt_on_result=interrupt_on_result,
        )
        try:
            await self.push_frame(
                OutputTransportMessageUrgentFrame(message=message)
            )
            if response_wait_mode == "session":
                return await future
            return await asyncio.wait_for(future, timeout=timeout_seconds)
        except TimeoutError as exc:
            await self._emit_trace(
                "client_tool_failed",
                toolCallId=tool_call_id,
                functionName=function_name,
                error="客户端工具调用超时",
            )
            raise ClientToolError(f"客户端工具调用超时: {function_name}") from exc
        except ClientToolError:
            raise
        except Exception as exc:
            raise ClientToolError(f"客户端工具调用失败: {function_name}") from exc
        finally:
            self._pending.pop(tool_call_id, None)

    async def process_frame(self, frame, direction: FrameDirection):
        await super().process_frame(frame, direction)

        if isinstance(frame, (EndFrame, CancelFrame)):
            self._close(
                "会话已结束" if isinstance(frame, EndFrame) else "会话已取消"
            )
            await self.push_frame(frame, direction)
            return

        if isinstance(frame, StopFrame):
            # StopFrame keeps processors reusable, so release only the calls
            # owned by the stopped run without permanently closing the broker.
            self._fail_pending("管线已停止")
            await self.push_frame(frame, direction)
            return

        if not isinstance(frame, InputTransportMessageFrame):
            await self.push_frame(frame, direction)
            return

        message = frame.message
        if not isinstance(message, dict) or message.get("type") != "client-tool-result":
            await self.push_frame(frame, direction)
            return

        tool_call_id = str(message.get("tool_call_id") or "")
        pending = self._pending.get(tool_call_id)
        if pending is None or pending.future.done():
            logger.debug(f"忽略未知或过期的客户端工具结果: {tool_call_id}")
            return

        future = pending.future
        status = str(message.get("status") or "error")
        result_data = message.get("data")
        user_action = (
            str(result_data.get("action") or "")
            if isinstance(result_data, dict)
            else ""
        )
        await self._emit_trace(
            "client_tool_completed" if status == "ok" else "client_tool_failed",
            toolCallId=tool_call_id,
            functionName=pending.function_name,
            status=status,
            internalStage=pending.function_name == "show_message",
            **({"userAction": user_action} if user_action else {}),
            **(
                {"error": str(message.get("message") or "客户端工具执行失败")}
                if status != "ok"
                else {}
            ),
        )
        if status == "ok":
            response = {
                "status": "ok",
                "data": result_data,
            }
            if pending.interrupt_on_result:
                # Match text input exactly: broadcasting only starts the
                # interruption. The assistant aggregator calls
                # on_interruption_processed() after the old response state has
                # actually been cleared; only then may the workflow continue.
                deferred = _DeferredClientToolResult(
                    future=future,
                    response=response,
                )
                self._deferred_results.append(deferred)
                try:
                    if self._interrupt_handler is not None:
                        await self._interrupt_handler()
                    else:
                        await self.broadcast_interruption()
                except Exception as exc:
                    try:
                        self._deferred_results.remove(deferred)
                    except ValueError:
                        # The downstream acknowledgement may have won the
                        # race before an upstream broadcast failure arrived.
                        pass
                    if not future.done():
                        future.set_exception(
                            ClientToolError("客户端工具结果中断处理失败")
                        )
                    logger.exception(f"客户端工具结果中断失败: {exc}")
                return
            future.set_result(response)
        else:
            future.set_result(
                {
                    "status": "error",
                    "message": str(message.get("message") or "客户端工具执行失败"),
                    "data": message.get("data"),
                }
            )

    async def _emit_trace(self, event: str, **details: Any) -> None:
        """Publish sanitized client interaction metadata for conversation history."""
        try:
            await self.push_frame(
                OutputTransportMessageUrgentFrame(
                    message={
                        "type": "workflow-event",
                        "eventId": f"wfe_{uuid4().hex[:20]}",
                        "event": event,
                        "timestamp": time_now_iso8601(),
                        **details,
                    }
                )
            )
        except Exception as exc:  # noqa: BLE001 - tracing must not alter the tool call
            logger.warning(f"发送客户端工具轨迹失败，不影响工具调用: {exc}")

    def on_interruption_processed(self) -> None:
        """Release one result after the aggregator acknowledges interruption."""
        while self._deferred_results:
            deferred = self._deferred_results.popleft()
            if deferred.future.done():
                continue
            deferred.future.set_result(deferred.response)
            logger.debug("客户端工具结果已通过中断处理边界")
            return

    def _fail_pending(self, message: str) -> None:
        self._deferred_results.clear()
        for pending in self._pending.values():
            future = pending.future
            if not future.done():
                future.set_exception(ClientToolError(message))
        self._pending.clear()

    def _close(self, message: str) -> None:
        """Close once and reject both current and future client-tool calls."""
        if self._closed_message is None:
            self._closed_message = message
        self._fail_pending(self._closed_message)

    async def cleanup(self) -> None:
        self._close("客户端工具通道已关闭")
        await super().cleanup()