clean warning, clean code

src/asr/aws.rs

@@ -2,7 +2,7 @@ use std::error::Error;
 use std::fmt::{Debug, Display, Formatter};
 use async_stream::stream;
 use async_trait::async_trait;
-use aws_config::SdkConfig;
+use aws_config::BehaviorVersion;
 use aws_sdk_transcribestreaming::operation::start_stream_transcription::StartStreamTranscriptionOutput;
 use aws_sdk_transcribestreaming::primitives::Blob;
 use aws_sdk_transcribestreaming::types::{
@@ -14,21 +14,21 @@ use tokio_stream::Stream;
 use futures_util::TryStreamExt;
 use crate::asr::{ASR, Event};
 
-pub struct AWS_ASR {
+pub struct AwsAsr {
     speaker_voice_channel: tokio::sync::mpsc::Sender<Vec<i16>>,
     speaker_transcript: tokio::sync::broadcast::Sender<Event>,
     drop_handler: Option<tokio::sync::oneshot::Sender<()>>,
 }
 
-impl Debug for AWS_ASR {
+impl Debug for AwsAsr {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         write!(f, "AWS_ASR")
     }
 }
 
-impl AWS_ASR {
+impl AwsAsr {
     pub async fn from_env(lang: LanguageCode) -> anyhow::Result<Self> {
-        let config = aws_config::load_from_env().await;
+        let config = aws_config::load_defaults(BehaviorVersion::latest()).await;
         let transcript_client = aws_sdk_transcribestreaming::Client::new(&config);
 
         let (speaker_voice_channel, mut speaker_voice_rx) = tokio::sync::mpsc::channel::<Vec<i16>>(128);
@@ -87,12 +87,12 @@ fn slice_i16_to_u8(slice: &[i16]) -> Vec<u8> {
     slice
         .iter()
         .flat_map(|&sample| {
-            [(sample >> 8) as u8, sample as u8]
+            [sample as u8, (sample >> 8) as u8]
         })
         .collect()
 }
 
-impl Drop for AWS_ASR {
+impl Drop for AwsAsr {
     fn drop(&mut self) {
         if let Some(drop_handler) = self.drop_handler.take() {
             let _ = drop_handler.send(());
@@ -102,7 +102,7 @@ impl Drop for AWS_ASR {
 
 
 #[async_trait]
-impl ASR for AWS_ASR {
+impl ASR for AwsAsr {
     async fn frame(&mut self, frame: Vec<i16>) -> anyhow::Result<()> {
         Ok(self.speaker_voice_channel.send(frame).await?)
     }

src/asr/whisper.rs

@@ -16,19 +16,19 @@ lazy_static! {
         .expect("Failed to initialize whisper context");
 }
 
-pub struct Whisper_ASR {
+pub struct WhisperAsr {
     whisper: WhisperHandler,
     tx: tokio::sync::broadcast::Sender<Event>,
 }
 
-impl Debug for Whisper_ASR {
+impl Debug for WhisperAsr {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         write!(f, "Whisper_ASR")
     }
 }
 
-impl Whisper_ASR {
-    pub async fn from_config() -> Result<Whisper_ASR, Error> {
+impl WhisperAsr {
+    pub async fn from_config() -> Result<WhisperAsr, Error> {
         let whisper = CONTEXT.create_handler(&SETTINGS.whisper, "".to_string())?;
         let mut output_rx = whisper.subscribe();
         let (tx, _) = tokio::sync::broadcast::channel(64);
@@ -71,7 +71,7 @@ impl Whisper_ASR {
 }
 
 #[async_trait]
-impl ASR for Whisper_ASR {
+impl ASR for WhisperAsr {
     async fn frame(&mut self, frame: Vec<i16>) -> anyhow::Result<()> {
         Ok(self.whisper.send_i16(frame).await?)
     }

src/lesson.rs

@@ -5,18 +5,18 @@ use aws_sdk_transcribestreaming::types::{LanguageCode};
 use futures_util::future::try_join;
 use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
-use std::fmt::{Debug, Display, Formatter};
+use std::fmt::{Debug, Formatter};
 use std::io::BufRead;
 use std::ops::Deref;
 use std::sync::{Arc, Weak};
 use tokio::sync::RwLock;
-use tracing::{error, warn};
+use tracing::{error};
 
 use tokio::select;
-use crate::asr::{Event, aws::AWS_ASR, ASR};
+use crate::asr::{Event, aws::AwsAsr, ASR};
 
 #[cfg(feature = "whisper")]
-use crate::asr::whisper::Whisper_ASR;
+use crate::asr::whisper::WhisperAsr;
 
 pub struct InnerLessonsManager {
     translate_client: aws_sdk_translate::Client,
@@ -43,23 +43,24 @@ impl Deref for LessonsManager {
     }
 }
 
-pub(crate) enum ASR_Engine {
+pub(crate) enum AsrEngine {
     AWS,
+    #[allow(dead_code)]
     #[cfg(feature = "whisper")]
     Whisper,
 }
 
-impl ASR_Engine {
+impl AsrEngine {
     async fn create(self, lang: LanguageCode) -> Box<dyn ASR + Send> {
         match self {
-            ASR_Engine::AWS => Box::new(
-                AWS_ASR::from_env(lang)
+            AsrEngine::AWS => Box::new(
+                AwsAsr::from_env(lang)
                     .await
                     .expect("Failed to initialize AWS ASR")
             ),
             #[cfg(feature = "whisper")]
-            ASR_Engine::Whisper => Box::new(
-                Whisper_ASR::from_config()
+            AsrEngine::Whisper => Box::new(
+                WhisperAsr::from_config()
                     .await
                     .expect("Failed to initialize Whisper ASR")
             ),
@@ -79,7 +80,7 @@ impl LessonsManager {
         LessonsManager { inner: Arc::new(inner) }
     }
 
-    pub(crate) async fn create_lesson(&self, id: u32, engine: ASR_Engine, speaker_lang: LanguageCode) -> Lesson {
+    pub(crate) async fn create_lesson(&self, id: u32, 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());
@@ -153,7 +154,7 @@ pub(crate) struct InnerLesson {
 }
 
 impl InnerLesson {
-    async fn new(parent: LessonsManager, engine: ASR_Engine, speaker_lang: LanguageCode) -> InnerLesson {
+    async fn new(parent: LessonsManager, engine: AsrEngine, speaker_lang: LanguageCode) -> InnerLesson {
         let (speaker_transcript, _) = tokio::sync::broadcast::channel::<Event>(128);
         let shared_speaker_transcript = speaker_transcript.clone();
         let (speaker_voice_channel, mut speaker_voice_rx) = tokio::sync::mpsc::channel::<Vec<i16>>(128);

src/main.rs

@@ -7,7 +7,10 @@
 
 #[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},
@@ -44,7 +47,7 @@ async fn main() -> Result<(), std::io::Error> {
 
     debug!("Transcribe client version: {}", PKG_VERSION);
 
-    let shared_config = aws_config::load_from_env().await;
+    let shared_config = aws_config::load_defaults(BehaviorVersion::latest()).await;
     let ctx = Context {
         lessons_manager: LessonsManager::new(&shared_config),
     };
@@ -57,7 +60,9 @@ async fn main() -> Result<(), std::io::Error> {
                 .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(
             "lesson-speaker",
             StaticFileEndpoint::new("./static/index.html"),
@@ -83,7 +88,8 @@ async fn main() -> Result<(), std::io::Error> {
 #[derive(Deserialize, Debug)]
 pub struct LessonSpeakerQuery {
     id: u32,
-    lang: String,
+    language: String,
+    #[allow(dead_code)] // TODO: use this in the future
     prompt: Option<String>,
 }
 
@@ -93,51 +99,58 @@ async fn stream_speaker(
     query: Query<LessonSpeakerQuery>,
     ws: WebSocket,
 ) -> impl IntoResponse {
-    let lesson = ctx
-        .lessons_manager
-        .create_lesson(
-            query.id,
-            ASR_Engine::AWS,
-            query.lang.clone().parse().expect("Not supported lang"),
-        )
-        .await;
-    let prompt = query.prompt.clone().unwrap_or_default();
-
+    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();
-        loop {
-            select! {
-                msg = socket.next() => {
-                    match msg.as_ref() {
-                        Some(Ok(Message::Binary(bin))) => {
-                            let frame = u8_to_i16(bin);
-                            if let Err(e) = lesson.send(frame).await {
-                                tracing::warn!("failed to send voice: {}", e);
-                                break;
-                            }
-                        },
-                        Some(Ok(_)) => {
+        let fut = async {
+            loop {
+                select! {
+                    msg = socket.next() => {
+                        let Some(res) = msg else { break };
+                        let msg = res?;
+                        let Message::Binary(bin) = msg else {
                             tracing::warn!("Other: {:?}", msg);
-                        },
-                        Some(Err(e)) => {
-                            tracing::warn!("Error: {:?}", e);
-                        },
-                        None => {
-                            if let Err(e) = socket.close().await {
-                                tracing::debug!("Message: {:?}, {}", msg, e);
-                            }
-                            break;
-                        }
-                    }
-                },
-                output = transcribe_rx.recv() => {
-                    if let Ok(evt) = output {
+                            continue
+                        };
+                        let frame = u8_to_i16(&bin);
+                        lesson.send(frame).await?
+                    },
+                    output = transcribe_rx.recv() => {
+                        let evt = output?;
                         tracing::trace!("Transcribed: {}", evt.transcript);
-                        let evt = LiveLessonTextEvent::Transcription { text: evt.transcript };
-                        let json = serde_json::to_string(&evt).expect("failed to serialize");
-                        let _ = socket.send(Message::Text(json)).await.expect("failed to send");
-                    }
-                },
+                        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 res: anyhow::Result<()> = fut.await;
+        match res {
+            Ok(()) => {
+                tracing::info!("lesson speaker closed");
+            }
+            Err(e) => {
+                tracing::warn!("lesson speaker error: {}", e);
             }
         }
     })
@@ -146,26 +159,29 @@ async fn stream_speaker(
 #[derive(Deserialize, Debug)]
 pub struct LessonListenerQuery {
     id: u32,
-    lang: String,
+    language: String,
     voice: String,
 }
 
-#[derive(Serialize)]
+#[derive(Serialize, Debug)]
 #[serde(tag = "type")]
 enum LiveLessonTextEvent {
-    Transcription { text: String },
-    Translation { text: String },
+    #[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 lesson_opt = ctx.lessons_manager.get_lesson(query.id).await;
-    debug!("listener param = {:?}", query);
+    let lessons_manager = ctx.lessons_manager.clone();
 
     ws.on_upgrade(|mut socket| async move {
         let voice_id = match query.voice.parse() {
@@ -177,6 +193,9 @@ async fn stream_listener(
                 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()))
@@ -184,54 +203,62 @@ async fn stream_listener(
             return;
         };
         let mut transcript_rx = lesson.transcript_channel();
-        let mut lang_lesson = lesson.get_or_init(query.lang.clone()).await;
+        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();
 
-        loop {
-            select! {
-                transcript = transcript_rx.recv() => {
-                    if let Ok(evt) = transcript {
-                        let evt = LiveLessonTextEvent::Transcription { text: evt.transcript };
-                        match serde_json::to_string(&evt) {
-                            Ok(json) => {
-                                tracing::debug!("Transcribed: {}", json);
-                                let _ = socket.send(Message::Text(json)).await;
-                            },
-                            Err(e) => {
-                                tracing::error!("failed to serialize: {}", e);
-                            }
-                        }
-                    }
-                },
-                translated = translate_rx.recv() => {
-                    if let Ok(translated) = translated {
-                        let evt = LiveLessonTextEvent::Translation { text: translated };
-                        match serde_json::to_string(&evt) {
-                            Ok(json) => {
-                                tracing::debug!("Translated: {}", json);
-                                let _ = socket.send(Message::Text(json)).await;
-                            },
-                            Err(e) => {
-                                tracing::error!("failed to serialize: {}", e);
-                            }
-                        }
-                    }
-                },
-                voice = voice_rx.recv() => {
-                    if let Ok(voice) = voice {
-                        let _ = socket.send(Message::Binary(voice)).await;
-                    }
-                },
-                visemes = lip_sync_rx.recv() => {
-                    if let Ok(visemes) = visemes {
+        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 json = serde_json::to_string(&evt).expect("failed to serialize");
-                        let _ = socket.send(Message::Text(json)).await;
-                    }
-                },
+                        let Ok(json) = serde_json::to_string(&evt) else {
+                            tracing::warn!("failed to serialize: {:?}", evt);
+                            continue
+                        };
+                        socket.send(Message::Text(json)).await?
+                    },
+                }
+            }
+        };
+
+        let res: anyhow::Result<()> = fut.await;
+        match res {
+            Ok(()) => {
+                tracing::info!("lesson listener closed");
+            }
+            Err(e) => {
+                tracing::warn!("lesson listener error: {}", e);
             }
         }
     })

static/index.html

@@ -5,7 +5,7 @@
     <link rel="icon" type="image/svg+xml" href="/vite.svg" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <title>Vite + React</title>
-    <script type="module" crossorigin src="/assets/index-34871d29.js"></script>
+    <script type="module" crossorigin src="/assets/index-b40ca4ea.js"></script>
     <link rel="stylesheet" href="/assets/index-983f9492.css">
   </head>
   <body>

whisper/src/config.rs

@@ -9,7 +9,6 @@ pub struct WhisperConfig {
     pub(crate) keep_ms: usize,
     pub model: String,
     pub(crate) max_prompt_tokens: usize,
-    pub(crate) context_confidence_threshold: f32,
 }
 
 #[allow(dead_code)]

whisper/src/handler.rs

@@ -6,7 +6,7 @@ use std::{
 };
 use fvad::SampleRate;
 
-use tokio::sync::{broadcast, mpsc, oneshot, OnceCell};
+use tokio::sync::{broadcast, mpsc, oneshot};
 use tokio::time::Instant;
 use tracing::{warn};
 use whisper_rs::{convert_integer_to_float_audio, WhisperContext, WhisperError, WhisperState, WhisperToken, WhisperTokenData};
@@ -30,18 +30,6 @@ impl <'a> Context {
     }
 }
 
-static WHISPER_CONTEXT: OnceCell<WhisperContext> = OnceCell::const_new();
-
-async fn initialize_whisper_context(model: String) -> WhisperContext {
-    tokio::task::spawn_blocking(move || {
-        WhisperContext::new(&model).expect("failed to create WhisperContext")
-    }).await.expect("failed to spawn")
-}
-
-async fn get_whisper_context(model: String) -> &'static WhisperContext {
-    WHISPER_CONTEXT.get_or_init(|| initialize_whisper_context(model)).await
-}
-
 #[derive(Debug)]
 pub enum Error {
     WhisperError {