Real-Time Voice AI: Streaming Whisper Transcription with WebRTC & WebSockets in Rust

// client/src/audioStreamer.ts
export class AudioStreamer {
  private audioContext: AudioContext | null = null;
  private mediaStream: MediaStream | null = null;
  private audioInput: MediaStreamAudioSourceNode | null = null;
  private processor: AudioWorkletNode | ScriptProcessorNode | null = null;
  private ws: WebSocket | null = null;
  private readonly sampleRate = 16000; // Target sample rate for Whisper
  private readonly bufferSize = 4096; // Audio buffer size for processing

  constructor(private websocketUrl: string) {}

  public async start(): Promise<void> {
    if (this.audioContext) {
      console.warn("AudioStreamer already started.");
      return;
    }

    this.audioContext = new (window.AudioContext || (window as any).webkitAudioContext)({
      sampleRate: this.sampleRate,
    });

    try {
      this.mediaStream = await navigator.mediaDevices.getUserMedia({ audio: true });
      this.audioInput = this.audioContext.createMediaStreamSource(this.mediaStream);

      // Use AudioWorklet for better performance and off-main-thread processing
      // Fallback to ScriptProcessorNode if AudioWorklet is not available
      if (this.audioContext.audioWorklet) {
        await this.audioContext.audioWorklet.addModule('/audio-processor.js');
        this.processor = new AudioWorkletNode(this.audioContext, 'audio-processor');
      } else {
        console.warn("AudioWorklet not supported, falling back to ScriptProcessorNode.");
        this.processor = this.audioContext.createScriptProcessor(this.bufferSize, 1, 1);
        (this.processor as ScriptProcessorNode).onaudioprocess = this.handleAudioProcess.bind(this);
      }

      this.audioInput.connect(this.processor);
      this.processor.connect(this.audioContext.destination); // Connect to destination to keep it alive

      this.ws = new WebSocket(this.websocketUrl);
      this.ws.binaryType = 'arraybuffer';

      this.ws.onopen = () => console.log('WebSocket connected.');
      this.ws.onclose = () => console.log('WebSocket disconnected.');
      this.ws.onerror = (error) => console.error('WebSocket error:', error);

      if (this.processor instanceof AudioWorkletNode) {
        this.processor.port.onmessage = (event) => {
          if (this.ws?.readyState === WebSocket.OPEN) {
            this.ws.send(event.data); // Send raw PCM float32 data
          }
        };
      }

    } catch (error) {
      console.error('Error starting audio stream:', error);
      this.stop();
      throw error;
    }
  }

  private handleAudioProcess(event: AudioProcessingEvent): void {
    if (this.ws?.readyState === WebSocket.OPEN) {
      const inputBuffer = event.inputBuffer.getChannelData(0);
      this.ws.send(inputBuffer.buffer); // Send raw PCM float32 data
    }
  }

  public stop(): void {
    if (this.processor) {
      this.processor.disconnect();
      this.processor = null;
    }
    if (this.audioInput) {
      this.audioInput.disconnect();
      this.audioInput = null;
    }
    if (this.mediaStream) {
      this.mediaStream.getTracks().forEach(track => track.stop());
      this.mediaStream = null;
    }
    if (this.audioContext) {
      this.audioContext.close();
      this.audioContext = null;
    }
    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
    console.log('AudioStreamer stopped.');
  }
}

// client/public/audio-processor.js (AudioWorklet module)
// This file needs to be served by your web server
class AudioProcessor extends AudioWorkletProcessor {
  process(inputs, outputs, parameters) {
    const input = inputs[0];
    if (input.length > 0) {
      const channelData = input[0]; // Get the first channel (mono)
      this.port.postMessage(channelData.buffer, [channelData.buffer]); // Transfer ownership
    }
    return true; // Keep the processor alive
  }
}
registerProcessor('audio-processor', AudioProcessor);

# Cargo.toml
[dependencies]
tokio = { version = "1", features = ["full"] }
warp = "0.3" # Or axum = { version = "0.6", features = ["ws"] }
futures-util = "0.3"
bytes = "1"
log = "0.4"
env_logger = "0.10"
# For whisper.cpp bindings
whisper-rs = { version = "0.1.1", features = ["full"] } # Ensure whisper.cpp is built with GGML_CUDA=1 if using GPU
# For VAD
silero-vad = "0.1.0" # Or a custom VAD implementation
# For audio processing
symphonia = { version = "0.5", features = ["all"] } # For potential audio decoding/resampling if not 16kHz PCM
# For serialization/deserialization
serde = { version = "1", features = ["derive"] }
serde_json = "1"

// src/main.rs
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use futures_util::{StreamExt, SinkExt};
use warp::ws::{Message, WebSocket};
use warp::Filter;
use std::sync::{Arc, Mutex};
use std::collections::VecDeque;
use std::time::{Instant, Duration};

// --- Whisper and VAD related imports ---
use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters};
use silero_vad::{Vad, VadNode};

const SAMPLE_RATE: u32 = 16000;
const WHISPER_MODEL_PATH: &str = "path/to/ggml-medium.en.bin"; // Path to your quantized Whisper model
const VAD_MODEL_PATH: &str = "path/to/silero_vad.onnx"; // Path to your Silero VAD ONNX model

// Audio buffer for a single client
struct ClientAudioBuffer {
    buffer: VecDeque<f32>,
    last_audio_activity: Instant,
    is_speaking: bool,
}

impl ClientAudioBuffer {
    fn new() -> Self {
        ClientAudioBuffer {
            buffer: VecDeque::new(),
            last_audio_activity: Instant::now(),
            is_speaking: false,
        }
    }

    fn push_audio(&mut self, audio_data: &[f32]) {
        self.buffer.extend(audio_data);
        self.last_audio_activity = Instant::now();
    }

    fn get_audio_chunk(&mut self, duration_ms: u64) -> Option<Vec<f32>> {
        let samples_needed = (SAMPLE_RATE as f32 * duration_ms as f32 / 1000.0) as usize;
        if self.buffer.len() >= samples_needed {
            let chunk: Vec<f32> = self.buffer.drain(0..samples_needed).collect();
            Some(chunk)
        } else {
            None
        }
    }

    fn clear_buffer(&mut self) {
        self.buffer.clear();
    }
}

#[tokio::main]
async fn main() {
    env_logger::init();

    // Load Whisper model once
    let whisper_context = Arc::new(
        WhisperContext::new_with_params(
            WHISPER_MODEL_PATH,
            WhisperContextParameters::default(),
        )
        .expect("Failed to load Whisper model"),
    );
    log::info!("Whisper model loaded: {}", WHISPER_MODEL_PATH);

    // Load VAD model once
    let vad_model = Arc::new(
        Vad::builder()
            .set_model_path(VAD_MODEL_PATH)
            .set_sample_rate(SAMPLE_RATE)
            .build()
            .expect("Failed to load Silero VAD model"),
    );
    log::info!("Silero VAD model loaded: {}", VAD_MODEL_PATH);


    let whisper_context_filter = warp::any().map(move || Arc::clone(&whisper_context));
    let vad_model_filter = warp::any().map(move || Arc::clone(&vad_model));

    let ws_route = warp::path("ws")
        .and(warp::ws())
        .and(whisper_context_filter)
        .and(vad_model_filter)
        .map(|ws: warp::ws::Ws, whisper_ctx: Arc<WhisperContext>, vad_model: Arc<Vad>| {
            ws.on_upgrade(move |websocket| handle_websocket(websocket, whisper_ctx, vad_model))
        });

    let routes = ws_route.with(warp::log("websocket_server"));

    log::info!("Server started on 127.0.0.1:8080");
    warp::serve(routes).run(([127, 0, 0, 1], 8080)).await;
}

async fn handle_websocket(
    websocket: WebSocket,
    whisper_ctx: Arc<WhisperContext>,
    vad_model: Arc<Vad>,
) {
    let (mut client_ws_tx, mut client_ws_rx) = websocket.split();

    let (audio_tx, audio_rx) = mpsc::channel::<Vec<f32>>(100); // Channel for raw audio chunks
    let (transcription_tx, mut transcription_rx) = mpsc::channel::<String>(10); // Channel for transcription results

    // Spawn a task to send transcription results back to the client
    tokio::spawn(async move {
        while let Some(transcription) = transcription_rx.recv().await {
            if let Err(e) = client_ws_tx.send(Message::text(transcription)).await {
                log::error!("Failed to send transcription to client: {}", e);
                break;
            }
        }
        log::info!("Transcription sender task terminated.");
    });

    // Spawn a task to process audio and run VAD/Whisper
    let whisper_ctx_clone = Arc::clone(&whisper_ctx);
    let vad_model_clone = Arc::clone(&vad_model);
    tokio::spawn(async move {
        let mut client_audio_buffer = ClientAudioBuffer::new();
        let mut vad_node = VadNode::new(vad_model_clone, SAMPLE_RATE, 512); // VAD processing frame size

        let mut current_speech_buffer: Vec<f32> = Vec::new();
        let mut last_vad_activity = Instant::now();
        let vad_timeout = Duration::from_secs(2); // How long to wait after speech ends before transcribing

        let mut whisper_session = whisper_ctx_clone.create_state().expect("Failed to create Whisper state");

        while let Some(audio_chunk) = audio_rx.recv().await {
            client_audio_buffer.push_audio(&audio_chunk);

            // Process audio in smaller VAD-friendly chunks
            while let Some(vad_chunk) = client_audio_buffer.get_audio_chunk(30) { // 30ms chunks for VAD
                let speech_prob = vad_node.process(&vad_chunk).expect("VAD processing failed");

                if speech_prob > 0.5 { // Threshold for speech detection
                    current_speech_buffer.extend(vad_chunk);
                    last_vad_activity = Instant::now();
                    client_audio_buffer.is_speaking = true;
                } else {
                    // If not speaking, but we were recently, check for timeout
                    if client_audio_buffer.is_speaking && last_vad_activity.elapsed() > vad_timeout {
                        // Speech has ended, transcribe the accumulated buffer
                        if !current_speech_buffer.is_empty() {
                            log::info!("Speech ended, transcribing {} samples.", current_speech_buffer.len());
                            let transcription = run_whisper_inference(
                                &whisper_ctx_clone,
                                &mut whisper_session,
                                &current_speech_buffer,
                            );
                            if let Err(e) = transcription_tx.send(transcription).await {
                                log::error!("Failed to send final transcription: {}", e);
                            }
                            current_speech_buffer.clear();
                        }
                        client_audio_buffer.is_speaking = false;
                    } else if client_audio_buffer.is_speaking {
                        // Still within timeout, keep accumulating non-speech for context
                        current_speech_buffer.extend(vad_chunk);
                    }
                }
            }

            // Periodically transcribe partial results if speaking
            if client_audio_buffer.is_speaking && current_speech_buffer.len() > (SAMPLE_RATE as usize * 1) { // Transcribe every 1 second of speech
                let partial_transcription = run_whisper_inference(
                    &whisper_ctx_clone,
                    &mut whisper_session,
                    &current_speech_buffer,
                );
                if let Err(e) = transcription_tx.send(format!("[partial] {}", partial_transcription)).await {
                    log::error!("Failed to send partial transcription: {}", e);
                }
            }
        }
        // Handle any remaining speech in buffer when audio stream ends
        if !current_speech_buffer.is_empty() {
            log::info!("Stream ended, transcribing remaining {} samples.", current_speech_buffer.len());
            let transcription = run_whisper_inference(
                &whisper_ctx_clone,
                &mut whisper_session,
                &current_speech_buffer,
            );
            if let Err(e) = transcription_tx.send(transcription).await {
                log::error!("Failed to send final transcription on stream end: {}", e);
            }
        }
        log::info!("Audio processor task terminated.");
    });

    // Receive audio data from client
    while let Some(result) = client_ws_rx.next().await {
        match result {
            Ok(msg) => {
                if msg.is_binary() {
                    let audio_bytes = msg.as_bytes();
                    // Assuming client sends f32 raw PCM
                    let audio_data: Vec<f32> = audio_bytes
                        .chunks_exact(4)
                        .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
                        .collect();

                    if let Err(e) = audio_tx.send(audio_data).await {
                        log::error!("Failed to send audio chunk to processor: {}", e);
                        break;
                    }
                } else if msg.is_text() {
                    log::debug!("Received text message from client: {}", msg.to_str().unwrap_or_default());
                    // Handle control messages if any
                }
            }
            Err(e) => {
                log::error!("WebSocket receive error: {}", e);
                break;
            }
        }
    }
    log::info!("Client WebSocket disconnected.");
}

fn run_whisper_inference(
    ctx: &WhisperContext,
    state: &mut whisper_rs::WhisperState,
    audio_data: &[f32],
) -> String {
    let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 });
    params.set_print_progress(false);
    params.set_print_special(false);
    params.set_print_realtime(false);
    params.set_print_timestamps(false);
    params.set_language(Some("en"));
    params.set_n_threads(4); // Adjust based on CPU cores

    // Run the inference
    state.full(params, audio_data).expect("Failed to run Whisper inference");

    // Iterate over the segments and collect the text
    let mut result = String::new();
    let num_segments = state.full_n_segments().expect("Failed to get number of segments");
    for i in 0..num_segments {
        let text = state.full_get_segment_text(i).expect("Failed to get segment text");
        result.push_str(&text);
    }
    result.trim().to_string()
}