Realtime audio and video understanding

# -- coding: utf-8 --
import os
import asyncio
import time
import threading
import queue
import pyaudio
from omni_realtime_client import OmniRealtimeClient, TurnDetectionMode


class AudioPlayer:
  """Real-time audio player class"""

  def __init__(self, sample_rate=24000, channels=1, sample_width=2):
    self.sample_rate = sample_rate
    self.channels = channels
    self.sample_width = sample_width  # 2 bytes for 16-bit
    self.audio_queue = queue.Queue()
    self.is_playing = False
    self.play_thread = None
    self.pyaudio_instance = None
    self.stream = None
    self._lock = threading.Lock()  # Add a lock for synchronized access
    self._last_data_time = time.time()  # Record the time the last data was received
    self._response_done = False  # Add a flag to indicate response completion
    self._waiting_for_response = False  # Flag to indicate if waiting for a server response
    # Record the time the last data was written to the audio stream and the duration of the most recent audio chunk for more accurate playback end detection
    self._last_play_time = time.time()
    self._last_chunk_duration = 0.0

  def start(self):
    """Start the audio player"""
    with self._lock:
      if self.is_playing:
        return

      self.is_playing = True

      try:
        self.pyaudio_instance = pyaudio.PyAudio()

        # Create an audio output stream
        self.stream = self.pyaudio_instance.open(
          format=pyaudio.paInt16,  # 16-bit
          channels=self.channels,
          rate=self.sample_rate,
          output=True,
          frames_per_buffer=1024
        )

        # Start the playback thread
        self.play_thread = threading.Thread(target=self._play_audio)
        self.play_thread.daemon = True
        self.play_thread.start()

        print("Audio player started")
      except Exception as e:
        print(f"Failed to start audio player: {e}")
        self._cleanup_resources()
        raise

  def stop(self):
    """Stop the audio player"""
    with self._lock:
      if not self.is_playing:
        return

      self.is_playing = False

    # Clear the queue
    while not self.audio_queue.empty():
      try:
        self.audio_queue.get_nowait()
      except queue.Empty:
        break

    # Wait for the playback thread to finish (wait outside the lock to avoid deadlock)
    if self.play_thread and self.play_thread.is_alive():
      self.play_thread.join(timeout=2.0)

    # Acquire the lock again to clean up resources
    with self._lock:
      self._cleanup_resources()

    print("Audio player stopped")

  def _cleanup_resources(self):
    """Clean up audio resources (must be called within the lock)"""
    try:
      # Close the audio stream
      if self.stream:
        if not self.stream.is_stopped():
          self.stream.stop_stream()
        self.stream.close()
        self.stream = None
    except Exception as e:
      print(f"Error closing audio stream: {e}")

    try:
      if self.pyaudio_instance:
        self.pyaudio_instance.terminate()
        self.pyaudio_instance = None
    except Exception as e:
      print(f"Error terminating PyAudio: {e}")

  def add_audio_data(self, audio_data):
    """Add audio data to the playback queue"""
    if self.is_playing and audio_data:
      self.audio_queue.put(audio_data)
      with self._lock:
        self._last_data_time = time.time()  # Update the time the last data was received
        self._waiting_for_response = False  # Data received, no longer waiting

  def stop_receiving_data(self):
    """Mark that no more new audio data will be received"""
    with self._lock:
      self._response_done = True
      self._waiting_for_response = False  # Response ended, no longer waiting

  def prepare_for_next_turn(self):
    """Reset the player state for the next conversation turn."""
    with self._lock:
      self._response_done = False
      self._last_data_time = time.time()
      self._last_play_time = time.time()
      self._last_chunk_duration = 0.0
      self._waiting_for_response = True  # Start waiting for the next response

    # Clear any remaining audio data from the previous turn
    while not self.audio_queue.empty():
      try:
        self.audio_queue.get_nowait()
      except queue.Empty:
        break

  def is_finished_playing(self):
    """Check if all audio data has been played"""
    with self._lock:
      queue_size = self.audio_queue.qsize()
      time_since_last_data = time.time() - self._last_data_time
      time_since_last_play = time.time() - self._last_play_time

      # ---------------------- Smart end detection ----------------------
      # 1. Preferred: If the server has marked completion and the playback queue is empty.
      #    Wait for the most recent audio chunk to finish playing (chunk duration + 0.1s tolerance).
      if self._response_done and queue_size == 0:
        min_wait = max(self._last_chunk_duration + 0.1, 0.5)  # Wait at least 0.5s
        if time_since_last_play >= min_wait:
          return True

      # 2. Fallback: If no new data has been received for a long time and the playback queue is empty.
      #    This logic serves as a safeguard if the server does not explicitly send `response.done`.
      if not self._waiting_for_response and queue_size == 0 and time_since_last_data > 1.0:
        print("\n(No new audio received for a while, assuming playback is finished)")
        return True

      return False

  def _play_audio(self):
    """Worker thread for playing audio data"""
    while True:
      # Check if it should stop
      with self._lock:
        if not self.is_playing:
          break
        stream_ref = self.stream  # Get a reference to the stream

      try:
        # Get audio data from the queue, with a timeout of 0.1 seconds
        audio_data = self.audio_queue.get(timeout=0.1)

        # Check the status and stream validity again
        with self._lock:
          if self.is_playing and stream_ref and not stream_ref.is_stopped():
            try:
              # Play the audio data
              stream_ref.write(audio_data)
              # Update the latest playback information
              self._last_play_time = time.time()
              self._last_chunk_duration = len(audio_data) / (
                    self.channels * self.sample_width) / self.sample_rate
            except Exception as e:
              print(f"Error writing to audio stream: {e}")
              break

        # Mark this data block as processed
        self.audio_queue.task_done()

      except queue.Empty:
        # Continue waiting if the queue is empty
        continue
      except Exception as e:
        print(f"Error playing audio: {e}")
        break


class MicrophoneRecorder:
  """Real-time microphone recorder"""

  def __init__(self, sample_rate=16000, channels=1, chunk_size=3200):
    self.sample_rate = sample_rate
    self.channels = channels
    self.chunk_size = chunk_size
    self.pyaudio_instance = None
    self.stream = None
    self.frames = []
    self._is_recording = False
    self._record_thread = None

  def _recording_thread(self):
    """Recording worker thread"""
    # Continuously read data from the audio stream while _is_recording is True
    while self._is_recording:
      try:
        # Use exception_on_overflow=False to avoid crashing due to buffer overflow
        data = self.stream.read(self.chunk_size, exception_on_overflow=False)
        self.frames.append(data)
      except (IOError, OSError) as e:
        # Reading from the stream might raise an error when it's closed
        print(f"Error reading from recording stream, it might be closed: {e}")
        break

  def start(self):
    """Start recording"""
    if self._is_recording:
      print("Recording is already in progress.")
      return

    self.frames = []
    self._is_recording = True

    try:
      self.pyaudio_instance = pyaudio.PyAudio()
      self.stream = self.pyaudio_instance.open(
        format=pyaudio.paInt16,
        channels=self.channels,
        rate=self.sample_rate,
        input=True,
        frames_per_buffer=self.chunk_size
      )

      self._record_thread = threading.Thread(target=self._recording_thread)
      self._record_thread.daemon = True
      self._record_thread.start()
      print("Microphone recording started...")
    except Exception as e:
      print(f"Failed to start microphone: {e}")
      self._is_recording = False
      self._cleanup()
      raise

  def stop(self):
    """Stop recording and return the audio data"""
    if not self._is_recording:
      return None

    self._is_recording = False

    # Wait for the recording thread to exit safely
    if self._record_thread:
      self._record_thread.join(timeout=1.0)

    self._cleanup()

    print("Microphone recording stopped.")
    return b''.join(self.frames)

  def _cleanup(self):
    """Safely clean up PyAudio resources"""
    if self.stream:
      try:
        if self.stream.is_active():
          self.stream.stop_stream()
        self.stream.close()
      except Exception as e:
        print(f"Error closing audio stream: {e}")

    if self.pyaudio_instance:
      try:
        self.pyaudio_instance.terminate()
      except Exception as e:
        print(f"Error terminating PyAudio instance: {e}")

    self.stream = None
    self.pyaudio_instance = None


async def interactive_test():
  """
  Interactive test script: Allows for multi-turn conversations, with audio and images sent in each turn.
  """
  # ------------------- 1. Initialization and connection (one-time) -------------------
  api_key = os.environ.get("DASHSCOPE_API_KEY")
  if not api_key:
    print("Please set the DASHSCOPE_API_KEY environment variable.")
    return

  print("--- Real-time Multimodal Audio/Video Chat Client ---")
  print("Initializing audio player and client...")

  audio_player = AudioPlayer()
  audio_player.start()

  def on_audio_received(audio_data):
    audio_player.add_audio_data(audio_data)

  def on_response_done(event):
    print("\n(Received response end marker)")
    audio_player.stop_receiving_data()

  realtime_client = OmniRealtimeClient(
    base_url="wss://maas.qwencloudapi.com/api-ws/v1/realtime",
    api_key=api_key,
    model="qwen3.8-omni-flash-realtime",
    voice="Tina",
    instructions="You are Xiaoyun, a personal assistant. Please answer the user's questions accurately and friendly, always responding with a helpful attitude.", # Set the model role
    on_text_delta=lambda text: print(f"Assistant reply: {text}", end="", flush=True),
    on_audio_delta=on_audio_received,
    turn_detection_mode=TurnDetectionMode.MANUAL,
    extra_event_handlers={"response.done": on_response_done}
  )

  message_handler_task = None
  try:
    await realtime_client.connect()
    print("Connected to the server. Enter 'q' or 'quit' to exit at any time.")
    message_handler_task = asyncio.create_task(realtime_client.handle_messages())
    await asyncio.sleep(0.5)

    turn_counter = 1
    # ------------------- 2. Multi-turn conversation loop -------------------
    while True:
      print(f"\n--- Turn {turn_counter} ---")
      audio_player.prepare_for_next_turn()

      recorded_audio = None
      image_paths = []

      # --- Get user input: Record from microphone ---
      loop = asyncio.get_event_loop()
      recorder = MicrophoneRecorder(sample_rate=16000)  # 16k sample rate is recommended for speech recognition

      print("Ready to record. Press Enter to start recording (or enter 'q' to exit)...")
      user_input = await loop.run_in_executor(None, input)
      if user_input.strip().lower() in ['q', 'quit']:
        print("User requested to exit...")
        return

      try:
        recorder.start()
      except Exception:
        print("Could not start recording. Please check your microphone permissions and device. Skipping this turn.")
        continue

      print("Recording... Press Enter again to stop.")
      await loop.run_in_executor(None, input)

      recorded_audio = recorder.stop()

      if not recorded_audio or len(recorded_audio) == 0:
        print("No valid audio was recorded. Please start this turn again.")
        continue

      # --- 3. Send data and get response ---
      print("\n--- Input Confirmation ---")
      print(f"Audio to process: 1 (from microphone), Images: {len(image_paths)}")
      print("------------------")

      # 3.1 Send the recorded audio
      try:
        print(f"Sending microphone recording ({len(recorded_audio)} bytes)")
        await realtime_client.stream_audio(recorded_audio)
        await asyncio.sleep(0.1)
      except Exception as e:
        print(f"Failed to send microphone recording: {e}")
        continue

      # 3.3 Submit and wait for response
      print("Submitting all inputs, requesting server response...")
      await realtime_client.commit_audio_buffer()
      await realtime_client.create_response()

      print("Waiting for and playing server response audio...")
      start_time = time.time()
      max_wait_time = 60
      while not audio_player.is_finished_playing():
        if time.time() - start_time > max_wait_time:
          print(f"\nWait timed out ({max_wait_time} seconds). Moving to the next turn.")
          break
        await asyncio.sleep(0.2)

      print("\nAudio playback for this turn is complete!")
      turn_counter += 1

  except (asyncio.CancelledError, KeyboardInterrupt):
    print("\nProgram was interrupted.")
  except Exception as e:
    print(f"An unhandled error occurred: {e}")
  finally:
    # ------------------- 4. Clean up resources -------------------
    print("\nClosing connection and cleaning up resources...")
    if message_handler_task and not message_handler_task.done():
      message_handler_task.cancel()

    if 'realtime_client' in locals() and realtime_client.ws and not realtime_client.ws.close:
      await realtime_client.close()
      print("Connection closed.")

    audio_player.stop()
    print("Program exited.")


if __name__ == "__main__":
  try:
    asyncio.run(interactive_test())
  except KeyboardInterrupt:
    print("\nProgram was forcibly exited by the user.")