Function Calling

import os
import uuid
import threading
import traceback
import json
import base64
import signal
import sys
import time
from typing import Dict, Any, Optional, List
import pyaudio
import queue
import contextlib
import dashscope
from dashscope.audio.qwen_omni import *

# ==================== 常量定義 ====================
VOICE = 'Tina'
MODEL = "qwen3.5-omni-plus-realtime"
# 如果需要訪問北京地區，請WS_URL將替換為：wss://{WorkspaceId}.cn-beijing.maas.aliyuncs.com/api-ws/v1/realtime
WS_URL = "wss://{WorkspaceId}.ap-southeast-1.maas.aliyuncs.com/api-ws/v1/realtime"
# 配置 API Key，若沒有設定環境變數，請用 API Key 將下行替換為 dashscope.api_key = "sk-xxx"
dashscope.api_key = os.getenv('DASHSCOPE_API_KEY')
AUDIO_SAMPLE_RATE = 16000
AUDIO_CHUNK_SIZE = 3200
OUTPUT_AUDIO_SAMPLE_RATE = 24000

# ==================== 工具定義 ====================
def get_train_price(src: str, dst: str) -> str:
    """查詢火車票價格"""
    return f"{src}到{dst}的火車票價格為100~200元。"

def get_flight_price(src: str, dst: str) -> str:
    """查詢飛機票價格"""
    return f"{src}到{dst}的機票價格為200~300美元。"

def get_current_weather(location: str) -> str:
    """查詢指定城市天氣"""
    return f"{location}今天天氣為霾轉晴，氣溫4/-4℃，微風"

# 統一的 OpenAI 格式工具定義
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_current_weather",
            "description": "當你想查詢指定城市的天氣時非常有用。",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "城市或縣區，比如北京市、杭州市、餘杭區等。",
                    }
                },
                "required": ["location"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_flight_price",
            "description": "當你想查詢飛機票價格時非常有用。",
            "parameters": {
                "type": "object",
                "properties": {
                    "src": {
                        "type": "string",
                        "description": "飛機起飛的城市，比如北京市、杭州市等。",
                    },
                    "dst": {
                        "type": "string",
                        "description": "飛機降落的城市，比如北京市、杭州市區等。",
                    },
                },
                "required": ["src", "dst"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_train_price",
            "description": "當你想查詢火車票價格時非常有用。",
            "parameters": {
                "type": "object",
                "properties": {
                    "src": {
                        "type": "string",
                        "description": "火車出發的城市，比如北京市、杭州市等。",
                    },
                    "dst": {
                        "type": "string",
                        "description": "火車到達的城市，比如北京市、杭州市區等。",
                    },
                },
                "required": ["src", "dst"],
            },
        },
    },
]

# 工具名稱到函數的映射
TOOL_FUNCTIONS = {
    "get_current_weather": get_current_weather,
    "get_flight_price": get_flight_price,
    "get_train_price": get_train_price,
}

# ==================== 工具調用處理 ====================
def handle_tool_call(tool_call_response: Dict[str, Any]) -> Dict[str, Any]:
    """
    處理工具調用請求

    Args:
        tool_call_response: 包含 name, arguments, call_id 的工具調用資訊

    Returns:
        更新後的工具調用響應，包含 output 欄位
    """
    try:
        function_name = tool_call_response['name']
        tool_call_arguments = json.loads(tool_call_response['arguments'])

        print(f'[Tool Call] 開始處理: name={function_name}, args={tool_call_arguments}')

        # 尋找對應的函數
        if function_name not in TOOL_FUNCTIONS:
            tool_call_response['output'] = f"用戶端未找到工具: {function_name}"
            print(f'[Tool Call] 錯誤: 未找到工具 {function_name}')
            return tool_call_response

        # 調用函數
        func = TOOL_FUNCTIONS[function_name]
        result = func(**tool_call_arguments)
        tool_call_response['output'] = result

        print(f'[Tool Call] 完成: {result}')
        return tool_call_response

    except Exception as e:
        error_msg = f"工具調用失敗: {str(e)}"
        tool_call_response['output'] = error_msg
        print(f'[Tool Call] 異常: {error_msg}')
        traceback.print_exc()
        return tool_call_response

def send_tool_call_response(conversation: OmniRealtimeConversation, response: Dict[str, Any]) -> None:
    """發送工具調用結果到服務端"""
    conversation.create_item({
        "id": 'item_' + uuid.uuid4().hex,
        "type": "function_call_output",
        "call_id": response['call_id'],
        "output": response["output"],
    })

# ==================== PCM 音頻播放器 ====================
class PCMPlayer:
    """
    PCM 音頻播放器

    使用雙線程架構實現即時音頻播放：
    - 解碼線程：將 base64 編碼的音頻資料解碼為原始 PCM 資料
    - 播放線程：將 PCM 資料寫入音訊輸出裝置

    支援動態添加音頻資料、取消播放、儲存音頻檔案等功能。
    """

    def __init__(self, pya: pyaudio.PyAudio, sample_rate=24000, chunk_size_ms=100, save_file=False):
        """
        初始化 PCM 播放器

        Args:
            pya: pyaudio.PyAudio 執行個體
            sample_rate: 音頻採樣率（Hz），預設 24000
            chunk_size_ms: 音頻塊大小（毫秒），影響取消播放的延遲，預設 100ms
            save_file: 是否儲存播放的音頻到檔案（result.pcm），預設 False
        """

        self.pya = pya
        self.sample_rate = sample_rate
        self.chunk_size_bytes = chunk_size_ms * sample_rate * 2 // 1000
        self.player_stream = pya.open(format=pyaudio.paInt16,
                                       channels=1,
                                       rate=sample_rate,
                                       output=True)

        self.raw_audio_buffer: queue.Queue = queue.Queue()
        self.b64_audio_buffer: queue.Queue = queue.Queue()
        self.status_lock = threading.Lock()
        self.status = 'playing'
        self.decoder_thread = threading.Thread(target=self.decoder_loop)
        self.player_thread = threading.Thread(target=self.player_loop)
        self.decoder_thread.start()
        self.player_thread.start()
        self.complete_event: threading.Event = None
        self.save_file = save_file
        if self.save_file:
            self.out_file = open('result.pcm', 'wb')

    def decoder_loop(self):
        """解碼線程：將 base64 音頻資料解碼為 PCM 未經處理資料"""
        while self.status != 'stop':
            recv_audio_b64 = None
            with contextlib.suppress(queue.Empty):
                recv_audio_b64 = self.b64_audio_buffer.get(timeout=0.1)
            if recv_audio_b64 is None:
                continue
            recv_audio_raw = base64.b64decode(recv_audio_b64)
            # push raw audio data into queue by chunk
            for i in range(0, len(recv_audio_raw), self.chunk_size_bytes):
                chunk = recv_audio_raw[i:i + self.chunk_size_bytes]
                self.raw_audio_buffer.put(chunk)
                if self.save_file:
                    self.out_file.write(chunk)

    def player_loop(self):
        """播放線程：將 PCM 資料寫入音訊輸出裝置"""
        while self.status != 'stop':
            recv_audio_raw = None
            with contextlib.suppress(queue.Empty):
                recv_audio_raw = self.raw_audio_buffer.get(timeout=0.1)
            if recv_audio_raw is None:
                if self.complete_event:
                    self.complete_event.set()
                continue
            # write chunk to pyaudio audio player, wait until finish playing this chunk.
            self.player_stream.write(recv_audio_raw)

    def cancel_playing(self):
        """取消播放：清空所有緩衝隊列"""
        self.b64_audio_buffer.queue.clear()
        self.raw_audio_buffer.queue.clear()

    def add_data(self, data):
        """添加 base64 編碼的音頻資料到播放隊列"""
        self.b64_audio_buffer.put(data)

    def wait_for_complete(self):
        """等待播放完成"""
        self.complete_event = threading.Event()
        self.complete_event.wait()
        self.complete_event = None

    def shutdown(self):
        """關閉播放器並釋放資源"""
        self.status = 'stop'
        self.decoder_thread.join()
        self.player_thread.join()
        self.player_stream.close()
        if self.save_file:
            self.out_file.close()

# ==================== 音頻管理器 ====================
class AudioManager:
    """管理音頻輸入輸出資源"""

    def __init__(self):
        self.pya: Optional[pyaudio.PyAudio] = None
        self.mic_stream: Optional[pyaudio.Stream] = None
        self.player: Optional[PCMPlayer] = None

    def initialize(self) -> None:
        """初始化音訊裝置"""
        print('初始化音訊裝置...')
        self.pya = pyaudio.PyAudio()
        self.mic_stream = self.pya.open(
            format=pyaudio.paInt16,
            channels=1,
            rate=AUDIO_SAMPLE_RATE,
            input=True
        )
        self.player = PCMPlayer(self.pya, sample_rate=OUTPUT_AUDIO_SAMPLE_RATE)
        print('音訊裝置初始化完成')

    def read_audio_chunk(self) -> Optional[bytes]:
        """讀取音頻資料區塊"""
        if not self.mic_stream:
            return None
        try:
            return self.mic_stream.read(AUDIO_CHUNK_SIZE, exception_on_overflow=False)
        except Exception as e:
            print(f'[Error] 讀取音頻資料失敗: {e}')
            return None

    def cleanup(self) -> None:
        """清理音頻資源"""
        print('清理音頻資源...')
        if self.player:
            self.player.shutdown()
        if self.mic_stream:
            self.mic_stream.close()
        if self.pya:
            self.pya.terminate()
        print('音頻資源清理完成')

# ==================== 回調處理器 ====================
class OmniCallback(OmniRealtimeCallback):
    """Omni 即時對話回調處理器"""

    def __init__(self, audio_manager: AudioManager):
        self.audio_manager = audio_manager
        self.tool_calls: Dict[str, Dict[str, Any]] = {}
        self.all_response_text: str = ''
        self.last_package_time: float = 0
        self.is_first_text: bool = True
        self.is_first_audio: bool = True
        self.conversation: Optional[OmniRealtimeConversation] = None

    def set_conversation(self, conversation: OmniRealtimeConversation) -> None:
        """設定對話執行個體引用"""
        self.conversation = conversation

    def on_open(self) -> None:
        """串連建立時的回調"""
        print('串連已建立')
        self.audio_manager.initialize()
        self.last_package_time = time.time() * 1000
        self.is_first_text = True
        self.is_first_audio = True
        self.tool_calls = {}
        self.all_response_text = ''

    def on_close(self, close_status_code: int, close_msg: str) -> None:
        """串連關閉時的回調"""
        print(f'串連已關閉: code={close_status_code}, msg={close_msg}')
        self.audio_manager.cleanup()
        sys.exit(0)

    def on_event(self, response: Dict[str, Any]) -> None:
        """處理事件回調"""
        try:
            event_type = response.get('type', '')

            # 會話建立
            if event_type == 'session.created':
                print(f'會話已啟動: {response["session"]["id"]}')

            # 語音轉文本完成
            elif event_type == 'conversation.item.input_audio_transcription.completed':
                print(f'使用者問題: {response.get("transcript", "")}')

            # 文本增量響應
            elif event_type in ('response.audio_transcript.delta', 'response.text.delta'):
                if self.is_first_text:
                    self.is_first_text = False
                    latency = time.time() * 1000 - self.last_package_time
                    print(f'首字延遲 (VAD結束): {latency:.0f} ms')

                text = response.get('delta', '')
                self.all_response_text += text

            # 音頻增量響應
            elif event_type == 'response.audio.delta':
                if self.is_first_audio:
                    self.is_first_audio = False
                    latency = time.time() * 1000 - self.last_package_time
                    print(f'首音延遲 (VAD結束): {latency:.0f} ms')

                audio_interval = time.time() * 1000 - self.last_package_time
                print(f'音頻間隔: {audio_interval:.0f} ms')
                self.last_package_time = time.time() * 1000

                recv_audio_b64 = response.get('delta', '')
                if self.audio_manager.player:
                    self.audio_manager.player.add_data(recv_audio_b64)

            # VAD 檢測到語音開始
            elif event_type == 'input_audio_buffer.speech_started':
                print('====== VAD 檢測到語音開始 ======')
                if self.audio_manager.player:
                    self.audio_manager.player.cancel_playing()

            # VAD 檢測到語音結束
            elif event_type == 'input_audio_buffer.speech_stopped':
                print('====== VAD 檢測到語音結束 ======')
                self.last_package_time = time.time() * 1000
                self.is_first_text = True
                self.is_first_audio = True
                self.tool_calls = {}

            # 函數調用參數完成
            elif event_type == 'response.function_call_arguments.done':
                print('====== 收到工具調用請求 ======')
                call_id = response.get('call_id', '')
                self.tool_calls[call_id] = response.copy()
                self.tool_calls[call_id]['processed'] = False

            # 響應完成
            elif event_type == 'response.done':
                print('====== 響應完成 ======')
                print(f'完整回複: {self.all_response_text}')

                if self.conversation:
                    response_id = self.conversation.get_last_response_id()
                    text_delay = self.conversation.get_last_first_text_delay()
                    audio_delay = self.conversation.get_last_first_audio_delay()

                    # 只有當所有指標都可用時才列印詳細指標
                    if response_id is not None and text_delay is not None and audio_delay is not None:
                        print(f'[Metric] 響應ID: {response_id}, '
                              f'首字延遲: {text_delay:.0f}ms, '
                              f'首音延遲: {audio_delay:.0f}ms')
                    else:
                        print('[Metric] 指標資訊暫不可用（可能是工具調用後的響應）')

                self.all_response_text = ''

        except Exception as e:
            print(f'[Error] 處理事件異常: {e}')
            traceback.print_exc()

    def process_pending_tool_calls(self) -> bool:
        """
        處理待處理的工具調用

        Returns:
            是否有新的工具調用需要響應
        """
        has_pending = False

        for call_id, tool_call in self.tool_calls.items():
            if not tool_call.get('processed', False):
                has_pending = True
                tool_call['processed'] = True

                # 處理工具調用
                result = handle_tool_call(tool_call)

                # 發送結果到服務端
                if self.conversation:
                    send_tool_call_response(self.conversation, result)

        return has_pending

# ==================== 主程式 ====================
def main():
    """主函數"""
    print('正在初始化 Omni 即時對話...')

    # 建立音頻管理器
    audio_manager = AudioManager()

    # 建立回調處理器
    callback = OmniCallback(audio_manager)

    # 建立對話執行個體
    conversation = OmniRealtimeConversation(
        api_key=dashscope.api_key,
        url=WS_URL,
        model=MODEL,
        callback=callback,
    )

    # 設定回調中的對話引用
    callback.set_conversation(conversation)

    # 建立串連
    conversation.connect()

    # 配置會話參數
    omni_output_modalities = [MultiModality.AUDIO, MultiModality.TEXT]

    conversation.update_session(
        output_modalities=omni_output_modalities,
        voice=VOICE,
        input_audio_format=AudioFormat.PCM_16000HZ_MONO_16BIT,
        output_audio_format=AudioFormat.PCM_24000HZ_MONO_16BIT,
        enable_input_audio_transcription=True,
        enable_turn_detection=True,
        turn_detection_type='server_vad',
        tools=TOOLS,
    )

    # 設定訊號處理
    def signal_handler(sig, frame):
        print('\n接收到 Ctrl+C，正在停止...')
        conversation.close()
        audio_manager.cleanup()
        print('Omni 即時對話已停止')
        sys.exit(0)

    signal.signal(signal.SIGINT, signal_handler)
    print("按 Ctrl+C 停止對話...\n")

    # 主迴圈：持續發送音頻並檢查工具調用
    try:
        while True:
            # 處理待處理的工具調用
            has_tool_calls = callback.process_pending_tool_calls()

            if has_tool_calls:
                print("*** 工具調用完成，建立新響應 ***")
                conversation.create_response(
                    instructions=None,
                    output_modalities=omni_output_modalities
                )
                print('====== 工具調用處理完成 ======\n')

            # 讀取並發送音頻資料
            audio_data = audio_manager.read_audio_chunk()
            if audio_data:
                audio_b64 = base64.b64encode(audio_data).decode('ascii')
                conversation.append_audio(audio_b64)
            else:
                break

    except KeyboardInterrupt:
        signal_handler(signal.SIGINT, None)
    except Exception as e:
        print(f'[Error] 主迴圈異常: {e}')
        traceback.print_exc()
    finally:
        conversation.close()
        audio_manager.cleanup()

if __name__ == '__main__':
    main()