Done!

src/lesson.rs

@@ -1,11 +1,17 @@
-use serde::Deserialize;
 use std::sync::{Arc, Weak};
 use tokio::sync::RwLock;
 use std::collections::BTreeMap;
+use async_stream::stream;
 use aws_config::SdkConfig;
+use aws_sdk_polly::types::VoiceId;
+use aws_sdk_transcribestreaming::primitives::Blob;
+use aws_sdk_transcribestreaming::types::{AudioEvent, AudioStream, LanguageCode, MediaEncoding, TranscriptResultStream};
+use futures_util::{StreamExt, TryStreamExt};
+
 use tokio::select;
+use crate::to_stream;
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct LessonsManager {
     translate_client: aws_sdk_translate::Client,
     polly_client: aws_sdk_polly::Client,
@@ -28,9 +34,9 @@ impl LessonsManager {
 
     pub(crate) async fn create_lesson(&self,
                                       id: u32,
-                                      speaker_lang: String) -> Lesson {
+                                      speaker_lang: LanguageCode) -> Lesson {
         let mut map = self.lessons.write().await;
-        let lesson: Lesson = InnerLesson::new(id, speaker_lang).into();
+        let lesson: Lesson = InnerLesson::new(self.clone(), id, speaker_lang).into();
         map.insert(id, lesson.clone());
         lesson
     }
@@ -48,22 +54,34 @@ pub(crate) struct Lesson {
 
 impl Lesson {
     pub(crate) async fn get_or_init(&self, lang: String) -> LangLesson {
-        let map = self.inner.lang_lessons.read().await;
-        if let Some(lang_lesson) = map.get(&lang).and_then(|weak| weak.upgrade()) {
-            return lang_lesson.into();
+        {
+            let map = self.inner.lang_lessons.read().await;
+            if let Some(lang_lesson) = map.get(&lang).and_then(|weak| weak.upgrade()) {
+                return lang_lesson.into();
+            }
         }
-        let mut map = self.inner.lang_lessons.write().await;
-        if let Some(lang_lesson) = map.get(&lang).and_then(|weak| weak.upgrade()) {
-            lang_lesson.into()
-        } else {
-            let lang_lesson = LangLesson::new(
-                self.clone(),
-                lang.clone(),
-            );
-            map.insert(lang.clone(), Arc::downgrade(&lang_lesson.inner));
-            lang_lesson
+        {
+            let mut map = self.inner.lang_lessons.write().await;
+            if let Some(lang_lesson) = map.get(&lang).and_then(|weak| weak.upgrade()) {
+                lang_lesson.into()
+            } else {
+                let lang_lesson = LangLesson::new(
+                    self.clone(),
+                    lang.clone(),
+                );
+                map.insert(lang.clone(), Arc::downgrade(&lang_lesson.inner));
+                lang_lesson
+            }
         }
     }
+
+    pub(crate) fn voice_channel(&self) -> tokio::sync::mpsc::Sender<Vec<u8>> {
+        self.inner.speaker_voice_channel.clone()
+    }
+
+    pub(crate) fn transcript_channel(&self) -> tokio::sync::broadcast::Receiver<String> {
+        self.inner.speaker_transcript.subscribe()
+    }
 }
 
 impl From<InnerLesson> for Lesson {
@@ -76,8 +94,9 @@ impl From<InnerLesson> for Lesson {
 
 #[derive(Debug)]
 struct InnerLesson {
+    parent: LessonsManager,
     id: u32,
-    speaker_lang: String,
+    speaker_lang: LanguageCode,
     speaker_voice_channel: tokio::sync::mpsc::Sender<Vec<u8>>,
     speaker_transcript: tokio::sync::broadcast::Sender<String>,
     lang_lessons: RwLock<BTreeMap<String, Weak<InnerLangLesson>>>,
@@ -86,30 +105,59 @@ struct InnerLesson {
 
 impl InnerLesson {
     fn new(
+        parent: LessonsManager,
         id: u32,
-        speaker_lang: String
+        speaker_lang: LanguageCode
     ) -> InnerLesson {
         let (speaker_transcript, _) = tokio::sync::broadcast::channel::<String>(128);
+        let shared_speaker_transcript = speaker_transcript.clone();
         let (speaker_voice_channel, mut speaker_voice_rx) = tokio::sync::mpsc::channel(128);
-        let (drop_handler, mut drop_rx) = tokio::sync::oneshot::channel::<Signal>();
+        let (drop_handler, drop_rx) = tokio::sync::oneshot::channel::<Signal>();
+        let transcript_client = parent.transcript_client.clone();
+        let shared_speak_lang = speaker_lang.clone();
 
         tokio::spawn(async move {
             let fut = async {
-                loop {
-                    tokio::time::sleep(std::time::Duration::from_secs(1)).await;
-                }
+                let input_stream = stream! {
+                    while let Some(raw) = speaker_voice_rx.recv().await {
+                        yield Ok(AudioStream::AudioEvent(AudioEvent::builder().audio_chunk(Blob::new(raw)).build()));
+                    }
+                };
+                let output = transcript_client
+                    .start_stream_transcription()
+                    .language_code(shared_speak_lang)//LanguageCode::EnGb
+                    .media_sample_rate_hertz(16000)
+                    .media_encoding(MediaEncoding::Pcm)
+                    .audio_stream(input_stream.into())
+                    .send()
+                    .await
+                    .map_err(|e| crate::StreamTranscriptionError::EstablishStreamError(Box::new(e)))?;
+
+                let mut output_stream = to_stream(output);
+                output_stream
+                    .try_for_each(|text| async {
+                        let _ = shared_speaker_transcript.send(text);
+                        Ok(())
+                    })
+                    .await?;
+                Ok(()) as Result<(), crate::StreamTranscriptionError>
             };
             select! {
-                _ = fut => {}
+                res = fut => {
+                    if let Err(e) = res {
+                        println!("Error: {:?}", e);
+                    }
+                }
                 _ = drop_rx => {}
             }
         });
 
         InnerLesson {
+            parent,
             id,
             speaker_lang,
             speaker_voice_channel,
-            speaker_transcript,
+            speaker_transcript: speaker_transcript,
             lang_lessons: RwLock::new(BTreeMap::new()),
             drop_handler: Some(drop_handler),
         }
@@ -128,8 +176,9 @@ impl Drop for InnerLesson {
 struct InnerLangLesson {
     parent: Lesson,
     lang: String,
-    translated: tokio::sync::broadcast::Sender<String>,
-    voice_lessons: RwLock<BTreeMap<String, Weak<InnerVoiceLesson>>>
+    translated_tx: tokio::sync::broadcast::Sender<String>,
+    voice_lessons: RwLock<BTreeMap<VoiceId, Weak<InnerVoiceLesson>>>,
+    drop_handler: Option<tokio::sync::oneshot::Sender<Signal>>,
 }
 
 #[derive(Clone)]
@@ -137,6 +186,12 @@ pub(crate) struct LangLesson {
     inner: Arc<InnerLangLesson>
 }
 
+impl LangLesson {
+    pub(crate) fn translated_channel(&self) -> tokio::sync::broadcast::Receiver<String> {
+        self.inner.translated_tx.subscribe()
+    }
+}
+
 impl From<InnerLangLesson> for LangLesson {
     fn from(inner: InnerLangLesson) -> Self {
         LangLesson {
@@ -158,39 +213,91 @@ impl LangLesson {
         parent: Lesson,
         lang: String,
     ) -> Self {
-        let (translated, _) = tokio::sync::broadcast::channel::<String>(128);
+        let shared_lang = lang.clone();;
+        let shared_speaker_lang = parent.inner.speaker_lang.clone();
+        let (translated_tx, _) = tokio::sync::broadcast::channel::<String>(128);
+        let shared_translated_tx = translated_tx.clone();
+        let mut transcript_rx = parent.inner.speaker_transcript.subscribe();
+        let translate_client = parent.inner.parent.translate_client.clone();
+        let (drop_handler, drop_rx) = tokio::sync::oneshot::channel::<Signal>();
+        tokio::spawn(async move {
+            let fut = async {
+                while let Ok(text) = transcript_rx.recv().await {
+                    let output = translate_client
+                        .translate_text()
+                        .text(text)
+                        .source_language_code(shared_speaker_lang.as_str())
+                        .target_language_code(shared_lang.clone())
+                        .send()
+                        .await;
+                    match output {
+                        Ok(res) => {
+                            if let Some(translated) = res.translated_text {
+                                let _ = shared_translated_tx.send(translated);
+                            }
+                        },
+                        Err(e) => {
+                            return Err(e);
+                        }
+                        _ => {}
+                    }
+                }
+                Ok(())
+            };
+
+            select! {
+                res = fut => {
+                    if let Err(e) = res {
+                        println!("Error: {:?}", e);
+                    }
+                }
+                _ = drop_rx => {}
+            }
+        });
         InnerLangLesson {
             parent,
             lang,
-            translated,
+            translated_tx,
             voice_lessons: RwLock::new(BTreeMap::new()),
+            drop_handler: Some(drop_handler),
         }.into()
     }
 
-    async fn get_or_init(&mut self, voice: String) -> VoiceLesson {
-        let map = self.inner.voice_lessons.read().await;
-        if let Some(voice_lesson) = map.get(&voice).and_then(|weak| weak.upgrade()) {
-            return voice_lesson.into();
+    pub(crate) async fn get_or_init(&mut self, voice: VoiceId) -> VoiceLesson {
+        {
+            let map = self.inner.voice_lessons.read().await;
+            if let Some(voice_lesson) = map.get(&voice).and_then(|weak| weak.upgrade()) {
+                return voice_lesson.into();
+            }
         }
-        let mut map = self.inner.voice_lessons.write().await;
-        if let Some(voice_lesson) = map.get(&voice).and_then(|weak| weak.upgrade()) {
-            voice_lesson.into()
-        } else {
-            let voice_lesson = Arc::new(InnerVoiceLesson::new(
-                self.clone(),
-                voice.clone(),
-            ));
-            map.insert(voice.clone(), Arc::downgrade(&voice_lesson));
-            voice_lesson.into()
+
+        {
+            let mut map = self.inner.voice_lessons.write().await;
+            if let Some(voice_lesson) = map.get(&voice).and_then(|weak| weak.upgrade()) {
+                voice_lesson.into()
+            } else {
+                let voice_lesson = Arc::new(InnerVoiceLesson::new(
+                    self.clone(),
+                    voice.clone(),
+                ));
+                map.insert(voice, Arc::downgrade(&voice_lesson));
+                voice_lesson.into()
+            }
         }
     }
 }
 
 #[derive(Clone)]
-struct VoiceLesson {
+pub(crate) struct VoiceLesson {
     inner: Arc<InnerVoiceLesson>
 }
 
+impl VoiceLesson {
+    pub(crate) fn voice_channel(&self) -> tokio::sync::broadcast::Receiver<Vec<u8>> {
+        self.inner.voice_lesson.subscribe()
+    }
+}
+
 impl From<InnerVoiceLesson> for VoiceLesson {
     fn from(inner: InnerVoiceLesson) -> Self {
         VoiceLesson {
@@ -209,7 +316,7 @@ impl From<Arc<InnerVoiceLesson>> for VoiceLesson {
 
 struct InnerVoiceLesson {
     parent: LangLesson,
-    voice: String,
+    voice: VoiceId,
     voice_lesson: tokio::sync::broadcast::Sender<Vec<u8>>,
     drop_handler: Option<tokio::sync::oneshot::Sender<Signal>>,
 }
@@ -222,15 +329,38 @@ enum Signal {
 impl InnerVoiceLesson {
     fn new(
         parent: LangLesson,
-        voice: String,
+        voice: VoiceId,
     ) -> InnerVoiceLesson {
+        let shared_voice_id: VoiceId = voice.clone();
         let (tx, rx) = tokio::sync::oneshot::channel::<Signal>();
+        let mut translate_rx = parent.inner.translated_tx.subscribe();
         let (voice_lesson, _) = tokio::sync::broadcast::channel::<Vec<u8>>(128);
+        let shared_voice_lesson = voice_lesson.clone();
+        let client = parent.inner.parent.inner.parent.polly_client.clone();
+        // let lang: LanguageCode = parent.inner.lang.clone().parse().expect("Invalid language code");
         tokio::spawn(async move {
             let fut = async {
-                loop {
-                    tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+                while let Ok(translated) = translate_rx.recv().await {
+                    let res = client.synthesize_speech()
+                        .set_text(Some(translated))
+                        .voice_id(shared_voice_id.clone())
+                        .output_format("pcm".into())
+                        // .language_code(lang)
+                        // .language_code("cmn-CN".into())
+                        .send()
+                        .await;
+                    match res {
+                        Ok(mut synthesized) => {
+                            while let Some(Ok(bytes)) = synthesized.audio_stream.next().await {
+                                let _ = &shared_voice_lesson.send(bytes.to_vec());
+                            }
+                        },
+                        Err(e) => {
+                            return Err(e);
+                        }
+                    }
                 }
+                Ok(())
             };
             select! {
                 _ = fut => {}

src/main.rs

@@ -20,7 +20,7 @@ use clap::Parser;
 
 use poem::{Endpoint, EndpointExt, get, handler, IntoResponse, listener::TcpListener, Route, Server};
 use futures_util::{Sink, SinkExt, TryFutureExt, TryStreamExt};
-use poem::endpoint::StaticFilesEndpoint;
+use poem::endpoint::{StaticFileEndpoint, StaticFilesEndpoint};
 use poem::web::websocket::{Message, WebSocket};
 use futures_util::stream::StreamExt;
 use poem::web::{Data, Query};
@@ -171,8 +171,10 @@ async fn main() -> Result<(), std::io::Error> {
             .index_file("index.html"),
         )
         .at("/translate", get(stream_translate))
-        .at("/lesson-speaker", get(stream_speaker))
-        .at("/lesson-listener", get(stream_listener))
+        .at("/ws/lesson-speaker", get(stream_speaker))
+        .at("/ws/lesson-listener", get(stream_listener))
+        .at("lesson-speaker", StaticFileEndpoint::new("./static/index.html"))
+        .at("lesson-listener", StaticFileEndpoint::new("./static/index.html"))
         .data(ctx);
     let listener = TcpListener::bind("[::]:8080");
     let server = Server::new(listener);
@@ -189,9 +191,40 @@ pub struct LessonSpeakerQuery {
 
 #[handler]
 async fn stream_speaker(ctx: Data<&Context>, query: Query<LessonSpeakerQuery>, ws: WebSocket) -> impl IntoResponse {
-    let lesson = ctx.lessons_manager.create_lesson(query.id, query.lang.clone()).await;
-    println!("{:?}", lesson);
-    println!("{:?}", query);
+    let lesson = ctx.lessons_manager.create_lesson(query.id, query.lang.clone().parse().expect("Not supported lang")).await;
+
+    ws.on_upgrade(|mut socket| async move {
+        let origin_tx = lesson.voice_channel();
+        let mut transcribe_rx = lesson.transcript_channel();
+        loop {
+            select! {
+                msg = socket.next() => {
+                    match msg.as_ref() {
+                        Some(Ok(Message::Binary(bin))) => {
+                            origin_tx.send(bin.to_vec()).await.expect("failed to send");
+                        },
+                        Some(Ok(_)) => {
+                            println!("Other: {:?}", msg);
+                        },
+                        Some(Err(e)) => {
+                            println!("Error: {:?}", e);
+                        },
+                        None => {
+                            socket.close().await.expect("failed to close");
+                            println!("Other: {:?}", msg);
+                            break;
+                        }
+                    }
+                },
+                output = transcribe_rx.recv() => {
+                    if let Ok(transcript) = output {
+                        println!("Transcribed: {}", transcript);
+                        socket.send(Message::Text(transcript)).await.expect("failed to send");
+                    }
+                },
+            }
+        }
+    })
 }
 
 
@@ -204,9 +237,51 @@ pub struct LessonListenerQuery {
 
 #[handler]
 async fn stream_listener(ctx: Data<&Context>, query: Query<LessonListenerQuery>, ws: WebSocket) -> impl IntoResponse {
-    let lesson = ctx.lessons_manager.get_lesson(query.id).await;
-    println!("{:?}", lesson);
+    let lesson_opt = ctx.lessons_manager.get_lesson(query.id).await;
     println!("{:?}", query);
+    let voice_id = query.voice.parse().expect("Not supported voice");
+
+    ws.on_upgrade(|mut socket| async move {
+        let Some(lesson) = lesson_opt else {
+            let _ = socket.send(Message::Text("lesson not found".to_string())).await;
+            return
+        };
+
+        println!("lesson found");
+        let mut transcript_rx = lesson.transcript_channel();
+        println!("transcribe start");
+
+        let mut lang_lesson = lesson.get_or_init(query.lang.clone()).await;
+        let mut translate_rx = lang_lesson.translated_channel();
+        println!("translate start");
+
+        let mut voice_lesson = lang_lesson.get_or_init(voice_id).await;
+        let mut voice_rx = voice_lesson.voice_channel();
+        println!("synthesize start");
+
+        loop {
+            select! {
+                transcript = transcript_rx.recv() => {
+                    if let Ok(transcript) = transcript {
+                        println!("Transcribed: {}", transcript);
+                        let _ = socket.send(Message::Text(transcript)).await;
+                    }
+                },
+                translated = translate_rx.recv() => {
+                    if let Ok(translated) = translated {
+                        println!("Translated: {}", translated);
+                        let _ = socket.send(Message::Text(translated)).await;
+                    }
+                },
+                voice = voice_rx.recv() => {
+                    if let Ok(voice) = voice {
+                        println!("Synthesized: {:?}", voice.len());
+                        let _ = socket.send(Message::Binary(voice)).await;
+                    }
+                },
+            }
+        }
+    })
 }
 
 #[handler]

static/client.js

@@ -13,7 +13,8 @@ if (location.protocol === "https:") {
     websocket_uri = "ws:"
 }
 websocket_uri += "//" + location.host
-websocket_uri += location.pathname + "translate"
+websocket_uri += "/ws" + location.pathname
+websocket_uri += location.search
 const socket = new WebSocket(websocket_uri);
 
 //================= CONFIG =================
@@ -122,10 +123,15 @@ function stopRecording() {
     // videoElement.srcObject = null;
 }
 
+
+const audioQueue = new rxjs.Subject();
+audioQueue
+    .pipe(rxjs.concatMap(playAudio))
+    .subscribe(_ => console.log('played audio'));
 //================= SOCKET IO =================
 socket.onmessage = function (msg) {
     if (msg.data instanceof Blob) {
-        playAudio(msg.data)
+        audioQueue.next(msg.data)
     } else {
         // text
         onSpeechData(msg.data)
@@ -258,8 +264,11 @@ function capitalize(s) {
     return s.charAt(0).toUpperCase() + s.slice(1);
 }
 
+const audioContext = new (window.AudioContext || window.webkitAudioContext)();
+
+let nextStartTime = audioContext.currentTime;
+
 async function playAudio(chunk) {
-    const audioContext = new (window.AudioContext || window.webkitAudioContext)();
     const totalLength = chunk.size;
 
     // Create an AudioBuffer of enough size
@@ -267,14 +276,20 @@ async function playAudio(chunk) {
     const output = audioBuffer.getChannelData(0);
 
     // Copy the PCM samples into the AudioBuffer
-    const int16Array = new Int16Array(await chunk.arrayBuffer())
+    const arrayBuf = await chunk.arrayBuffer();
+    const int16Array = new Int16Array(arrayBuf, 0, Math.floor(arrayBuf.byteLength / 2))
     for(let i = 0; i < int16Array.length; i++) {
         output[i] = int16Array[i] / 32768.0;  // Convert to [-1, 1] float32 range
     }
 
     // 3. Play the audio using Web Audio API
+
     const source = audioContext.createBufferSource();
     source.buffer = audioBuffer;
     source.connect(audioContext.destination);
-    source.start();
+    source.start(nextStartTime);
+    nextStartTime = Math.max(nextStartTime, audioContext.currentTime) + audioBuffer.duration;
+    source.onended = () => {
+        console.log('audio slice ended');
+    }
 }

static/index.html

@@ -18,6 +18,7 @@
   <audio></audio>
 
   <br>
+  <button id="startButton" type="button"> Start listening</button>
   <button id="startRecButton" type="button"> Start recording</button>
   <button id="stopRecButton" type="button"> Stop recording</button>
   <div id="recordingStatus">&nbsp;</div>
@@ -40,6 +41,7 @@
 <!-- Socket -->
 <!--<script src="assets/js/socket.io.js"></script>-->
 
+<script src="https://unpkg.com/rxjs@%5E7/dist/bundles/rxjs.umd.min.js"></script>
 <!-- Client -->
 <script src="client.js"></script>
 </body>