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 *

# ==================== Constant Definitions ====================
VOICE = 'Tina'
MODEL = "qwen3.5-omni-plus-realtime"
# To access the Beijing region, replace WS_URL with: 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"
# Configure the API key. If you have not set the environment variable, replace the following line with your 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

# ==================== Tool Definitions ====================
def get_train_price(src: str, dst: str) -> str:
    """Query train ticket prices"""
    return f"The train ticket price from {src} to {dst} is 100-200 CNY."

def get_flight_price(src: str, dst: str) -> str:
    """Query flight ticket prices"""
    return f"The flight ticket price from {src} to {dst} is 200-300 USD."

def get_current_weather(location: str) -> str:
    """Query the weather in a specific city"""
    return f"The weather in {location} today is changing from haze to sunny, with a temperature of 4/-4°C and a light breeze."

# Unified OpenAI format tool definitions
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_current_weather",
            "description": "Useful for when you want to query the weather in a specific city.",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "A city or district, such as Beijing, Hangzhou, or Yuhang.",
                    }
                },
                "required": ["location"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_flight_price",
            "description": "Useful for when you want to query flight ticket prices.",
            "parameters": {
                "type": "object",
                "properties": {
                    "src": {
                        "type": "string",
                        "description": "The departure city of the flight, such as Beijing or Hangzhou.",
                    },
                    "dst": {
                        "type": "string",
                        "description": "The arrival city of the flight, such as Beijing or Hangzhou.",
                    },
                },
                "required": ["src", "dst"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_train_price",
            "description": "Useful for when you want to query train ticket prices.",
            "parameters": {
                "type": "object",
                "properties": {
                    "src": {
                        "type": "string",
                        "description": "The departure city of the train, such as Beijing or Hangzhou.",
                    },
                    "dst": {
                        "type": "string",
                        "description": "The arrival city of the train, such as Beijing or Hangzhou.",
                    },
                },
                "required": ["src", "dst"],
            },
        },
    },
]

# Mapping of tool names to functions
TOOL_FUNCTIONS = {
    "get_current_weather": get_current_weather,
    "get_flight_price": get_flight_price,
    "get_train_price": get_train_price,
}

# ==================== Tool Call Handling ====================
def handle_tool_call(tool_call_response: Dict[str, Any]) -> Dict[str, Any]:
    """
    Handles a tool call request

    Args:
        tool_call_response: Tool call information containing name, arguments, and call_id

    Returns:
        Updated tool call response containing the output field
    """
    try:
        function_name = tool_call_response['name']
        tool_call_arguments = json.loads(tool_call_response['arguments'])

        print(f'[Tool Call] Start processing: name={function_name}, args={tool_call_arguments}')

        # Find the corresponding function
        if function_name not in TOOL_FUNCTIONS:
            tool_call_response['output'] = f"Client did not find the tool: {function_name}"
            print(f'[Tool Call] Error: Tool not found {function_name}')
            return tool_call_response

        # Call the function
        func = TOOL_FUNCTIONS[function_name]
        result = func(**tool_call_arguments)
        tool_call_response['output'] = result

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

    except Exception as e:
        error_msg = f"Tool call failed: {str(e)}"
        tool_call_response['output'] = error_msg
        print(f'[Tool Call] Exception: {error_msg}')
        traceback.print_exc()
        return tool_call_response

def send_tool_call_response(conversation: OmniRealtimeConversation, response: Dict[str, Any]) -> None:
    """Sends the tool call result to the server"""
    conversation.create_item({
        "id": 'item_' + uuid.uuid4().hex,
        "type": "function_call_output",
        "call_id": response['call_id'],
        "output": response["output"],
    })

# ==================== PCM Audio Player ====================
class PCMPlayer:
    """
    PCM Audio Player

    Uses a dual-thread architecture for real-time audio playback:
    - Decoding thread: Decodes base64-encoded audio data into raw PCM data
    - Playback thread: Writes PCM data to the audio output device

    Supports dynamically adding audio data, canceling playback, saving audio files, etc.
    """

    def __init__(self, pya: pyaudio.PyAudio, sample_rate=24000, chunk_size_ms=100, save_file=False):
        """
        Initializes the PCM player

        Args:
            pya: pyaudio.PyAudio instance
            sample_rate: Audio sampling rate (Hz), default 24000
            chunk_size_ms: Audio chunk size (milliseconds), affects playback cancellation latency, default 100ms
            save_file: Whether to save the played audio to a file (result.pcm), default 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):
        """Decoding thread: Decodes base64 audio data into raw PCM data"""
        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):
        """Playback thread: Writes PCM data to the audio output device"""
        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):
        """Cancel playback: Clear all buffer queues"""
        self.b64_audio_buffer.queue.clear()
        self.raw_audio_buffer.queue.clear()

    def add_data(self, data):
        """Add base64-encoded audio data to the playback queue"""
        self.b64_audio_buffer.put(data)

    def wait_for_complete(self):
        """Wait for playback to complete"""
        self.complete_event = threading.Event()
        self.complete_event.wait()
        self.complete_event = None

    def shutdown(self):
        """Shut down the player and release resources"""
        self.status = 'stop'
        self.decoder_thread.join()
        self.player_thread.join()
        self.player_stream.close()
        if self.save_file:
            self.out_file.close()

# ==================== Audio Manager ====================
class AudioManager:
    """Manages audio input and output resources"""

    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:
        """Initialize audio devices"""
        print('Initializing audio devices...')
        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('Audio devices initialized')

    def read_audio_chunk(self) -> Optional[bytes]:
        """Read an audio data chunk"""
        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] Failed to read audio data: {e}')
            return None

    def cleanup(self) -> None:
        """Clean up audio resources"""
        print('Cleaning up audio resources...')
        if self.player:
            self.player.shutdown()
        if self.mic_stream:
            self.mic_stream.close()
        if self.pya:
            self.pya.terminate()
        print('Audio resources cleaned up')

# ==================== Callback Handler ====================
class OmniCallback(OmniRealtimeCallback):
    """Omni real-time conversation callback handler"""

    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:
        """Set the conversation instance reference"""
        self.conversation = conversation

    def on_open(self) -> None:
        """Callback on connection establishment"""
        print('Connection established')
        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:
        """Callback on connection closure"""
        print(f'Connection closed: code={close_status_code}, msg={close_msg}')
        self.audio_manager.cleanup()
        sys.exit(0)

    def on_event(self, response: Dict[str, Any]) -> None:
        """Handle event callbacks"""
        try:
            event_type = response.get('type', '')

            # Session created
            if event_type == 'session.created':
                print(f'Session started: {response["session"]["id"]}')

            # Speech-to-text completed
            elif event_type == 'conversation.item.input_audio_transcription.completed':
                print(f'User question: {response.get("transcript", "")}')

            # Incremental text response
            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'Time to first token (VAD end): {latency:.0f} ms')

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

            # Incremental audio response
            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'Time to first audio (VAD end): {latency:.0f} ms')

                audio_interval = time.time() * 1000 - self.last_package_time
                print(f'Audio interval: {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 detected speech start
            elif event_type == 'input_audio_buffer.speech_started':
                print('====== VAD detected speech start ======')
                if self.audio_manager.player:
                    self.audio_manager.player.cancel_playing()

            # VAD detected speech end
            elif event_type == 'input_audio_buffer.speech_stopped':
                print('====== VAD detected speech end ======')
                self.last_package_time = time.time() * 1000
                self.is_first_text = True
                self.is_first_audio = True
                self.tool_calls = {}

            # Function call arguments completed
            elif event_type == 'response.function_call_arguments.done':
                print('====== Received tool call request ======')
                call_id = response.get('call_id', '')
                self.tool_calls[call_id] = response.copy()
                self.tool_calls[call_id]['processed'] = False

            # Response completed
            elif event_type == 'response.done':
                print('====== Response completed ======')
                print(f'Full response: {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()

                    # Print detailed metrics only when all are available
                    if response_id is not None and text_delay is not None and audio_delay is not None:
                        print(f'[Metric] Response ID: {response_id}, '
                              f'Time to first token: {text_delay:.0f}ms, '
                              f'Time to first audio: {audio_delay:.0f}ms')
                    else:
                        print('[Metric] Metric information is temporarily unavailable (possibly a response after a tool call)')

                self.all_response_text = ''

        except Exception as e:
            print(f'[Error] Exception handling event: {e}')
            traceback.print_exc()

    def process_pending_tool_calls(self) -> bool:
        """
        Processes pending tool calls

        Returns:
            Whether there are new tool calls that need a response
        """
        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

                # Handle the tool call
                result = handle_tool_call(tool_call)

                # Send the result to the server
                if self.conversation:
                    send_tool_call_response(self.conversation, result)

        return has_pending

# ==================== Main Program ====================
def main():
    """Main function"""
    print('Initializing Omni real-time conversation...')

    # Create an audio manager
    audio_manager = AudioManager()

    # Create a callback handler
    callback = OmniCallback(audio_manager)

    # Create a conversation instance
    conversation = OmniRealtimeConversation(
        api_key=dashscope.api_key,
        url=WS_URL,
        model=MODEL,
        callback=callback,
    )

    # Set the conversation reference in the callback
    callback.set_conversation(conversation)

    # Establish the connection
    conversation.connect()

    # Configure session parameters
    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,
    )

    # Set up signal handling
    def signal_handler(sig, frame):
        print('\nReceived Ctrl+C, stopping...')
        conversation.close()
        audio_manager.cleanup()
        print('Omni real-time conversation stopped')
        sys.exit(0)

    signal.signal(signal.SIGINT, signal_handler)
    print("Press Ctrl+C to stop the conversation...\n")

    # Main loop: continuously send audio and check for tool calls
    try:
        while True:
            # Process pending tool calls
            has_tool_calls = callback.process_pending_tool_calls()

            if has_tool_calls:
                print("*** Tool call completed, creating new response ***")
                conversation.create_response(
                    instructions=None,
                    output_modalities=omni_output_modalities
                )
                print('====== Tool call processing completed ======\n')

            # Read and send audio data
            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] Main loop exception: {e}')
        traceback.print_exc()
    finally:
        conversation.close()
        audio_manager.cleanup()

if __name__ == '__main__':
    main()