336 lines
11 KiB
Python
336 lines
11 KiB
Python
"""Provider-neutral function calling for speech-to-speech sessions."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from collections import deque
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Literal
|
|
|
|
from loguru import logger
|
|
|
|
|
|
RealtimeToolFollowUp = Literal["continue", "none", "suppress"]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RealtimeTool:
|
|
"""Provider-neutral JSON-schema function definition."""
|
|
|
|
name: str
|
|
description: str
|
|
properties: dict[str, Any] = field(default_factory=dict)
|
|
required: tuple[str, ...] = ()
|
|
|
|
def openai_schema(self) -> dict[str, Any]:
|
|
"""Serialize the flat schema used by OpenAI-compatible providers."""
|
|
|
|
return {
|
|
"type": "function",
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": self.properties,
|
|
"required": list(self.required),
|
|
},
|
|
}
|
|
|
|
def qwen_schema(self) -> dict[str, Any]:
|
|
"""Serialize Qwen-Audio's documented nested function schema."""
|
|
|
|
flat = self.openai_schema()
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": flat["name"],
|
|
"description": flat["description"],
|
|
"parameters": flat["parameters"],
|
|
},
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RealtimeToolResult:
|
|
"""Tool output and the action taken after every call in one response."""
|
|
|
|
output: dict[str, Any]
|
|
follow_up: RealtimeToolFollowUp = "continue"
|
|
after_output: Callable[[], Awaitable[None]] | None = None
|
|
|
|
|
|
RealtimeToolDispatcher = Callable[
|
|
[str, dict[str, Any], str], Awaitable[RealtimeToolResult]
|
|
]
|
|
SendProviderEvent = Callable[[dict[str, Any]], Awaitable[None]]
|
|
SuppressProviderResponse = Callable[[], Awaitable[None]]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _PendingCall:
|
|
name: str
|
|
call_id: str
|
|
arguments: dict[str, Any]
|
|
output_index: int
|
|
argument_error: str = ""
|
|
|
|
|
|
@dataclass
|
|
class _ResponseCalls:
|
|
calls: dict[str, _PendingCall] = field(default_factory=dict)
|
|
cancelled: bool = False
|
|
processing: bool = False
|
|
|
|
|
|
class RealtimeToolSession:
|
|
"""Coordinate all function calls emitted by one provider response.
|
|
|
|
Qwen may emit more than one function call in a response. Calls are collected
|
|
until ``response.done``, executed in output order, and followed by at most
|
|
one new model response. Duplicate argument/output-item events are ignored by
|
|
call id.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
send_event: SendProviderEvent,
|
|
*,
|
|
suppress_response: SuppressProviderResponse | None = None,
|
|
) -> None:
|
|
self._send_event = send_event
|
|
self._suppress_response = suppress_response
|
|
self._dispatcher: RealtimeToolDispatcher | None = None
|
|
self._handled_call_ids: set[str] = set()
|
|
self._handled_call_order: deque[str] = deque()
|
|
self._finished_response_ids: set[str] = set()
|
|
self._finished_response_order: deque[str] = deque()
|
|
self._responses: dict[str, _ResponseCalls] = {}
|
|
self._tasks: set[asyncio.Task[None]] = set()
|
|
|
|
def set_dispatcher(self, dispatcher: RealtimeToolDispatcher | None) -> None:
|
|
self._dispatcher = dispatcher
|
|
|
|
def submit_call(
|
|
self,
|
|
*,
|
|
response_id: str,
|
|
name: str,
|
|
call_id: str,
|
|
arguments: str | dict[str, Any] | None,
|
|
output_index: int = 0,
|
|
) -> None:
|
|
"""Register one complete call without yielding to the receive loop."""
|
|
|
|
if not call_id:
|
|
logger.warning("Realtime function call 缺少 call_id，已忽略")
|
|
return
|
|
if call_id in self._handled_call_ids:
|
|
return
|
|
self._handled_call_ids.add(call_id)
|
|
self._handled_call_order.append(call_id)
|
|
if len(self._handled_call_order) > 4096:
|
|
expired_call_id = self._handled_call_order.popleft()
|
|
self._handled_call_ids.discard(expired_call_id)
|
|
group_id = response_id or f"response_{call_id}"
|
|
if group_id in self._finished_response_ids:
|
|
return
|
|
group = self._responses.setdefault(group_id, _ResponseCalls())
|
|
if group.cancelled or group.processing:
|
|
return
|
|
parsed_arguments, argument_error = self._parse_arguments(arguments)
|
|
group.calls[call_id] = _PendingCall(
|
|
name=name,
|
|
call_id=call_id,
|
|
arguments=parsed_arguments,
|
|
output_index=output_index,
|
|
argument_error=argument_error,
|
|
)
|
|
|
|
def finish_response(self, response_id: str, *, cancelled: bool = False) -> None:
|
|
"""Seal one response and process its calls in a background task."""
|
|
|
|
if not response_id:
|
|
if len(self._responses) != 1:
|
|
return
|
|
response_id = next(iter(self._responses))
|
|
self._remember_finished_response(response_id)
|
|
group = self._responses.get(response_id)
|
|
if group is None:
|
|
return
|
|
group.cancelled = group.cancelled or cancelled
|
|
if group.processing:
|
|
return
|
|
group.processing = True
|
|
task = asyncio.create_task(
|
|
self._process_response(response_id, group),
|
|
name=f"realtime_tools_{response_id}",
|
|
)
|
|
self._tasks.add(task)
|
|
task.add_done_callback(self._tasks.discard)
|
|
|
|
async def wait_pending(self) -> None:
|
|
"""Wait for current batches; primarily useful for orderly shutdown/tests."""
|
|
|
|
tasks = tuple(self._tasks)
|
|
if tasks:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
def cancel_response(self, response_id: str) -> None:
|
|
"""Prevent unstarted calls and any follow-up for a stale response."""
|
|
|
|
group = self._responses.get(response_id)
|
|
if group is not None:
|
|
group.cancelled = True
|
|
|
|
def invalidate_pending_follow_ups(self) -> None:
|
|
"""Make in-flight tool groups stale without cancelling active I/O."""
|
|
|
|
for group in self._responses.values():
|
|
group.cancelled = True
|
|
|
|
def cancel_all(self) -> None:
|
|
"""Cancel pending work when the provider connection closes."""
|
|
|
|
for group in self._responses.values():
|
|
group.cancelled = True
|
|
for task in tuple(self._tasks):
|
|
task.cancel()
|
|
self._responses.clear()
|
|
|
|
async def _process_response(
|
|
self,
|
|
response_id: str,
|
|
group: _ResponseCalls,
|
|
) -> None:
|
|
try:
|
|
results: list[RealtimeToolResult] = []
|
|
calls = sorted(
|
|
group.calls.values(),
|
|
key=lambda item: (item.output_index, item.call_id),
|
|
)
|
|
for call in calls:
|
|
if group.cancelled:
|
|
await self._send_output(
|
|
call.call_id,
|
|
{"status": "cancelled", "message": "调用已被新的用户输入取代"},
|
|
)
|
|
continue
|
|
result = await self._dispatch(call)
|
|
results.append(result)
|
|
await self._send_output(call.call_id, result.output)
|
|
|
|
if group.cancelled:
|
|
return
|
|
for result in results:
|
|
if result.after_output is not None:
|
|
try:
|
|
await result.after_output()
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc: # noqa: BLE001 - isolate callbacks
|
|
logger.exception(f"Realtime 工具结果回调失败:{exc}")
|
|
follow_up = self._combined_follow_up(results)
|
|
if follow_up == "suppress" and self._suppress_response is not None:
|
|
await self._suppress_response()
|
|
elif follow_up == "continue":
|
|
await self._send_event({"type": "response.create"})
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc: # noqa: BLE001 - background provider task
|
|
logger.exception(f"Realtime response 工具处理失败:{exc}")
|
|
finally:
|
|
self._responses.pop(response_id, None)
|
|
|
|
async def _send_output(self, call_id: str, output: dict[str, Any]) -> None:
|
|
await self._send_event(
|
|
{
|
|
"type": "conversation.item.create",
|
|
"item": {
|
|
"type": "function_call_output",
|
|
"call_id": call_id,
|
|
"output": json.dumps(
|
|
output,
|
|
ensure_ascii=False,
|
|
separators=(",", ":"),
|
|
),
|
|
},
|
|
}
|
|
)
|
|
|
|
async def _dispatch(self, call: _PendingCall) -> RealtimeToolResult:
|
|
if call.argument_error:
|
|
return RealtimeToolResult(
|
|
{"status": "error", "message": call.argument_error}
|
|
)
|
|
if self._dispatcher is None:
|
|
return RealtimeToolResult(
|
|
{"status": "error", "message": "当前会话未注册工具处理器"}
|
|
)
|
|
try:
|
|
return await self._dispatcher(
|
|
call.name,
|
|
call.arguments,
|
|
call.call_id,
|
|
)
|
|
except Exception as exc: # noqa: BLE001 - return stable errors to provider
|
|
logger.exception(f"Realtime 工具 {call.name} 执行失败:{exc}")
|
|
return RealtimeToolResult(
|
|
{
|
|
"status": "error",
|
|
"message": f"工具执行失败:{type(exc).__name__}",
|
|
}
|
|
)
|
|
|
|
@staticmethod
|
|
def _combined_follow_up(
|
|
results: list[RealtimeToolResult],
|
|
) -> RealtimeToolFollowUp:
|
|
actions = {result.follow_up for result in results}
|
|
if "suppress" in actions:
|
|
return "suppress"
|
|
if "continue" in actions:
|
|
return "continue"
|
|
return "none"
|
|
|
|
def clear(self) -> None:
|
|
self.cancel_all()
|
|
self._handled_call_ids.clear()
|
|
self._handled_call_order.clear()
|
|
self._finished_response_ids.clear()
|
|
self._finished_response_order.clear()
|
|
|
|
async def shutdown(self) -> None:
|
|
"""Cancel and join background calls before releasing a connection."""
|
|
|
|
tasks = tuple(self._tasks)
|
|
self.clear()
|
|
if tasks:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
def _remember_finished_response(self, response_id: str) -> None:
|
|
if response_id in self._finished_response_ids:
|
|
return
|
|
self._finished_response_ids.add(response_id)
|
|
self._finished_response_order.append(response_id)
|
|
if len(self._finished_response_order) > 512:
|
|
expired = self._finished_response_order.popleft()
|
|
self._finished_response_ids.discard(expired)
|
|
|
|
@staticmethod
|
|
def _parse_arguments(
|
|
arguments: str | dict[str, Any] | None,
|
|
) -> tuple[dict[str, Any], str]:
|
|
if isinstance(arguments, dict):
|
|
return dict(arguments), ""
|
|
if arguments is None or arguments == "":
|
|
return {}, ""
|
|
try:
|
|
parsed = json.loads(arguments)
|
|
except (TypeError, json.JSONDecodeError):
|
|
return {}, "工具参数不是有效的 JSON 对象"
|
|
if not isinstance(parsed, dict):
|
|
return {}, "工具参数必须是 JSON 对象"
|
|
return dict(parsed), ""
|