feat: support opus

Author: L-jasmine

Cargo.lock

@@ -23,8 +23,11 @@ axum-extra = { version = "0.10.0", features = ["typed-header"] }
 tokio = { version = "1", features = ["full"] }
 
 reqwest = { version = "0.12", features = ["multipart", "json", "stream"] }
+
 hound = "3.5.1"
 wav_io = "0.1.15"
+opus = "0.3.0"
+
 rand = "0.9.0"
 uuid = { version = "1.14", features = [
     "v4",       # Lets you generate random UUIDs

Cargo.toml

@@ -17,6 +17,8 @@ pub struct ConnectQueryParams {
     reconnect: bool,
     #[serde(default)]
     record: bool,
+    #[serde(default)]
+    opus: bool,
 }
 
 pub async fn mixed_handler(
@@ -37,6 +39,7 @@ pub async fn mixed_handler(
             Path(id),
             Query(ws::ConnectQueryParams {
                 reconnect: params.reconnect,
+                opus: params.opus,
             }),
         )
         .await
@@ -62,6 +65,7 @@ pub async fn v2_mixed_handler(
             Path(id),
             Query(ws::ConnectQueryParams {
                 reconnect: params.reconnect,
+                opus: params.opus,
             }),
         )
         .await

src/services/mod.rs

@@ -80,6 +80,8 @@ impl WsSetting {
 pub struct ConnectQueryParams {
     #[serde(default)]
     pub reconnect: bool,
+    #[serde(default)]
+    pub opus: bool,
 }
 
 pub async fn ws_handler(
@@ -929,22 +931,35 @@ async fn process_socket_io(
     rx: &mut WsRx,
     audio_tx: ClientTx,
     socket: &mut WebSocket,
+    enable_opus: bool,
 ) -> anyhow::Result<()> {
+    let mut opus_encoder =
+        opus::Encoder::new(SAMPLE_RATE, opus::Channels::Mono, opus::Application::Voip)
+            .map_err(|e| anyhow::anyhow!("opus encoder error: {e}"))?;
+    let mut ret_audio = Vec::new();
+
     loop {
         let r = tokio::select! {
             cmd = rx.recv() => {
                 cmd.map(|cmd| WsEvent::Command(cmd))
             }
             message = socket.recv() => {
-                message.map(|message| match message{
+                message.map(|message| match message {
                     Ok(message) => WsEvent::Message(Ok(message)),
                     Err(e) => WsEvent::Message(Err(anyhow::anyhow!("recv ws error: {e}"))),
                 })
             }
         };
 
         match r {
-            Some(WsEvent::Command(cmd)) => process_command(socket, cmd).await?,
+            Some(WsEvent::Command(cmd)) => {
+                if enable_opus {
+                    process_command_with_opus(socket, cmd, &mut opus_encoder, &mut ret_audio)
+                        .await?
+                } else {
+                    process_command(socket, cmd).await?
+                }
+            }
             Some(WsEvent::Message(Ok(msg))) => match process_message(msg) {
                 ProcessMessageResult::Audio(d) => audio_tx
                     .send(ClientMsg::AudioChunk(d))
@@ -1213,14 +1228,20 @@ async fn handle_socket(
     }
 
     log::info!("`{}` starting socket io processing", id);
-    process_socket_io(&mut cmd_rx, audio_tx, &mut socket).await?;
+    process_socket_io(&mut cmd_rx, audio_tx, &mut socket, connect_params.opus).await?;
 
     Ok(())
 }
 
 pub const SAMPLE_RATE: u32 = 16000;
 pub const SAMPLE_RATE_BUFFER_SIZE: usize = 2 * (SAMPLE_RATE as usize) / 10;
 
+pub const fn sample_120ms(sample_rate: u32) -> usize {
+    (sample_rate as usize) * 12 / 100
+}
+
+pub const SAMPLE_RATE_120MS: usize = sample_120ms(SAMPLE_RATE);
+
 async fn process_command(ws: &mut WebSocket, cmd: WsCommand) -> anyhow::Result<()> {
     match cmd {
         WsCommand::AsrResult(texts) => {
@@ -1272,6 +1293,86 @@ async fn process_command(ws: &mut WebSocket, cmd: WsCommand) -> anyhow::Result<(
     Ok(())
 }
 
+async fn process_command_with_opus(
+    ws: &mut WebSocket,
+    cmd: WsCommand,
+    opus_encode: &mut opus::Encoder,
+    ret_audio: &mut Vec<i16>,
+) -> anyhow::Result<()> {
+    match cmd {
+        WsCommand::AsrResult(texts) => {
+            let asr = rmp_serde::to_vec(&crate::protocol::ServerEvent::ASR {
+                text: texts.join("\n"),
+            })
+            .expect("Failed to serialize ASR ServerEvent");
+            ws.send(Message::binary(asr)).await?;
+        }
+
+        WsCommand::Action { action } => {
+            let action = rmp_serde::to_vec(&crate::protocol::ServerEvent::Action { action })
+                .expect("Failed to serialize Action ServerEvent");
+            ws.send(Message::binary(action)).await?;
+        }
+        WsCommand::StartAudio(text) => {
+            log::trace!("StartAudio: {text:?}");
+            opus_encode
+                .reset_state()
+                .map_err(|e| anyhow::anyhow!("opus reset state error: {e}"))?;
+            let start_audio = rmp_serde::to_vec(&crate::protocol::ServerEvent::StartAudio { text })
+                .expect("Failed to serialize StartAudio ServerEvent");
+            ws.send(Message::binary(start_audio)).await?;
+        }
+        WsCommand::Audio(data) => {
+            log::trace!("Audio chunk size: {}", data.len());
+            for chunk in data.chunks_exact(2) {
+                let sample = i16::from_le_bytes([chunk[0], chunk[1]]);
+                ret_audio.push(sample);
+            }
+
+            // 120ms per chunk
+            for chunk in ret_audio.chunks(sample_120ms(SAMPLE_RATE)) {
+                if chunk.len() < sample_120ms(SAMPLE_RATE) {
+                    *ret_audio = chunk.to_vec();
+                    break;
+                }
+                let data = opus_encode.encode_vec(chunk, 2 * sample_120ms(SAMPLE_RATE) / 3)?;
+
+                let audio_chunk =
+                    rmp_serde::to_vec(&crate::protocol::ServerEvent::AudioChunk { data })
+                        .expect("Failed to serialize AudioChunk ServerEvent");
+                ws.send(Message::binary(audio_chunk)).await?;
+            }
+        }
+        WsCommand::EndAudio => {
+            log::trace!("EndAudio");
+            if !ret_audio.is_empty() {
+                let padded_audio_len = sample_120ms(SAMPLE_RATE) - ret_audio.len();
+                ret_audio.extend(vec![0i16; padded_audio_len]);
+                let data = opus_encode.encode_vec(&ret_audio, 2 * sample_120ms(SAMPLE_RATE) / 3)?;
+                let audio_chunk =
+                    rmp_serde::to_vec(&crate::protocol::ServerEvent::AudioChunk { data })
+                        .expect("Failed to serialize AudioChunk ServerEvent");
+                log::info!("Sending final audio chunk of size: {}", audio_chunk.len());
+                ws.send(Message::binary(audio_chunk)).await?;
+                ret_audio.clear();
+            }
+            let end_audio = rmp_serde::to_vec(&crate::protocol::ServerEvent::EndAudio)
+                .expect("Failed to serialize EndAudio ServerEvent");
+            ws.send(Message::binary(end_audio)).await?;
+        }
+        WsCommand::Video(_) => {
+            log::warn!("video command is not implemented yet");
+        }
+        WsCommand::EndResponse => {
+            log::debug!("EndResponse");
+            let end_response = rmp_serde::to_vec(&crate::protocol::ServerEvent::EndResponse)
+                .expect("Failed to serialize JsonCommand");
+            ws.send(Message::binary(end_response)).await?;
+        }
+    }
+    Ok(())
+}
+
 enum ProcessMessageResult {
     Audio(Bytes),
     Submit,

src/services/ws.rs

@@ -68,7 +68,7 @@ async fn handle_socket(
         })
         .map_err(|e| anyhow::anyhow!("send session error: {}", e))?;
 
-    super::process_socket_io(&mut cmd_rx, client_tx, &mut socket).await?;
+    super::process_socket_io(&mut cmd_rx, client_tx, &mut socket, params.opus).await?;
     Ok(())
 }