re-org code structure

Cargo.lock

@@ -696,9 +696,9 @@ dependencies = [
 
 [[package]]
 name = "config"
-version = "0.13.3"
+version = "0.13.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d379af7f68bfc21714c6c7dea883544201741d2ce8274bb12fa54f89507f52a7"
+checksum = "23738e11972c7643e4ec947840fc463b6a571afcd3e735bdfce7d03c7a784aca"
 dependencies = [
  "async-trait",
  "json5",
@@ -838,9 +838,9 @@ checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
 
 [[package]]
 name = "form_urlencoded"
-version = "1.2.0"
+version = "1.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a62bc1cf6f830c2ec14a513a9fb124d0a213a629668a4186f329db21fe045652"
+checksum = "e13624c2627564efccf4934284bdd98cbaa14e79b0b5a141218e507b3a823456"
 dependencies = [
  "percent-encoding",
 ]
@@ -853,9 +853,9 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
 
 [[package]]
 name = "futures"
-version = "0.3.28"
+version = "0.3.29"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "23342abe12aba583913b2e62f22225ff9c950774065e4bfb61a19cd9770fec40"
+checksum = "da0290714b38af9b4a7b094b8a37086d1b4e61f2df9122c3cad2577669145335"
 dependencies = [
  "futures-channel",
  "futures-core",
@@ -1179,9 +1179,9 @@ dependencies = [
 
 [[package]]
 name = "idna"
-version = "0.4.0"
+version = "0.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7d20d6b07bfbc108882d88ed8e37d39636dcc260e15e30c45e6ba089610b917c"
+checksum = "634d9b1461af396cad843f47fdba5597a4f9e6ddd4bfb6ff5d85028c25cb12f6"
 dependencies = [
  "unicode-bidi",
  "unicode-normalization",
@@ -1487,9 +1487,9 @@ checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099"
 
 [[package]]
 name = "percent-encoding"
-version = "2.3.0"
+version = "2.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9b2a4787296e9989611394c33f193f676704af1686e70b8f8033ab5ba9a35a94"
+checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e"
 
 [[package]]
 name = "pest"
@@ -2535,9 +2535,9 @@ checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
 
 [[package]]
 name = "url"
-version = "2.4.1"
+version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "143b538f18257fac9cad154828a57c6bf5157e1aa604d4816b5995bf6de87ae5"
+checksum = "31e6302e3bb753d46e83516cae55ae196fc0c309407cf11ab35cc51a4c2a4633"
 dependencies = [
  "form_urlencoded",
  "idna",

Cargo.toml

@@ -3,6 +3,14 @@ name = "polyhedron"
 version = "0.1.0"
 edition = "2021"
 
+[lib]
+name = "polyhedron"
+path = "src/lib.rs"
+
+[[bin]]
+name = "polyhedron"
+path = "src/main.rs"
+
 [dependencies]
 anyhow = "1.0"
 async-stream = "0.3"

src/asr/aws.rs

@@ -12,6 +12,7 @@ use tokio::select;
 use tokio::sync::broadcast::Receiver;
 use tokio_stream::Stream;
 use futures_util::TryStreamExt;
+use tracing::trace;
 use crate::asr::{ASR, Event, slice_i16_to_u8};
 
 pub struct AwsAsr {
@@ -58,6 +59,7 @@ impl AwsAsr {
                 let output_stream = to_stream(output);
                 output_stream
                     .try_for_each(|text| async {
+                        trace!("here");
                         let _ = shared_speaker_transcript.send(text);
                         Ok(())
                     })

src/asr/mod.rs

@@ -18,8 +18,7 @@ pub(crate) trait ASR {
 }
 
 
-#[allow(dead_code)]
-pub(crate) fn slice_i16_to_u8(slice: &[i16]) -> Vec<u8> {
+pub fn slice_i16_to_u8(slice: &[i16]) -> Vec<u8> {
     slice
         .iter()
         .flat_map(|&sample| {

src/lesson.rs

@@ -1,9 +1,7 @@
 use aws_config::SdkConfig;
-use aws_sdk_polly::primitives::ByteStream;
 use aws_sdk_polly::types::{Engine, OutputFormat, SpeechMarkType, VoiceId};
 use aws_sdk_transcribestreaming::types::{LanguageCode};
 use futures_util::future::try_join;
-use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 use std::fmt::{Debug, Formatter};
 use std::io::BufRead;
@@ -17,11 +15,14 @@ use crate::asr::{Event, aws::AwsAsr, ASR};
 
 #[cfg(feature = "whisper")]
 use crate::asr::whisper::WhisperAsr;
+use crate::Viseme;
+
+pub type LessonID = String;
 
 pub struct InnerLessonsManager {
     translate_client: aws_sdk_translate::Client,
     polly_client: aws_sdk_polly::Client,
-    lessons: Arc<RwLock<BTreeMap<u32, Lesson>>>,
+    lessons: Arc<RwLock<BTreeMap<LessonID, Lesson>>>,
 }
 
 #[derive(Clone)]
@@ -43,7 +44,7 @@ impl Deref for LessonsManager {
     }
 }
 
-pub(crate) enum AsrEngine {
+pub enum AsrEngine {
     AWS,
     #[allow(dead_code)]
     #[cfg(feature = "whisper")]
@@ -69,7 +70,7 @@ impl AsrEngine {
 }
 
 impl LessonsManager {
-    pub(crate) fn new(sdk_config: &SdkConfig) -> Self {
+    pub fn new(sdk_config: &SdkConfig) -> Self {
         let translate_client = aws_sdk_translate::Client::new(sdk_config);
         let polly_client = aws_sdk_polly::Client::new(sdk_config);
         let inner = InnerLessonsManager {
@@ -80,14 +81,14 @@ impl LessonsManager {
         LessonsManager { inner: Arc::new(inner) }
     }
 
-    pub(crate) async fn create_lesson(&self, id: u32, engine: AsrEngine, speaker_lang: LanguageCode) -> Lesson {
+    pub(crate) async fn create_lesson(&self, id: LessonID, engine: AsrEngine, speaker_lang: LanguageCode) -> Lesson {
         let mut map = self.lessons.write().await;
         let lesson: Lesson = InnerLesson::new(self.clone(), engine, speaker_lang).await.into();
         map.insert(id, lesson.clone());
         lesson
     }
 
-    pub(crate) async fn get_lesson(&self, id: u32) -> Option<Lesson> {
+    pub(crate) async fn get_lesson(&self, id: LessonID) -> Option<Lesson> {
         let map = self.lessons.read().await;
         map.get(&id).cloned()
     }
@@ -434,12 +435,6 @@ impl Drop for InnerVoiceLesson {
 }
 
 
-// {"time":180,"type":"viseme","value":"r"}
-#[derive(Debug, Deserialize, Clone, Serialize)]
-pub(crate) struct Viseme {
-    time: u32,
-    value: String,
-}
 
 #[derive(Debug)]
 enum SynthesizeError {

src/lib.rs

@@ -0,0 +1,400 @@
+#[cfg(feature = "whisper")]
+extern crate whisper;
+
+pub mod config;
+pub mod lesson;
+pub mod asr;
+pub mod base64box;
+
+use aws_config::{SdkConfig};
+use aws_sdk_transcribestreaming::types::LanguageCode;
+use futures_util::{stream::StreamExt, SinkExt};
+use poem::{
+    endpoint::{StaticFileEndpoint, StaticFilesEndpoint},
+    get, handler,
+    listener::TcpListener,
+    web::{
+        websocket::{Message, WebSocket},
+        Data, Query,
+    },
+    EndpointExt, IntoResponse, Route, Server,
+};
+use serde::{Deserialize, Serialize};
+use tokio::select;
+use tracing::{debug, span};
+
+use crate::base64box::Base64Box;
+use crate::config::SETTINGS;
+use crate::lesson::{AsrEngine, LessonID, LessonsManager};
+
+#[derive(Clone)]
+pub struct Context {
+    lessons_manager: LessonsManager,
+}
+
+impl Context {
+    pub fn new(config: &SdkConfig) -> Self {
+        Self {
+            lessons_manager: LessonsManager::new(config),
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub enum SingleEvent {
+    #[serde(rename = "original")]
+    Transcription {
+        content: String,
+        #[serde(rename = "isFinal")]
+        is_final: bool
+    },
+    #[serde(rename = "translated")]
+    Translation { content: String },
+    #[serde(rename = "voice")]
+    Voice {
+        content: Base64Box
+    },
+}
+
+
+#[derive(Deserialize, Debug)]
+pub struct LessonListenerQuery {
+    id: String,
+    language: String,
+    voice: String,
+}
+
+#[derive(Serialize, Debug)]
+#[serde(tag = "type")]
+pub enum LiveLessonTextEvent {
+    #[serde(rename = "original")]
+    Transcription {
+        content: String,
+        #[serde(rename = "isFinal")]
+        is_final: bool
+    },
+    Translation { content: String },
+    LipSync { visemes: Vec<Viseme> },
+}
+
+
+// {"time":180,"type":"viseme","value":"r"}
+#[derive(Debug, Deserialize, Clone, Serialize)]
+pub struct Viseme {
+    time: u32,
+    value: String,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct LessonSpeakerQuery {
+    id: LessonID,
+    language: String,
+    #[allow(dead_code)] // TODO: use this in the future
+    prompt: Option<String>,
+}
+
+#[handler]
+async fn stream_speaker(
+    ctx: Data<&Context>,
+    query: Query<LessonSpeakerQuery>,
+    ws: WebSocket,
+) -> impl IntoResponse {
+    let lessons_manager = ctx.lessons_manager.clone();
+    ws.on_upgrade(|mut socket| async move {
+        let Ok(lang) = query.language.parse::<LanguageCode>() else {
+            let _ = socket
+                .send(Message::Text(format!("invalid language code: {}", query.language)))
+                .await;
+            return
+        };
+        let lesson = lessons_manager
+            .create_lesson(
+                query.id.clone(),
+                AsrEngine::AWS,
+                lang,
+            )
+            .await;
+
+        let mut transcribe_rx = lesson.transcript_channel();
+        let fut = async {
+            loop {
+                select! {
+                    msg = socket.next() => {
+                        let Some(res) = msg else { break };
+                        let msg = res?;
+                        if msg.is_close() {
+                            break
+                        }
+                        let Message::Binary(bin) = msg else {
+                            tracing::warn!("Other: {:?}", msg);
+                            continue
+                        };
+                        let frame = u8_to_i16(&bin);
+                        lesson.send(frame).await?
+                    },
+                    output = transcribe_rx.recv() => {
+                        let evt = output?;
+                        if evt.is_final {
+                            tracing::trace!("Transcribed: {}", evt.transcript);
+                        }
+                        let evt = LiveLessonTextEvent::Transcription { content: evt.transcript, is_final: evt.is_final };
+                        let Ok(json) = serde_json::to_string(&evt) else {
+                            tracing::warn!("failed to serialize json: {:?}", evt);
+                            continue
+                        };
+                        socket.send(Message::Text(json)).await?
+                    },
+                }
+            }
+            Ok(())
+        };
+
+        let span = span!(tracing::Level::TRACE, "lesson_speaker", lesson_id = query.id);
+        let _ = span.enter();
+        let res: anyhow::Result<()> = fut.await;
+        match res {
+            Ok(()) => {
+                tracing::info!("lesson speaker closed");
+            }
+            Err(e) => {
+                tracing::warn!("lesson speaker error: {}", e);
+            }
+        }
+    })
+}
+
+#[handler]
+async fn stream_listener(
+    ctx: Data<&Context>,
+    query: Query<LessonListenerQuery>,
+    ws: WebSocket,
+) -> impl IntoResponse {
+    let lessons_manager = ctx.lessons_manager.clone();
+
+    ws.on_upgrade(|mut socket| async move {
+        let voice_id = match query.voice.parse() {
+            Ok(id) => id,
+            Err(e) => {
+                let _ = socket
+                    .send(Message::Text(format!("invalid voice id: {}", e)))
+                    .await;
+                return;
+            }
+        };
+
+        let lesson_opt = lessons_manager.get_lesson(query.id.clone()).await;
+        debug!("listener param = {:?}", query);
+        let Some(lesson) = lesson_opt else {
+            let _ = socket
+                .send(Message::Text("lesson not found".to_string()))
+                .await;
+            return;
+        };
+        let mut transcript_rx = lesson.transcript_channel();
+        let mut lang_lesson = lesson.get_or_init(query.language.clone()).await;
+        let mut translate_rx = lang_lesson.translated_channel();
+        let mut voice_lesson = lang_lesson.get_or_init(voice_id).await;
+        let mut voice_rx = voice_lesson.voice_channel();
+        let mut lip_sync_rx = voice_lesson.lip_sync_channel();
+
+        let fut = async {
+            loop {
+                select! {
+                    transcript_poll = transcript_rx.recv() => {
+                        let transcript = transcript_poll?;
+                        let evt = LiveLessonTextEvent::Transcription {
+                            content: transcript.transcript,
+                            is_final: transcript.is_final
+                        };
+                        let Ok(json) = serde_json::to_string(&evt) else {
+                            tracing::warn!("failed to serialize: {:?}", evt);
+                            continue
+                        };
+                        tracing::debug!("Transcribed: {}", json);
+                        socket.send(Message::Text(json)).await?
+                    },
+                    translated_poll = translate_rx.recv() => {
+                        let translated = translated_poll?;
+                        let evt = LiveLessonTextEvent::Translation { content: translated };
+                        let Ok(json) = serde_json::to_string(&evt) else {
+                            tracing::warn!("failed to serialize: {:?}", evt);
+                            continue
+                        };
+                        tracing::debug!("Translated: {}", json);
+                        socket.send(Message::Text(json)).await?
+                    },
+                    voice_poll = voice_rx.recv() => {
+                        let voice = voice_poll?;
+                        socket.send(Message::Binary(voice)).await?
+                    },
+                    visemes_poll = lip_sync_rx.recv() => {
+                        let visemes = visemes_poll?;
+                        let evt = LiveLessonTextEvent::LipSync { visemes };
+                        let Ok(json) = serde_json::to_string(&evt) else {
+                            tracing::warn!("failed to serialize: {:?}", evt);
+                            continue
+                        };
+                        socket.send(Message::Text(json)).await?
+                    },
+                }
+            }
+        };
+
+        let span = span!(tracing::Level::TRACE, "lesson_listener", lesson_id = query.id);
+        let _ = span.enter();
+        let res: anyhow::Result<()> = fut.await;
+        match res {
+            Ok(()) => {
+                tracing::info!("lesson listener closed");
+            }
+            Err(e) => {
+                tracing::warn!("lesson listener error: {}", e);
+            }
+        }
+    })
+}
+
+#[derive(Deserialize, Debug)]
+pub struct SingleQuery {
+    id: LessonID,
+    from: String,
+    to: String,
+    voice: Option<String>,
+}
+
+
+#[handler]
+pub async fn stream_single(
+    ctx: Data<&Context>,
+    query: Query<SingleQuery>,
+    ws: WebSocket
+) -> impl IntoResponse {
+    let lessons_manager = ctx.lessons_manager.clone();
+    ws.on_upgrade(|mut socket| async move {
+        let Ok(lang) = query.from.parse::<LanguageCode>() else {
+            let _ = socket
+                .send(Message::Text(format!("invalid language code: {}", query.from)))
+                .await;
+            return
+        };
+        let lesson = lessons_manager
+            .create_lesson(
+                query.id.clone(),
+                AsrEngine::AWS,
+                lang,
+            )
+            .await;
+
+        let mut transcribe_rx = lesson.transcript_channel();
+        let mut lang_lesson = lesson.get_or_init(query.to.clone()).await;
+        let mut translate_rx = lang_lesson.translated_channel();
+        let Ok(voice_id) = query.voice.as_deref().unwrap_or("Amy").parse() else {
+            let _ = socket
+                .send(Message::Text(format!("invalid voice id: {:?}", query.voice)))
+                .await;
+            return
+        };
+        let mut voice_lesson = lang_lesson.get_or_init(voice_id).await;
+        let mut voice_rx = voice_lesson.voice_channel();
+        // let mut lip_sync_rx = voice_lesson.lip_sync_channel();
+
+        let fut = async {
+            loop {
+                let evt = select! {
+                    input = socket.next() => {
+                        let Some(res) = input else { break };
+                        let msg = res?;
+                        if msg.is_close() {
+                            break
+                        }
+                        let Message::Binary(bin) = msg else {
+                            tracing::warn!("Other: {:?}", msg);
+                            continue
+                        };
+                        let frame = u8_to_i16(&bin);
+                        lesson.send(frame).await?;
+                        continue
+                    },
+                    transcript_poll = transcribe_rx.recv() => {
+                        let evt = transcript_poll?;
+                        if evt.is_final {
+                            tracing::trace!("Transcribed: {}", evt.transcript);
+                        }
+                        SingleEvent::Transcription { content: evt.transcript, is_final: evt.is_final }
+                    },
+                    translated_poll = translate_rx.recv() => {
+                        let translated = translated_poll?;
+                        SingleEvent::Translation { content: translated }
+                    },
+                    voice_poll = voice_rx.recv() => {
+                        let voice = voice_poll?;
+                        SingleEvent::Voice { content: Base64Box(voice) }
+                    },
+                };
+
+                let Ok(json) = serde_json::to_string(&evt) else {
+                    tracing::warn!("failed to serialize json: {:?}", evt);
+                    continue
+                };
+                socket.send(Message::Text(json)).await?
+            }
+            Ok(())
+        };
+
+        let span = span!(tracing::Level::TRACE, "lesson_speaker", lesson_id = query.id);
+        let _ = span.enter();
+        let res: anyhow::Result<()> = fut.await;
+        match res {
+            Ok(()) => {
+                tracing::info!("lesson speaker closed");
+            }
+            Err(e) => {
+                tracing::warn!("lesson speaker error: {}", e);
+            }
+        }
+    })
+}
+
+fn u8_to_i16(input: &[u8]) -> Vec<i16> {
+    input
+        .chunks_exact(2)
+        .map(|chunk| {
+            let mut buf = [0u8; 2];
+            buf.copy_from_slice(chunk);
+            i16::from_le_bytes(buf)
+        })
+        .collect::<Vec<i16>>()
+}
+
+pub async fn app(config: &SdkConfig) -> Result<(), std::io::Error> {
+    let ctx = Context {
+        lessons_manager: LessonsManager::new(config),
+    };
+
+    let app = Route::new()
+        .nest(
+            "/",
+            StaticFilesEndpoint::new("./static")
+                .show_files_listing()
+                .index_file("index.html"),
+        )
+        .at("/ws/lesson-speaker", get(stream_speaker))
+        .at("/ws/teacher", get(stream_speaker))
+        .at("/ws/lesson-listener", get(stream_listener))
+        .at("/ws/student", get(stream_listener))
+        .at("/ws/voice", get(stream_single))
+        .at(
+            "lesson-speaker",
+            StaticFileEndpoint::new("./static/index.html"),
+        )
+        .at(
+            "lesson-listener",
+            StaticFileEndpoint::new("./static/index.html"),
+        )
+        .data(ctx);
+    let addr = format!("{}:{}", SETTINGS.server.host, SETTINGS.server.port);
+    let listener = TcpListener::bind(addr);
+    let server = Server::new(listener);
+
+    server.run(app).await
+}

src/main.rs

@@ -3,42 +3,14 @@
  * SPDX-License-Identifier: Apache-2.0.
  */
 
-#![allow(clippy::result_large_err)]
-
-#[cfg(feature = "whisper")]
-extern crate whisper;
-
 use aws_config::BehaviorVersion;
 use aws_sdk_transcribestreaming::meta::PKG_VERSION;
-use aws_sdk_transcribestreaming::types::LanguageCode;
-use futures_util::{stream::StreamExt, SinkExt};
-use poem::{
-    endpoint::{StaticFileEndpoint, StaticFilesEndpoint},
-    get, handler,
-    listener::TcpListener,
-    web::{
-        websocket::{Message, WebSocket},
-        Data, Query,
-    },
-    EndpointExt, IntoResponse, Route, Server,
-};
-use serde::{Deserialize, Serialize};
 use tokio::select;
-use tracing::{debug, span};
+use tracing::{debug};
 use tracing_subscriber::{fmt, prelude::*, EnvFilter};
 
-use crate::{config::*, lesson::*};
-use crate::base64box::Base64Box;
+use polyhedron::app;
 
-mod config;
-mod lesson;
-mod asr;
-mod base64box;
-
-#[derive(Clone)]
-struct Context {
-    lessons_manager: LessonsManager,
-}
 
 #[tokio::main]
 async fn main() -> Result<(), std::io::Error> {
@@ -50,476 +22,13 @@ async fn main() -> Result<(), std::io::Error> {
     debug!("Transcribe client version: {}", PKG_VERSION);
 
     let shared_config = aws_config::load_defaults(BehaviorVersion::latest()).await;
-    let ctx = Context {
-        lessons_manager: LessonsManager::new(&shared_config),
-    };
 
-    let app = Route::new()
-        .nest(
-            "/",
-            StaticFilesEndpoint::new("./static")
-                .show_files_listing()
-                .index_file("index.html"),
-        )
-        .at("/ws/lesson-speaker", get(stream_speaker))
-        .at("/ws/teacher", get(stream_speaker))
-        .at("/ws/lesson-listener", get(stream_listener))
-        .at("/ws/student", get(stream_listener))
-        .at("/ws/voice", get(stream_single))
-        .at(
-            "lesson-speaker",
-            StaticFileEndpoint::new("./static/index.html"),
-        )
-        .at(
-            "lesson-listener",
-            StaticFileEndpoint::new("./static/index.html"),
-        )
-        .data(ctx);
-    let addr = format!("{}:{}", SETTINGS.server.host, SETTINGS.server.port);
-    let listener = TcpListener::bind(addr);
-    let server = Server::new(listener);
 
     select! {
-        res = server.run(app) => res,
+        res = app(&shared_config) => res,
         _ = tokio::signal::ctrl_c() => {
             tracing::info!("Shutting down");
             Ok(())
         },
     }
 }
-
-#[derive(Deserialize, Debug)]
-pub struct LessonSpeakerQuery {
-    id: u32,
-    language: String,
-    #[allow(dead_code)] // TODO: use this in the future
-    prompt: Option<String>,
-}
-
-#[handler]
-async fn stream_speaker(
-    ctx: Data<&Context>,
-    query: Query<LessonSpeakerQuery>,
-    ws: WebSocket,
-) -> impl IntoResponse {
-    let lessons_manager = ctx.lessons_manager.clone();
-    ws.on_upgrade(|mut socket| async move {
-        let Ok(lang) = query.language.parse::<LanguageCode>() else {
-            let _ = socket
-                .send(Message::Text(format!("invalid language code: {}", query.language)))
-                .await;
-            return
-        };
-        let lesson = lessons_manager
-            .create_lesson(
-                query.id,
-                AsrEngine::AWS,
-                lang,
-            )
-            .await;
-
-        let mut transcribe_rx = lesson.transcript_channel();
-        let fut = async {
-            loop {
-                select! {
-                    msg = socket.next() => {
-                        let Some(res) = msg else { break };
-                        let msg = res?;
-                        if msg.is_close() {
-                            break
-                        }
-                        let Message::Binary(bin) = msg else {
-                            tracing::warn!("Other: {:?}", msg);
-                            continue
-                        };
-                        let frame = u8_to_i16(&bin);
-                        lesson.send(frame).await?
-                    },
-                    output = transcribe_rx.recv() => {
-                        let evt = output?;
-                        if evt.is_final {
-                            tracing::trace!("Transcribed: {}", evt.transcript);
-                        }
-                        let evt = LiveLessonTextEvent::Transcription { content: evt.transcript, is_final: evt.is_final };
-                        let Ok(json) = serde_json::to_string(&evt) else {
-                            tracing::warn!("failed to serialize json: {:?}", evt);
-                            continue
-                        };
-                        socket.send(Message::Text(json)).await?
-                    },
-                }
-            }
-            Ok(())
-        };
-
-        let span = span!(tracing::Level::TRACE, "lesson_speaker", lesson_id = query.id);
-        let _ = span.enter();
-        let res: anyhow::Result<()> = fut.await;
-        match res {
-            Ok(()) => {
-                tracing::info!("lesson speaker closed");
-            }
-            Err(e) => {
-                tracing::warn!("lesson speaker error: {}", e);
-            }
-        }
-    })
-}
-
-#[derive(Deserialize, Debug)]
-pub struct LessonListenerQuery {
-    id: u32,
-    language: String,
-    voice: String,
-}
-
-#[derive(Serialize, Debug)]
-#[serde(tag = "type")]
-enum LiveLessonTextEvent {
-    #[serde(rename = "original")]
-    Transcription {
-        content: String,
-        #[serde(rename = "isFinal")]
-        is_final: bool
-    },
-    Translation { content: String },
-    LipSync { visemes: Vec<Viseme> },
-}
-#[handler]
-async fn stream_listener(
-    ctx: Data<&Context>,
-    query: Query<LessonListenerQuery>,
-    ws: WebSocket,
-) -> impl IntoResponse {
-    let lessons_manager = ctx.lessons_manager.clone();
-
-    ws.on_upgrade(|mut socket| async move {
-        let voice_id = match query.voice.parse() {
-            Ok(id) => id,
-            Err(e) => {
-                let _ = socket
-                    .send(Message::Text(format!("invalid voice id: {}", e)))
-                    .await;
-                return;
-            }
-        };
-
-        let lesson_opt = lessons_manager.get_lesson(query.id).await;
-        debug!("listener param = {:?}", query);
-        let Some(lesson) = lesson_opt else {
-            let _ = socket
-                .send(Message::Text("lesson not found".to_string()))
-                .await;
-            return;
-        };
-        let mut transcript_rx = lesson.transcript_channel();
-        let mut lang_lesson = lesson.get_or_init(query.language.clone()).await;
-        let mut translate_rx = lang_lesson.translated_channel();
-        let mut voice_lesson = lang_lesson.get_or_init(voice_id).await;
-        let mut voice_rx = voice_lesson.voice_channel();
-        let mut lip_sync_rx = voice_lesson.lip_sync_channel();
-
-        let fut = async {
-            loop {
-                select! {
-                    transcript_poll = transcript_rx.recv() => {
-                        let transcript = transcript_poll?;
-                        let evt = LiveLessonTextEvent::Transcription {
-                            content: transcript.transcript,
-                            is_final: transcript.is_final
-                        };
-                        let Ok(json) = serde_json::to_string(&evt) else {
-                            tracing::warn!("failed to serialize: {:?}", evt);
-                            continue
-                        };
-                        tracing::debug!("Transcribed: {}", json);
-                        socket.send(Message::Text(json)).await?
-                    },
-                    translated_poll = translate_rx.recv() => {
-                        let translated = translated_poll?;
-                        let evt = LiveLessonTextEvent::Translation { content: translated };
-                        let Ok(json) = serde_json::to_string(&evt) else {
-                            tracing::warn!("failed to serialize: {:?}", evt);
-                            continue
-                        };
-                        tracing::debug!("Translated: {}", json);
-                        socket.send(Message::Text(json)).await?
-                    },
-                    voice_poll = voice_rx.recv() => {
-                        let voice = voice_poll?;
-                        socket.send(Message::Binary(voice)).await?
-                    },
-                    visemes_poll = lip_sync_rx.recv() => {
-                        let visemes = visemes_poll?;
-                        let evt = LiveLessonTextEvent::LipSync { visemes };
-                        let Ok(json) = serde_json::to_string(&evt) else {
-                            tracing::warn!("failed to serialize: {:?}", evt);
-                            continue
-                        };
-                        socket.send(Message::Text(json)).await?
-                    },
-                }
-            }
-        };
-
-        let span = span!(tracing::Level::TRACE, "lesson_listener", lesson_id = query.id);
-        let _ = span.enter();
-        let res: anyhow::Result<()> = fut.await;
-        match res {
-            Ok(()) => {
-                tracing::info!("lesson listener closed");
-            }
-            Err(e) => {
-                tracing::warn!("lesson listener error: {}", e);
-            }
-        }
-    })
-}
-
-#[derive(Serialize, Deserialize, Debug)]
-enum SingleEvent {
-    #[serde(rename = "original")]
-    Transcription {
-        content: String,
-        #[serde(rename = "isFinal")]
-        is_final: bool
-    },
-    #[serde(rename = "translated")]
-    Translation { content: String },
-    #[serde(rename = "voice")]
-    Voice {
-        content: Base64Box
-    },
-}
-
-
-#[derive(Deserialize, Debug)]
-pub struct SingleQuery {
-    id: u32,
-    from: String,
-    to: String,
-    voice: Option<String>,
-}
-
-
-#[handler]
-async fn stream_single(
-    ctx: Data<&Context>,
-    query: Query<SingleQuery>,
-    ws: WebSocket
-) -> impl IntoResponse {
-    let lessons_manager = ctx.lessons_manager.clone();
-    ws.on_upgrade(|mut socket| async move {
-        let Ok(lang) = query.from.parse::<LanguageCode>() else {
-            let _ = socket
-                .send(Message::Text(format!("invalid language code: {}", query.from)))
-                .await;
-            return
-        };
-        let lesson = lessons_manager
-            .create_lesson(
-                query.id,
-                AsrEngine::AWS,
-                lang,
-            )
-            .await;
-
-        let mut transcribe_rx = lesson.transcript_channel();
-        let mut lang_lesson = lesson.get_or_init(query.to.clone()).await;
-        let mut translate_rx = lang_lesson.translated_channel();
-        let Ok(voice_id) = query.voice.as_deref().unwrap_or("Amy").parse() else {
-            let _ = socket
-                .send(Message::Text(format!("invalid voice id: {:?}", query.voice)))
-                .await;
-          return
-        };
-        let mut voice_lesson = lang_lesson.get_or_init(voice_id).await;
-        let mut voice_rx = voice_lesson.voice_channel();
-        // let mut lip_sync_rx = voice_lesson.lip_sync_channel();
-
-        let fut = async {
-            loop {
-                let evt = select! {
-                    input = socket.next() => {
-                        let Some(res) = input else { break };
-                        let msg = res?;
-                        if msg.is_close() {
-                            break
-                        }
-                        let Message::Binary(bin) = msg else {
-                            tracing::warn!("Other: {:?}", msg);
-                            continue
-                        };
-                        let frame = u8_to_i16(&bin);
-                        lesson.send(frame).await?;
-                        continue
-                    },
-                    transcript_poll = transcribe_rx.recv() => {
-                        let evt = transcript_poll?;
-                        if evt.is_final {
-                            tracing::trace!("Transcribed: {}", evt.transcript);
-                        }
-                        SingleEvent::Transcription { content: evt.transcript, is_final: evt.is_final }
-                    },
-                    translated_poll = translate_rx.recv() => {
-                        let translated = translated_poll?;
-                        SingleEvent::Translation { content: translated }
-                    },
-                    voice_poll = voice_rx.recv() => {
-                        let voice = voice_poll?;
-                        SingleEvent::Voice { content: Base64Box(voice) }
-                    },
-                };
-
-                let Ok(json) = serde_json::to_string(&evt) else {
-                    tracing::warn!("failed to serialize json: {:?}", evt);
-                    continue
-                };
-                socket.send(Message::Text(json)).await?
-            }
-            Ok(())
-        };
-
-        let span = span!(tracing::Level::TRACE, "lesson_speaker", lesson_id = query.id);
-        let _ = span.enter();
-        let res: anyhow::Result<()> = fut.await;
-        match res {
-            Ok(()) => {
-                tracing::info!("lesson speaker closed");
-            }
-            Err(e) => {
-                tracing::warn!("lesson speaker error: {}", e);
-            }
-        }
-    })
-}
-
-fn u8_to_i16(input: &[u8]) -> Vec<i16> {
-    input
-        .chunks_exact(2)
-        .map(|chunk| {
-            let mut buf = [0u8; 2];
-            buf.copy_from_slice(chunk);
-            i16::from_le_bytes(buf)
-        })
-        .collect::<Vec<i16>>()
-}
-
-#[cfg(test)]
-mod test {
-    use std::time::Duration;
-    use async_stream::stream;
-    use poem::listener::{Acceptor, Listener};
-    use tokio::pin;
-    use tokio::time::sleep;
-    use tokio_stream::StreamExt;
-    use tokio_tungstenite::{
-        connect_async,
-        tungstenite::Message,
-    };
-    use tracing::{info, error};
-    use crate::asr::slice_i16_to_u8;
-    use super::*;
-
-    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-    #[tracing_test::traced_test]
-    async fn test_single() {
-        let shared_config = aws_config::load_defaults(BehaviorVersion::latest()).await;
-        let ctx = Context {
-            lessons_manager: LessonsManager::new(&shared_config),
-        };
-
-        let query = SingleQuery {
-            id: 1,
-            from: "zh-CN".to_string(),
-            to: "en-US".to_string(),
-            voice: None,
-        };
-
-        let acceptor = TcpListener::bind("[::]:0")
-            .into_acceptor()
-            .await
-            .unwrap();
-        let addr = acceptor
-            .local_addr()
-            .remove(0)
-            .as_socket_addr()
-            .cloned()
-            .unwrap();
-        let server = Server::new_with_acceptor(acceptor);
-        let handle = tokio::spawn(async move {
-            let _ = server.run(
-                Route::new()
-                    .at("/ws/voice", get(stream_single))
-                    .data(ctx)
-            ).await;
-        });
-
-        let url = format!(
-            "ws://{}/ws/voice?id={}&from={}&to={}&voice={}",
-            addr, query.id, query.from, query.to, query.voice.unwrap_or("Amy".to_string())
-        );
-        let (mut client_stream, _) = connect_async(url)
-            .await
-            .unwrap();
-
-        client_stream
-            .send(Message::Binary(Vec::new()))
-            .await
-            .unwrap();
-
-
-        let wav = hound::WavReader::open("whisper/samples/samples_jfk.wav")
-            .expect("failed to open wav");
-        let spec = wav.spec();
-        info!("{:?}", spec);
-        let samples = wav
-            .into_samples::<i16>()
-            .map(|s| s.unwrap())
-            .collect::<Vec<i16>>();
-        let chunks = samples.chunks(1600)
-            .map(|chunk| chunk.to_vec())
-            .into_iter();
-
-        let audio_stream = stream! {
-            for chunk in chunks {
-                yield slice_i16_to_u8(&chunk);
-                sleep(Duration::from_millis(10)).await;
-            }
-        };
-        pin!(audio_stream);
-
-        let recv_fut = async {
-            while let Some(voice_slice) = audio_stream.next().await {
-                client_stream.send(Message::Binary(voice_slice)).await?;
-            }
-            info!("sent all voice chunks");
-
-            while let Some(next_msg) = client_stream.next().await {
-                debug!(?next_msg);
-                let msg = next_msg?;
-                let Message::Text(json_str) = msg else { continue };
-                let Ok(evt) = serde_json::from_str::<SingleEvent>(&json_str) else { continue };
-                if let SingleEvent::Voice { .. } = evt {
-                    return Ok(())
-                }
-            }
-
-            Ok(()) as anyhow::Result<()>
-        };
-
-        select! {
-            res = recv_fut => {
-                if let Err(e) = res {
-                    error!("Error: {:?}", e);
-                    assert!(false, "Error: {}", e);
-                }
-            }
-            _ = sleep(Duration::from_secs(10)) => {
-                assert!(false, "timeout");
-            }
-        };
-
-        handle.abort();
-    }
-}

tests/stream_single.rs

@@ -0,0 +1,115 @@
+
+use aws_config::BehaviorVersion;
+use std::time::Duration;
+use async_stream::stream;
+use poem::{
+    get,
+    listener::{Listener, Acceptor, TcpListener},
+    EndpointExt, Route, Server,
+};
+use tokio::{pin, select};
+use tokio::time::sleep;
+use tokio_stream::StreamExt;
+use tokio_tungstenite::{
+    connect_async,
+    tungstenite::Message,
+};
+use futures_util::sink::SinkExt;
+use tracing::{info, error, debug};
+use polyhedron::{
+    stream_single,
+    SingleEvent,
+    asr::slice_i16_to_u8,
+    Context
+};
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[tracing_test::traced_test]
+async fn test_single() {
+    let shared_config = aws_config::load_defaults(BehaviorVersion::latest()).await;
+    let ctx = Context::new(&shared_config);
+
+    let acceptor = TcpListener::bind("[::]:0")
+        .into_acceptor()
+        .await
+        .unwrap();
+    let addr = acceptor
+        .local_addr()
+        .remove(0)
+        .as_socket_addr()
+        .cloned()
+        .unwrap();
+    let server = Server::new_with_acceptor(acceptor);
+    let handle = tokio::spawn(async move {
+        let _ = server.run(
+            Route::new()
+                .at("/ws/voice", get(stream_single))
+                .data(ctx)
+        ).await;
+    });
+
+    let url = format!(
+        "ws://{}/ws/voice?id=123abc&from=zh-CN&to=en-US&voice=Amy", addr
+    );
+    let (mut client_stream, _) = connect_async(url)
+        .await
+        .unwrap();
+
+    client_stream
+        .send(Message::Binary(Vec::new()))
+        .await
+        .unwrap();
+
+
+    let wav = hound::WavReader::open("whisper/samples/samples_jfk.wav")
+        .expect("failed to open wav");
+    let spec = wav.spec();
+    info!("{:?}", spec);
+    let samples = wav
+        .into_samples::<i16>()
+        .map(|s| s.unwrap())
+        .collect::<Vec<i16>>();
+    let chunks = samples.chunks(1600)
+        .map(|chunk| chunk.to_vec())
+        .into_iter();
+
+    let audio_stream = stream! {
+        for chunk in chunks {
+            yield slice_i16_to_u8(&chunk);
+            sleep(Duration::from_millis(10)).await;
+        }
+    };
+    pin!(audio_stream);
+
+    let recv_fut = async {
+        while let Some(voice_slice) = audio_stream.next().await {
+            client_stream.send(Message::Binary(voice_slice)).await?;
+        }
+        info!("sent all voice chunks");
+
+        while let Some(Ok(msg)) = client_stream.next().await {
+            debug!("recv: {:?}", msg);
+            let Message::Text(json_str) = msg else { continue };
+            let Ok(evt) = serde_json::from_str::<SingleEvent>(&json_str) else { continue };
+            if let SingleEvent::Voice { .. } = evt {
+                return Ok(())
+            }
+        }
+
+        Ok(()) as anyhow::Result<()>
+    };
+
+    select! {
+        res = recv_fut => {
+            if let Err(e) = res {
+                error!("Error: {:?}", e);
+                assert!(false, "Error: {}", e);
+            }
+        }
+        _ = sleep(Duration::from_secs(10)) => {
+            assert!(false, "timeout");
+        }
+    };
+
+    handle.abort();
+}

whisper/src/handler.rs

@@ -445,7 +445,6 @@ mod test {
 
         let mut rx = whisper_handler.subscribe();
         let send_fut = async {
-            // tokio::time::sleep(Duration::from_secs(5)).await;
             for chunk in chunks {
                 let _ = whisper_handler.send_i16(chunk).await.expect("failed to send sample");
                 tokio::time::sleep(Duration::from_millis(100)).await;