236 lines
8.4 KiB
Python
236 lines
8.4 KiB
Python
"""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 loguru import logger
|
|
from pipecat.frames.frames import (
|
|
CancelFrame,
|
|
EndFrame,
|
|
InputTransportMessageFrame,
|
|
OutputTransportMessageUrgentFrame,
|
|
StopFrame,
|
|
)
|
|
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
|
|
|
|
|
|
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]]
|
|
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):
|
|
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
|
|
|
|
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]:
|
|
from uuid import uuid4
|
|
|
|
if self._closed_message is not None:
|
|
raise ClientToolError(self._closed_message)
|
|
|
|
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:
|
|
raise ClientToolError(
|
|
f"客户端工具调用发送失败: {function_name}"
|
|
) from exc
|
|
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,
|
|
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:
|
|
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")
|
|
if status == "ok":
|
|
response = {
|
|
"status": "ok",
|
|
"data": message.get("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"),
|
|
}
|
|
)
|
|
|
|
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()
|