"""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), ""