rimelabs_sdk/stt/
stream.rs

1use super::{queue::TranscriptQueue, transport, TranscriptionOptions, TranscriptionUpdate};
2use crate::{client::Inner, Error, ErrorKind};
3use bytes::Bytes;
4use futures_core::Stream;
5use futures_util::FutureExt;
6use std::{
7    fmt,
8    pin::Pin,
9    sync::Arc,
10    task::{Context, Poll},
11    time::Duration,
12};
13use tokio::{sync::oneshot, task::JoinHandle, time::Instant};
14use tokio_util::sync::CancellationToken;
15
16/// A single-reader stream of replacement transcripts.
17///
18/// The first poll starts network work and the overall timeout. Dropping the stream
19/// cancels its source and RPC. Use [`Self::close`] to also wait for worker cleanup.
20pub struct TranscriptStream {
21    shared: Arc<TranscriptQueue>,
22    cancellation: CancellationToken,
23    start: Option<oneshot::Sender<Instant>>,
24    worker: Option<JoinHandle<()>>,
25    timeout: Option<Duration>,
26    deadline: Option<Instant>,
27    finished: bool,
28}
29impl fmt::Debug for TranscriptStream {
30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31        f.debug_struct("TranscriptStream")
32            .field("request_id", &self.request_id())
33            .field("finished", &self.finished)
34            .finish_non_exhaustive()
35    }
36}
37impl TranscriptStream {
38    /// The request identifier, when supplied by the service.
39    pub fn request_id(&self) -> Option<&str> {
40        self.shared.request_id()
41    }
42    /// Cancel only this utterance, including an unread final result.
43    pub fn cancel(&self) {
44        self.cancellation.cancel();
45    }
46    /// Cancel and wait for the worker to release the source and RPC.
47    pub async fn close(&mut self) {
48        self.cancel();
49        if let Some(worker) = self.worker.take() {
50            let _ = worker.await;
51        }
52        self.finished = true;
53    }
54}
55impl Drop for TranscriptStream {
56    fn drop(&mut self) {
57        self.cancel();
58    }
59}
60impl Stream for TranscriptStream {
61    type Item = Result<TranscriptionUpdate, Error>;
62    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
63        let this = self.get_mut();
64        if this.finished {
65            return Poll::Ready(None);
66        }
67        if this.cancellation.is_cancelled() {
68            this.shared
69                .fail(Error::new(ErrorKind::Cancelled, "transcription cancelled"));
70        } else if this
71            .deadline
72            .is_some_and(|deadline| Instant::now() >= deadline)
73        {
74            this.shared
75                .fail(Error::new(ErrorKind::Timeout, "transcription timed out"));
76            this.cancellation.cancel();
77        }
78        if let Some(start) = this.start.take() {
79            let now = Instant::now();
80            this.deadline = this.timeout.map(|timeout| now + timeout);
81            let _ = start.send(now);
82        }
83        match this.shared.poll_next(cx) {
84            Poll::Ready(result) => {
85                this.finished = matches!(&result, Err(_) | Ok(TranscriptionUpdate::Final { .. }));
86                Poll::Ready(Some(result))
87            }
88            Poll::Pending => Poll::Pending,
89        }
90    }
91}
92
93pub(super) fn start<S>(
94    client: Arc<Inner>,
95    source: S,
96    options: TranscriptionOptions,
97) -> Result<TranscriptStream, Error>
98where
99    S: Stream<Item = Result<Bytes, Error>> + Send + 'static,
100{
101    // Validate runtime time support before registering work, as the TTS API does.
102    let _timer = tokio::time::sleep(Duration::ZERO);
103    let shared = Arc::new(TranscriptQueue::default());
104    let cancellation = client.cancellation.child_token();
105    let (start, started) = oneshot::channel();
106    let operation = cancellation.clone();
107    let state = shared.clone();
108    let timeout = options.timeout;
109    let owner = client.clone();
110    let worker = owner.spawn(async move {
111        let work = async {
112            let began = tokio::select! {
113                biased;
114                _ = operation.cancelled() => return Err(Error::new(ErrorKind::Cancelled, "transcription cancelled")),
115                began = started => began.map_err(|_| Error::new(ErrorKind::Cancelled, "transcription dropped"))?,
116            };
117            let deadline = async {
118                match timeout {
119                    Some(timeout) => tokio::time::sleep_until(began + timeout).await,
120                    None => std::future::pending().await,
121                }
122            };
123            let execute = async {
124                let final_result = transport::run(client, source, options, state.clone()).await?;
125                state.put(final_result).await?;
126                // Keep cancellation and the overall deadline alive until final consumption.
127                state.wait_consumed().await;
128                Ok(())
129            };
130            tokio::select! {
131                biased;
132                _ = operation.cancelled() => Err(Error::new(ErrorKind::Cancelled, "transcription cancelled")),
133                _ = deadline => Err(Error::new(ErrorKind::Timeout, "transcription timed out")),
134                result = execute => result,
135            }
136        };
137        let result = std::panic::AssertUnwindSafe(work).catch_unwind().await
138            .unwrap_or_else(|_| Err(Error::new(ErrorKind::Stream, "audio source panicked")));
139        if let Err(error) = result { state.fail(error); }
140    })?;
141    Ok(TranscriptStream {
142        shared,
143        cancellation,
144        start: Some(start),
145        worker: Some(worker),
146        timeout,
147        deadline: None,
148        finished: false,
149    })
150}
rimelabs_sdk/stt/stream.rs

rimelabs_sdk/stt/
stream.rs
