Merge pull request #14 from second-state/concurrent_tts

 (v2) Concurrent tts

Author: L-jasmine

src/ai/bailian/cosyvoice.rs

@@ -59,8 +59,10 @@ impl Default for CosyVoiceVersion {
 }
 
 impl CosyVoiceTTS {
+    const WEBSOCKET_URL: &str = "wss://dashscope.aliyuncs.com/api-ws/v1/inference";
+
     pub async fn connect(token: String) -> anyhow::Result<Self> {
-        let url = format!("wss://dashscope.aliyuncs.com/api-ws/v1/inference");
+        let url = Self::WEBSOCKET_URL;
 
         let client = reqwest::Client::new();
         let response = client
@@ -79,6 +81,22 @@ impl CosyVoiceTTS {
         })
     }
 
+    pub async fn reconnect(&mut self) -> anyhow::Result<()> {
+        let url = Self::WEBSOCKET_URL;
+
+        let client = reqwest::Client::new();
+        let response = client
+            .get(url)
+            .bearer_auth(&self.token)
+            .header("X-DashScope-DataInspection", "enable")
+            .upgrade()
+            .send()
+            .await?;
+        self.websocket = response.into_websocket().await?;
+        self.synthesis_started = false;
+        Ok(())
+    }
+
     pub async fn start_synthesis(
         &mut self,
         model: CosyVoiceVersion,
@@ -222,4 +240,23 @@ async fn test_cosyvoice_tts() {
     };
     let wav = crate::util::pcm_to_wav(&audio_data, config);
     std::fs::write("./resources/test/cosyvoice_out.wav", wav).unwrap();
+
+    let text = "Hello, this is CosyVoice V2";
+    tts.start_synthesis(CosyVoiceVersion::V2, None, Some(24000), text)
+        .await
+        .unwrap();
+
+    let mut audio_data = bytes::BytesMut::new();
+    while let Ok(Some(chunk)) = tts.next_audio_chunk().await {
+        audio_data.extend_from_slice(&chunk);
+    }
+
+    println!("Audio data size: {} bytes", audio_data.len());
+    let config = crate::util::WavConfig {
+        channels: 1,
+        sample_rate: 24000,
+        bits_per_sample: 16,
+    };
+    let wav = crate::util::pcm_to_wav(&audio_data, config);
+    std::fs::write("./resources/test/cosyvoice_out2.wav", wav).unwrap();
 }

src/ai/bailian/realtime_asr.rs

@@ -75,6 +75,21 @@ impl ParaformerRealtimeV2Asr {
         })
     }
 
+    pub async fn reconnect(&mut self) -> anyhow::Result<()> {
+        let url = format!("wss://dashscope.aliyuncs.com/api-ws/v1/inference");
+
+        let client = reqwest::Client::new();
+        let response = client
+            .get(url)
+            .bearer_auth(&self.token)
+            .header("X-DashScope-DataInspection", "enable")
+            .upgrade()
+            .send()
+            .await?;
+        self.websocket = response.into_websocket().await?;
+        Ok(())
+    }
+
     pub async fn start_pcm_recognition(&mut self) -> anyhow::Result<()> {
         let task_id = Uuid::new_v4().to_string();
         log::info!("Starting asr task with ID: {}", task_id);
@@ -110,6 +125,13 @@ impl ParaformerRealtimeV2Asr {
                     log::debug!("Received message: {:?}", text);
 
                     let response: ResponseMessage = serde_json::from_str(&text)?;
+                    if response.header.task_id != self.task_id {
+                        log::warn!(
+                            "Received message for different task_id: {}",
+                            response.header.task_id
+                        );
+                        continue;
+                    }
 
                     if response.is_task_started() {
                         log::info!("Recognition task started");
@@ -187,6 +209,7 @@ impl ParaformerRealtimeV2Asr {
     }
 }
 
+// cargo test --package echokit_server --bin echokit_server -- ai::bailian::realtime_asr::test_paraformer_asr --exact --show-output
 #[tokio::test]
 async fn test_paraformer_asr() {
     env_logger::init();
@@ -202,6 +225,21 @@ async fn test_paraformer_asr() {
         .unwrap();
     asr.start_pcm_recognition().await.unwrap();
 
+    asr.send_audio(audio_data.clone()).await.unwrap();
+    asr.finish_task().await.unwrap();
+
+    loop {
+        if let Ok(Some(sentence)) = asr.next_result().await {
+            println!("{:?}", sentence);
+            if sentence.sentence_end {
+                println!();
+            }
+        } else {
+            break;
+        }
+    }
+
+    asr.start_pcm_recognition().await.unwrap();
     asr.send_audio(audio_data).await.unwrap();
     asr.finish_task().await.unwrap();
 

src/ai/elevenlabs/tts.rs

@@ -98,13 +98,21 @@ impl ElevenlabsTTS {
         token: String,
         voice: String,
         output_format: OutputFormat,
+    ) -> anyhow::Result<Self> {
+        let client = reqwest::Client::new();
+        Self::new_with_client(&client, token, voice, output_format).await
+    }
+
+    pub async fn new_with_client(
+        client: &reqwest::Client,
+        token: String,
+        voice: String,
+        output_format: OutputFormat,
     ) -> anyhow::Result<Self> {
         let url = format!(
             "wss://api.elevenlabs.io/v1/text-to-speech/{voice}/stream-input?model_id={MODEL_ID}&output_format={output_format}",
         );
 
-        let client = reqwest::Client::new();
-
         let response = client
             .get(url)
             .header("xi-api-key", &token)
@@ -203,6 +211,7 @@ impl ElevenlabsTTS {
     }
 }
 
+// cargo test --package echokit_server --bin echokit_server -- ai::elevenlabs::tts::test_elevenlabs_tts --exact --show-output
 #[tokio::test]
 async fn test_elevenlabs_tts() {
     env_logger::init();

src/ai/tts.rs

@@ -2,13 +2,13 @@ use bytes::Bytes;
 
 /// return: wav_audio: 16bit,32k,single-channel.
 pub async fn gsv(
+    client: &reqwest::Client,
     tts_url: &str,
     speaker: &str,
     text: &str,
     sample_rate: Option<usize>,
 ) -> anyhow::Result<Bytes> {
     log::debug!("speaker: {speaker}, text: {text}");
-    let client = reqwest::Client::new();
     let res = client
         .post(tts_url)
         .json(&serde_json::json!({"speaker": speaker, "input": text, "sample_rate": sample_rate}))
@@ -34,7 +34,10 @@ async fn test_gsv() {
     let tts_url = "http://localhost:8000/v1/audio/speech";
     let speaker = "ad";
     let text = "你好，我是胡桃";
-    let wav_audio = gsv(tts_url, speaker, text, Some(16000)).await.unwrap();
+    let client = reqwest::Client::new();
+    let wav_audio = gsv(&client, tts_url, speaker, text, Some(16000))
+        .await
+        .unwrap();
     let header = hound::WavReader::new(wav_audio.as_ref()).unwrap();
     let spec = header.spec();
     println!("wav header: {:?}", spec);
@@ -44,13 +47,13 @@ async fn test_gsv() {
 
 /// return: pcm_chunk: 16bit,32k,single-channel.
 pub async fn stream_gsv(
+    client: &reqwest::Client,
     tts_url: &str,
     speaker: &str,
     text: &str,
     sample_rate: Option<usize>,
 ) -> anyhow::Result<reqwest::Response> {
     log::debug!("speaker: {speaker}, text: {text}");
-    let client = reqwest::Client::new();
     let res = client
         .post(tts_url)
         .json(&serde_json::json!({"speaker": speaker, "input": text, "sample_rate": sample_rate}))
@@ -69,9 +72,14 @@ pub async fn stream_gsv(
 }
 
 /// return: wav_audio: 16bit,48k,single-channel.
-pub async fn groq(model: &str, token: &str, voice: &str, text: &str) -> anyhow::Result<Bytes> {
+pub async fn groq(
+    client: &reqwest::Client,
+    model: &str,
+    token: &str,
+    voice: &str,
+    text: &str,
+) -> anyhow::Result<Bytes> {
     log::debug!("groq tts. voice: {voice}, text: {text}");
-    let client = reqwest::Client::new();
     let res = client
         .post("https://api.groq.com/openai/v1/audio/speech")
         .bearer_auth(token)
@@ -102,7 +110,10 @@ async fn test_groq() {
     let token = std::env::var("GROQ_API_KEY").unwrap();
     let speaker = "Aaliyah-PlayAI";
     let text = "你好，我是胡桃";
-    let wav_audio = groq("playai-tts", &token, speaker, text).await.unwrap();
+    let client = reqwest::Client::new();
+    let wav_audio = groq(&client, "playai-tts", &token, speaker, text)
+        .await
+        .unwrap();
     let mut reader = wav_io::reader::Reader::from_vec(wav_audio.to_vec()).unwrap();
     let head = reader.read_header().unwrap();
     println!("wav header: {:?}", head);

src/main.rs

@@ -54,16 +54,15 @@ async fn routes(
     let mut tool_set = ai::openai::tool::ToolSet::default();
     let mut real_config: Option<StableRealtimeConfig> = None;
     match &config.config {
-        config::AIConfig::Stable {
-            llm,
-            tts,
-            asr: ASRConfig::Whisper(asr),
-        } => {
-            real_config = Some(StableRealtimeConfig {
-                llm: llm.clone(),
-                tts: tts.clone(),
-                asr: asr.clone(),
-            });
+        config::AIConfig::Stable { llm, tts, asr } => {
+            if let ASRConfig::Whisper(asr) = asr {
+                real_config = Some(StableRealtimeConfig {
+                    llm: llm.clone(),
+                    tts: tts.clone(),
+                    asr: asr.clone(),
+                });
+            }
+
             for server in &llm.mcp_server {
                 match server.type_ {
                     config::MCPType::SSE => {
@@ -87,22 +86,53 @@ async fn routes(
         _ => {}
     }
 
+    let record_config = Arc::new(services::ws_record::WsRecordSetting {
+        record_callback_url: config.record.callback_url,
+    });
+
+    let ws_setting = Arc::new(services::ws::WsSetting::new(
+        hello_wav.clone(),
+        config.config.clone(),
+        tool_set.clone(),
+    ));
+
     let mut router = Router::new()
         // .route("/", get(handler))
-        .route("/ws/{id}", any(services::mixed_handler))
+        .route("/v1/ws/{id}", any(services::mixed_handler))
         .route("/v1/chat/{id}", any(services::ws::ws_handler))
         .route("/v1/record/{id}", any(services::ws_record::ws_handler))
         .nest("/downloads", services::file::new_file_service("./record"))
-        .layer(axum::Extension(Arc::new(services::ws::WsSetting::new(
-            hello_wav,
-            config.config,
-            tool_set,
-        ))))
-        .layer(axum::Extension(Arc::new(
-            services::ws_record::WsRecordSetting {
-                record_callback_url: config.record.callback_url,
-            },
-        )));
+        .layer(axum::Extension(ws_setting.clone()))
+        .layer(axum::Extension(record_config.clone()));
+
+    if let config::AIConfig::Stable { llm, tts, asr } = config.config {
+        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+        // let tool_set = tool_set;
+        tokio::spawn(async move {
+            if let Err(e) =
+                crate::services::ws::stable::run_session_manager(&llm, &tts, &asr, &tool_set, rx)
+                    .await
+            {
+                log::error!("Stable session manager exited with error: {}", e);
+            }
+        });
+
+        router = router
+            .route("/ws/{id}", any(services::v2_mixed_handler))
+            .route("/v2/stable_ws/{id}", any(services::ws::stable::ws_handler))
+            .layer(axum::Extension(Arc::new(
+                services::ws::stable::StableWsSetting {
+                    sessions: tx,
+                    hello_wav,
+                },
+            )))
+            .layer(axum::Extension(record_config.clone()));
+    } else {
+        router = router
+            .route("/ws/{id}", any(services::mixed_handler))
+            .layer(axum::Extension(ws_setting.clone()))
+            .layer(axum::Extension(record_config.clone()));
+    }
 
     if let Some(real_config) = real_config {
         log::info!(

src/services/mod.rs

@@ -43,3 +43,28 @@ pub async fn mixed_handler(
         .into_response()
     }
 }
+
+pub async fn v2_mixed_handler(
+    Extension(record_setting): Extension<Arc<ws_record::WsRecordSetting>>,
+    Extension(pool): Extension<Arc<ws::stable::StableWsSetting>>,
+    ws: WebSocketUpgrade,
+    Path(id): Path<String>,
+    Query(params): Query<ConnectQueryParams>,
+) -> Response {
+    if params.record {
+        ws_record::ws_handler(Extension(record_setting), ws, Path(id))
+            .await
+            .into_response()
+    } else {
+        ws::stable::ws_handler(
+            Extension(pool),
+            ws,
+            Path(id),
+            Query(ws::ConnectQueryParams {
+                reconnect: params.reconnect,
+            }),
+        )
+        .await
+        .into_response()
+    }
+}

src/services/realtime_ws.rs

@@ -1169,9 +1169,11 @@ async fn tts_and_send(
     item_id: Option<String>,
     text: String,
 ) -> anyhow::Result<()> {
+    let client = reqwest::Client::new();
     match tts_config {
         crate::config::TTSConfig::Stable(tts) => {
-            let wav_data = crate::ai::tts::gsv(&tts.url, &tts.speaker, &text, Some(24000)).await?;
+            let wav_data =
+                crate::ai::tts::gsv(&client, &tts.url, &tts.speaker, &text, Some(24000)).await?;
             let duration_sec = send_wav(tx, response_id, item_id, text, wav_data).await?;
             log::info!("Stable TTS duration: {:?}", duration_sec);
             Ok(())
@@ -1184,13 +1186,15 @@ async fn tts_and_send(
         }
         crate::config::TTSConfig::Groq(groq) => {
             let wav_data =
-                crate::ai::tts::groq(&groq.model, &groq.api_key, &groq.voice, &text).await?;
+                crate::ai::tts::groq(&client, &groq.model, &groq.api_key, &groq.voice, &text)
+                    .await?;
             let duration_sec = send_wav(tx, response_id, item_id, text, wav_data).await?;
             log::info!("Groq TTS duration: {:?}", duration_sec);
             Ok(())
         }
         crate::config::TTSConfig::StreamGSV(stream_tts) => {
             let resp = crate::ai::tts::stream_gsv(
+                &client,
                 &stream_tts.url,
                 &stream_tts.speaker,
                 &text,