2nd commit

src/main.rs

@@ -5,24 +5,30 @@
 
 #![allow(clippy::result_large_err)]
 
+use std::collections::BTreeMap;
 use std::default::Default;
+use std::error::Error;
+use std::fmt::{Debug, Display, Formatter};
+use std::future::Future;
 use tokio::sync::mpsc::channel;
 use async_stream::stream;
 use aws_config::meta::region::RegionProviderChain;
 use aws_sdk_transcribestreaming::primitives::Blob;
 use aws_sdk_transcribestreaming::types::{AudioStream, AudioEvent, LanguageCode, MediaEncoding, TranscriptResultStream};
-use aws_sdk_transcribestreaming::{config::Region, meta::PKG_VERSION, Client, Error};
-use bytes::BufMut;
+use aws_sdk_transcribestreaming::{config::Region, meta::PKG_VERSION};
+use aws_sdk_transcribestreaming::operation::start_stream_transcription::StartStreamTranscriptionOutput;
 use clap::Parser;
 
 use poem::{handler, listener::TcpListener, Server, get, Route, IntoResponse, Endpoint, EndpointExt};
-use futures_util::{Sink, SinkExt};
+use futures_util::{Sink, SinkExt, TryFutureExt, TryStreamExt};
 use poem::endpoint::StaticFilesEndpoint;
 use poem::web::websocket::{Message, WebSocket};
 use futures_util::stream::StreamExt;
+use poem::web::Data;
 
-use tokio::select;
+use tokio::{select};
 use tokio::sync::mpsc::{Receiver, Sender};
+use tokio_stream::Stream;
 
 
 #[derive(Debug, Parser)]
@@ -46,7 +52,6 @@ enum ReplyEvent {
     Synthesized(Vec<u8>),
 }
 
-const CHUNK_SIZE: usize = 8192;
 
 /// Transcribes an audio file to text.
 /// # Arguments
@@ -58,35 +63,11 @@ const CHUNK_SIZE: usize = 8192;
 ///   If not supplied, uses the value of the **AWS_REGION** environment variable.
 ///   If the environment variable is not set, defaults to **us-west-2**.
 /// * `[-v]` - Whether to display additional information.
-async fn stream_process(mut rx: Receiver<Vec<u8>>, tx: Sender<ReplyEvent>) -> Result<(), Error> {
-    tracing_subscriber::fmt::init();
-
-    let Opt {
-        region,
-        // audio_file,
-        verbose,
-    } = Opt::parse();
-
-    let region_provider = RegionProviderChain::first_try(region.map(Region::new))
-        .or_default_provider()
-        .or_else(Region::new("us-west-2"));
-
-    println!();
-
-    if verbose {
-        println!("Transcribe client version: {}", PKG_VERSION);
-        println!(
-            "Region:                    {}",
-            region_provider.region().await.unwrap().as_ref()
-        );
-        // println!("Audio filename:            {}", &audio_file);
-        println!();
-    }
-
-    let shared_config = aws_config::from_env().region(region_provider).load().await;
-    let transcript_client = Client::new(&shared_config);
-    let translate_client = aws_sdk_translate::Client::new(&shared_config);
-    let polly_client = aws_sdk_polly::Client::new(&shared_config);
+async fn stream_process(translate_client: aws_sdk_translate::Client,
+                        polly_client: aws_sdk_polly::Client,
+                        transcript_client: aws_sdk_transcribestreaming::Client,
+                        mut rx: Receiver<Vec<u8>>,
+                        tx: Sender<ReplyEvent>) -> Result<(), StreamTranscriptionError> {
 
     let input_stream = stream! {
         while let Some(raw) = rx.recv().await {
@@ -94,47 +75,24 @@ async fn stream_process(mut rx: Receiver<Vec<u8>>, tx: Sender<ReplyEvent>) -> Re
         }
     };
 
-    let mut output = transcript_client
+    let output = transcript_client
         .start_stream_transcription()
         .language_code(LanguageCode::ZhCn)//LanguageCode::EnGb
         .media_sample_rate_hertz(16000)
         .media_encoding(MediaEncoding::Pcm)
         .audio_stream(input_stream.into())
         .send()
-        .await?;
-
-    while let Some(event) = output.transcript_result_stream.recv().await? {
-        match event {
-            TranscriptResultStream::TranscriptEvent(transcript_event) => {
-                let transcript = transcript_event.transcript.expect("transcript");
-                for result in transcript.results.unwrap_or_default() {
-                    if result.is_partial {
-                        if verbose {
-                            println!("Partial: {:?}", result);
-                        }
-                    } else {
-                        let first_alternative = &result.alternatives.as_ref().expect("should have")[0];
-                        let slice = first_alternative.transcript.as_ref().expect("should have");
-                        println!("Line:   {:?}", slice);
-                        tx.send(ReplyEvent::Transcribed(slice.clone())).await.expect("failed to send");
-                        let lc = result.language_code.as_ref().map(|lc| lc.as_str().to_string());
-                        let translated = translate(&translate_client, first_alternative.transcript.clone(), lc).await;
-                        if let Some(has) = translated {
-                            tx.send(ReplyEvent::Transcribed(has.clone())).await.expect("failed to send");
-                            println!("Translated: {}", has);
-                            if let Some(synthesized) = synthesize(&polly_client, has).await {
-                                tx.send(ReplyEvent::Synthesized(synthesized)).await.expect("failed to send");
-                            }
-                        }
+        .await
+        .map_err(|e| StreamTranscriptionError::EstablishStreamError(Box::new(e)))?;
 
-                    }
-                }
-            }
-            otherwise => panic!("received unexpected event type: {:?}", otherwise),
-        }
-    }
+    let output_stream = to_stream(output);
 
-    Ok(())
+    output_stream
+        .flat_map(|res| {
+            process(translate_client.clone(), polly_client.clone(), res)
+        })
+        .try_for_each(|reply| tx.send(reply).map_err(|e| StreamTranscriptionError::Shutdown))
+        .await
 }
 
 async fn translate(client: &aws_sdk_translate::Client, transcript: Option<String>, source_lang_code: Option<String>) -> Option<String> {
@@ -158,24 +116,63 @@ async fn synthesize(client: &aws_sdk_polly::Client, transcript: String) -> Optio
     Some(bs.to_vec())
 }
 
-fn pcm_data(audio_file: &str) -> Vec<u8> {
-    let reader = hound::WavReader::open(audio_file).unwrap();
-    let samples_result: hound::Result<Vec<i16>> = reader.into_samples::<i16>().collect();
+#[derive(Clone)]
+struct Context {
+    translate_client: aws_sdk_translate::Client,
+    polly_client: aws_sdk_polly::Client,
+    transcript_client: aws_sdk_transcribestreaming::Client,
+}
+
+struct Lessons {
+    lessons: BTreeMap<u32, Lesson>
+}
+
+struct Lesson {
 
-    let mut pcm: Vec<u8> = Vec::new();
-    for sample in samples_result.unwrap() {
-        pcm.put_i16_le(sample);
-    }
-    pcm
 }
 
 #[tokio::main]
-async fn main() -> Result<(), std::io::Error> {    let app = Route::new().nest(
+async fn main() -> Result<(), std::io::Error> {
+    tracing_subscriber::fmt::init();
+
+    let Opt {
+        region,
+        verbose,
+    } = Opt::parse();
+
+    let region_provider = RegionProviderChain::first_try(region.map(Region::new))
+        .or_default_provider()
+        .or_else(Region::new("us-west-2"));
+
+    println!();
+
+    if verbose {
+        println!("Transcribe client version: {}", PKG_VERSION);
+        println!(
+            "Region:                    {}",
+            region_provider.region().await.unwrap().as_ref()
+        );
+        // println!("Audio filename:            {}", &audio_file);
+        println!();
+    }
+
+    let shared_config = aws_config::from_env().region(region_provider).load().await;
+    let transcript_client = aws_sdk_transcribestreaming::Client::new(&shared_config);
+    let translate_client = aws_sdk_translate::Client::new(&shared_config);
+    let polly_client = aws_sdk_polly::Client::new(&shared_config);
+    let ctx = Context {
+        translate_client,
+        polly_client,
+        transcript_client,
+    };
+
+    let app = Route::new().nest(
         "/",
         StaticFilesEndpoint::new("./static")
             .show_files_listing()
             .index_file("index.html"),
-    ).at("/translate", get(stream_translate));
+    ).at("/translate", get(stream_translate))
+        .data(ctx);
     let listener = TcpListener::bind("[::]:8080");
     let server = Server::new(listener);
 
@@ -184,11 +181,19 @@ async fn main() -> Result<(), std::io::Error> {    let app = Route::new().nest(
 
 
 #[handler]
-async fn stream_translate(ws: WebSocket) -> impl IntoResponse {
+async fn stream_translate(ctx: Data<&Context>, ws: WebSocket) -> impl IntoResponse {
+    let translate_client = ctx.translate_client.clone();
+    let polly_client = ctx.polly_client.clone();
+    let transcript_client = ctx.transcript_client.clone();
     ws.on_upgrade(|mut socket| async move {
         let (origin_tx, origin_rx) = channel::<Vec<u8>>(128);
         let (translate_tx, mut translate_rx) = channel::<ReplyEvent>(128);
-        let stream_fut = stream_process(origin_rx, translate_tx);
+        let stream_fut = stream_process(
+            translate_client,
+            polly_client,
+            transcript_client,
+            origin_rx,
+            translate_tx);
 
         let ws_fut = async {
             loop {
@@ -239,3 +244,81 @@ async fn stream_translate(ws: WebSocket) -> impl IntoResponse {
     })
 }
 
+#[derive(Debug)]
+enum StreamTranscriptionError {
+    EstablishStreamError(Box<dyn Error + Send + Sync>),
+    TranscriptResultStreamError(Box<dyn Error + Send + Sync>),
+    Shutdown,
+    Unknown
+}
+
+
+impl Display for StreamTranscriptionError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            StreamTranscriptionError::EstablishStreamError(e) => write!(f, "EstablishStreamError: {}", e),
+            StreamTranscriptionError::TranscriptResultStreamError(e) => write!(f, "TranscriptResultStreamError: {}", e),
+            StreamTranscriptionError::Shutdown => write!(f, "Shutdown"),
+            StreamTranscriptionError::Unknown => write!(f, "Unknown"),
+        }
+    }
+}
+
+impl Error for StreamTranscriptionError {
+    fn source(&self) -> Option<&(dyn Error + 'static)> {
+        match self {
+            StreamTranscriptionError::EstablishStreamError(e) => Some(e.as_ref()),
+            StreamTranscriptionError::TranscriptResultStreamError(e) => Some(e.as_ref()),
+            StreamTranscriptionError::Shutdown => None,
+            StreamTranscriptionError::Unknown => None,
+        }
+    }
+}
+
+fn to_stream(mut output: StartStreamTranscriptionOutput) -> impl Stream<Item=Result<String, StreamTranscriptionError>> {
+    stream! {
+        while let Some(event) = output
+            .transcript_result_stream
+            .recv()
+            .await
+            .map_err(|e| StreamTranscriptionError::TranscriptResultStreamError(Box::new(e)))? {
+            match event {
+                TranscriptResultStream::TranscriptEvent(transcript_event) => {
+                    let transcript = transcript_event.transcript.expect("transcript");
+                    for result in transcript.results.unwrap_or_default() {
+                        if !result.is_partial {
+                            let first_alternative = &result.alternatives.as_ref().expect("should have")[0];
+                            let slice = first_alternative.transcript.as_ref().expect("should have");
+                            yield Ok(slice.clone());
+                        }
+                    }
+                }
+                otherwise => yield Err(StreamTranscriptionError::Unknown),
+            }
+        }
+    }
+}
+
+fn process(translate_client: aws_sdk_translate::Client,
+               polly_client: aws_sdk_polly::Client,
+               res: Result<String, StreamTranscriptionError>) -> impl Stream<Item=Result<ReplyEvent, StreamTranscriptionError>> {
+    stream! {
+        match res {
+            Ok(transcription) => {
+                yield Ok(ReplyEvent::Transcribed(transcription.clone()));
+                let translated = translate(&translate_client, Some(transcription), Some("en".to_string())).await;
+                if let Some(has) = translated {
+                    yield Ok(ReplyEvent::Translated(has.clone()));
+                    println!("Translated: {}", has);
+                    if let Some(synthesized) = synthesize(&polly_client, has).await {
+                        yield Ok(ReplyEvent::Synthesized(synthesized));
+                    }
+                }
+            },
+            Err(e) => {
+                yield Err(e);
+            }
+        }
+
+    }
+}