137 lines
4.4 KiB
Python
137 lines
4.4 KiB
Python
"""Conversation-scoped bridge for tools implemented by the connected client."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any, Protocol
|
|
|
|
from loguru import logger
|
|
from pipecat.frames.frames import (
|
|
EndFrame,
|
|
InputTransportMessageFrame,
|
|
OutputTransportMessageUrgentFrame,
|
|
)
|
|
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
|
|
|
|
|
|
class ClientToolError(RuntimeError):
|
|
"""Raised when a client tool cannot be delivered or completed."""
|
|
|
|
|
|
class ClientToolPort(Protocol):
|
|
async def call(
|
|
self,
|
|
function_name: str,
|
|
arguments: dict[str, Any],
|
|
*,
|
|
timeout_seconds: float,
|
|
wait_for_response: bool = True,
|
|
) -> 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, asyncio.Future[dict[str, Any]]] = {}
|
|
|
|
async def call(
|
|
self,
|
|
function_name: str,
|
|
arguments: dict[str, Any],
|
|
*,
|
|
timeout_seconds: float,
|
|
wait_for_response: bool = True,
|
|
) -> dict[str, Any]:
|
|
from uuid import uuid4
|
|
|
|
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] = future
|
|
try:
|
|
await self.push_frame(
|
|
OutputTransportMessageUrgentFrame(message=message)
|
|
)
|
|
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):
|
|
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 "")
|
|
future = self._pending.get(tool_call_id)
|
|
if future is None or future.done():
|
|
logger.debug(f"忽略未知或过期的客户端工具结果: {tool_call_id}")
|
|
return
|
|
|
|
status = str(message.get("status") or "error")
|
|
if status == "ok":
|
|
future.set_result(
|
|
{
|
|
"status": "ok",
|
|
"data": message.get("data"),
|
|
}
|
|
)
|
|
else:
|
|
future.set_result(
|
|
{
|
|
"status": "error",
|
|
"message": str(message.get("message") or "客户端工具执行失败"),
|
|
"data": message.get("data"),
|
|
}
|
|
)
|
|
|
|
def _fail_pending(self, message: str) -> None:
|
|
for future in self._pending.values():
|
|
if not future.done():
|
|
future.set_exception(ClientToolError(message))
|
|
self._pending.clear()
|
|
|
|
async def cleanup(self) -> None:
|
|
self._fail_pending("客户端工具通道已关闭")
|
|
await super().cleanup()
|