1use super::{TranscriptionClient, TranscriptionOption, handle_wait_for_answer_with_audio_drop};
2use crate::{
3 event::{EventSender, SessionEvent},
4 media::{Sample, SourcePacket, TrackId},
5};
6use anyhow::{Result, anyhow};
7use async_trait::async_trait;
8use audio_codec::samples_to_bytes;
9use futures::{SinkExt, StreamExt};
10use serde::{Deserialize, Serialize};
11use std::{
12 future::Future,
13 pin::Pin,
14 sync::{
15 Arc,
16 atomic::{AtomicBool, Ordering},
17 },
18 time::Duration,
19};
20use tokio::{net::TcpStream, sync::mpsc};
21use tokio_tungstenite::tungstenite::client::IntoClientRequest;
22use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async, tungstenite::Message};
23use tokio_util::sync::CancellationToken;
24use tracing::{debug, info, warn};
25use url::Url;
26use uuid::Uuid;
27
28const DEEPGRAM_LISTEN_URL: &str = "wss://api.deepgram.com/v1/listen";
29
30type TranscriptionClientFuture =
31 Pin<Box<dyn Future<Output = Result<Box<dyn TranscriptionClient>>> + Send>>;
32
33struct DeepgramAsrClientInner {
34 option: TranscriptionOption,
35}
36
37pub struct DeepgramAsrClient {
38 audio_tx: mpsc::UnboundedSender<Vec<u8>>,
39 is_closed: Arc<AtomicBool>,
40}
41
42pub struct DeepgramAsrClientBuilder {
43 option: TranscriptionOption,
44 track_id: Option<String>,
45 token: Option<CancellationToken>,
46 event_sender: EventSender,
47}
48
49#[derive(Debug, Deserialize)]
50#[serde(default)]
51struct DeepgramResult {
52 #[serde(rename = "type")]
53 event_type: Option<String>,
54 channel: Option<DeepgramChannel>,
55 is_final: bool,
56 speech_final: bool,
57 start: Option<f32>,
58 duration: Option<f32>,
59 metadata: Option<DeepgramMetadata>,
60}
61
62impl Default for DeepgramResult {
63 fn default() -> Self {
64 Self {
65 event_type: None,
66 channel: None,
67 is_final: false,
68 speech_final: false,
69 start: None,
70 duration: None,
71 metadata: None,
72 }
73 }
74}
75
76#[derive(Debug, Deserialize)]
77struct DeepgramChannel {
78 alternatives: Vec<DeepgramAlternative>,
79}
80
81#[derive(Debug, Deserialize)]
82struct DeepgramAlternative {
83 transcript: String,
84 confidence: Option<f32>,
85}
86
87#[derive(Debug, Deserialize)]
88struct DeepgramMetadata {
89 request_id: Option<String>,
90}
91
92#[derive(Serialize)]
93struct CloseStreamCommand {
94 #[serde(rename = "type")]
95 event_type: &'static str,
96}
97
98impl DeepgramAsrClientBuilder {
99 pub fn create(
100 track_id: TrackId,
101 token: CancellationToken,
102 option: TranscriptionOption,
103 event_sender: EventSender,
104 ) -> TranscriptionClientFuture {
105 Box::pin(async move {
106 let builder = Self::new(option, event_sender);
107 builder
108 .with_cancel_token(token)
109 .with_track_id(track_id)
110 .build()
111 .await
112 .map(|client| Box::new(client) as Box<dyn TranscriptionClient>)
113 })
114 }
115
116 pub fn new(option: TranscriptionOption, event_sender: EventSender) -> Self {
117 Self {
118 option,
119 token: None,
120 track_id: None,
121 event_sender,
122 }
123 }
124
125 pub fn with_cancel_token(mut self, cancellation_token: CancellationToken) -> Self {
126 self.token = Some(cancellation_token);
127 self
128 }
129
130 pub fn with_track_id(mut self, track_id: String) -> Self {
131 self.track_id = Some(track_id);
132 self
133 }
134
135 pub async fn build(self) -> Result<DeepgramAsrClient> {
136 let (audio_tx, mut audio_rx) = mpsc::unbounded_channel();
137
138 let event_sender_rx = match self.option.start_when_answer {
139 Some(true) => Some(self.event_sender.subscribe()),
140 _ => None,
141 };
142
143 let track_id = self.track_id.unwrap_or_else(|| Uuid::new_v4().to_string());
144 let token = self.token.unwrap_or_default();
145 let event_sender = self.event_sender;
146 let inner = DeepgramAsrClientInner {
147 option: self.option,
148 };
149 let is_closed = Arc::new(AtomicBool::new(false));
150 let client = DeepgramAsrClient {
151 audio_tx,
152 is_closed: Arc::clone(&is_closed),
153 };
154
155 crate::spawn(async move {
156 let res = async move {
157 if event_sender_rx.is_some() {
158 handle_wait_for_answer_with_audio_drop(event_sender_rx, &mut audio_rx, &token)
159 .await;
160
161 if token.is_cancelled() {
162 debug!("Cancelled during wait for answer");
163 return Ok::<(), anyhow::Error>(());
164 }
165 }
166
167 let ws_stream = match inner.connect_websocket(&track_id).await {
168 Ok(stream) => stream,
169 Err(e) => {
170 warn!(
171 track_id,
172 "Failed to connect to Deepgram ASR WebSocket: {}", e
173 );
174 let _ = event_sender.send(SessionEvent::Error {
175 timestamp: crate::media::get_timestamp(),
176 track_id,
177 sender: "DeepgramAsrClient".to_string(),
178 error: format!("Failed to connect to Deepgram ASR WebSocket: {}", e),
179 code: Some(500),
180 });
181 return Err(e);
182 }
183 };
184
185 info!(%track_id, "Starting Deepgram ASR client");
186 match DeepgramAsrClient::handle_websocket_message(
187 track_id.clone(),
188 ws_stream,
189 audio_rx,
190 event_sender.clone(),
191 token,
192 inner.option.refer,
193 )
194 .await
195 {
196 Ok(_) => debug!(track_id, "Deepgram ASR websocket handling completed"),
197 Err(e) => {
198 info!(track_id, "Error in Deepgram ASR websocket handling: {}", e);
199 event_sender
200 .send(SessionEvent::Error {
201 track_id,
202 timestamp: crate::media::get_timestamp(),
203 sender: "deepgram_asr".to_string(),
204 error: e.to_string(),
205 code: None,
206 })
207 .ok();
208 }
209 }
210 Ok::<(), anyhow::Error>(())
211 }
212 .await;
213 is_closed.store(true, Ordering::SeqCst);
214 if let Err(e) = res {
215 debug!("Deepgram ASR task finished with error: {:?}", e);
216 }
217 });
218
219 Ok(client)
220 }
221}
222
223impl DeepgramAsrClientInner {
224 async fn connect_websocket(
225 &self,
226 track_id: &str,
227 ) -> Result<WebSocketStream<MaybeTlsStream<TcpStream>>> {
228 let api_key = self
229 .option
230 .secret_key
231 .clone()
232 .or_else(|| std::env::var("DEEPGRAM_API_KEY").ok())
233 .ok_or_else(|| anyhow!("No DEEPGRAM_API_KEY provided"))?;
234
235 let mut url = Url::parse(
236 self.option
237 .endpoint
238 .as_deref()
239 .unwrap_or(DEEPGRAM_LISTEN_URL),
240 )?;
241 let extra = self.option.extra.as_ref();
242 {
243 let mut query = url.query_pairs_mut();
244 if extra.map(|e| !e.contains_key("model")).unwrap_or(true) {
245 query.append_pair("model", self.option.model_type.as_deref().unwrap_or("nova-3"));
246 }
247 if extra.map(|e| !e.contains_key("language")).unwrap_or(true) {
248 if let Some(language) = self.option.language.as_deref() {
249 if language != "auto" {
250 query.append_pair("language", language);
251 }
252 }
253 }
254 if extra.map(|e| !e.contains_key("encoding")).unwrap_or(true) {
255 query.append_pair("encoding", "linear16");
256 }
257 if extra.map(|e| !e.contains_key("sample_rate")).unwrap_or(true) {
258 query.append_pair(
259 "sample_rate",
260 self.option.samplerate.unwrap_or(16000).to_string().as_str(),
261 );
262 }
263 if extra.map(|e| !e.contains_key("channels")).unwrap_or(true) {
264 query.append_pair("channels", "1");
265 }
266 if extra
267 .map(|e| !e.contains_key("interim_results"))
268 .unwrap_or(true)
269 {
270 query.append_pair("interim_results", "true");
271 }
272 if extra.map(|e| !e.contains_key("endpointing")).unwrap_or(true) {
273 query.append_pair("endpointing", "300");
274 }
275 if extra.map(|e| !e.contains_key("smart_format")).unwrap_or(true) {
276 query.append_pair("smart_format", "true");
277 }
278 if let Some(extra) = extra {
279 for (key, value) in extra {
280 query.append_pair(key, value);
281 }
282 }
283 }
284
285 let mut request = url.as_str().into_client_request()?;
286 request
287 .headers_mut()
288 .insert("Authorization", format!("Token {}", api_key).parse()?);
289 let (ws_stream, response) = connect_async(request).await?;
290 debug!(
291 track_id,
292 "Deepgram WebSocket connection established. Response: {}",
293 response.status()
294 );
295 Ok(ws_stream)
296 }
297}
298
299impl DeepgramAsrClient {
300 async fn handle_websocket_message(
301 track_id: TrackId,
302 ws_stream: WebSocketStream<MaybeTlsStream<TcpStream>>,
303 mut audio_rx: mpsc::UnboundedReceiver<Vec<u8>>,
304 event_sender: EventSender,
305 token: CancellationToken,
306 refer: Option<bool>,
307 ) -> Result<()> {
308 let (mut ws_sender, mut ws_receiver) = ws_stream.split();
309 let begin_time = crate::media::get_timestamp();
310 let mut final_text = String::new();
311 let mut final_start_time = None;
312 let mut index = 0u32;
313 let mut first_result_time = None;
314
315 let token_clone = token.clone();
316 crate::spawn(async move {
317 let mut keep_alive = tokio::time::interval(Duration::from_secs(5));
318 keep_alive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
319 loop {
320 tokio::select! {
321 _ = token_clone.cancelled() => {
322 break;
323 }
324 _ = keep_alive.tick() => {
325 if let Err(e) = ws_sender.send(Message::Text(r#"{"type":"KeepAlive"}"#.into())).await {
326 warn!("Failed to send Deepgram KeepAlive message: {}", e);
327 break;
328 }
329 }
330 audio_data = audio_rx.recv() => {
331 match audio_data {
332 Some(audio_data) => {
333 if audio_data.is_empty() {
334 continue;
335 }
336 if let Err(e) = ws_sender.send(Message::Binary(audio_data.into())).await {
337 warn!("Failed to send audio data to Deepgram: {}", e);
338 break;
339 }
340 }
341 None => break,
342 }
343 }
344 }
345 }
346
347 let close_msg = CloseStreamCommand {
348 event_type: "CloseStream",
349 };
350 if let Ok(msg_json) = serde_json::to_string(&close_msg) {
351 if let Err(e) = ws_sender.send(Message::Text(msg_json.into())).await {
352 warn!("Failed to send Deepgram CloseStream message: {}", e);
353 }
354 }
355 });
356
357 loop {
358 tokio::select! {
359 msg = ws_receiver.next() => {
360 match msg {
361 Some(Ok(Message::Text(text))) => {
362 match serde_json::from_str::<DeepgramResult>(&text) {
363 Ok(result) => {
364 if result.event_type.as_deref().is_some_and(|event_type| event_type != "Results") {
365 continue;
366 }
367
368 let Some(alternative) = result
369 .channel
370 .as_ref()
371 .and_then(|channel| channel.alternatives.first()) else {
372 continue;
373 };
374 let transcript = alternative.transcript.trim();
375 if transcript.is_empty() {
376 continue;
377 }
378
379 let start_time = result.start.map(|start| {
380 begin_time + (start * 1000.0).max(0.0) as u64
381 });
382 let end_time = result.start.zip(result.duration).map(|(start, duration)| {
383 begin_time + ((start + duration) * 1000.0).max(0.0) as u64
384 });
385
386 if result.is_final {
387 if final_start_time.is_none() {
388 final_start_time = start_time;
389 }
390 if !final_text.is_empty() {
391 final_text.push(' ');
392 }
393 final_text.push_str(transcript);
394 }
395
396 let text = if result.speech_final {
397 if !result.is_final {
398 if final_start_time.is_none() {
399 final_start_time = start_time;
400 }
401 if !final_text.is_empty() {
402 final_text.push(' ');
403 }
404 final_text.push_str(transcript);
405 }
406 let text = final_text.trim().to_string();
407 final_text.clear();
408 text
409 } else if result.is_final {
410 final_text.clone()
411 } else if final_text.is_empty() {
412 transcript.to_string()
413 } else {
414 format!("{} {}", final_text, transcript)
415 };
416
417 if text.is_empty() {
418 continue;
419 }
420
421 let timestamp = crate::media::get_timestamp();
422 let task_id = result
423 .metadata
424 .as_ref()
425 .and_then(|metadata| metadata.request_id.clone());
426 let event = if result.speech_final {
427 SessionEvent::AsrFinal {
428 track_id: track_id.clone(),
429 index,
430 text,
431 timestamp,
432 start_time: final_start_time.or(start_time),
433 end_time,
434 is_filler: None,
435 confidence: alternative.confidence,
436 task_id,
437 refer,
438 }
439 } else {
440 SessionEvent::AsrDelta {
441 track_id: track_id.clone(),
442 index,
443 text,
444 timestamp,
445 start_time: final_start_time.or(start_time),
446 end_time,
447 is_filler: None,
448 confidence: alternative.confidence,
449 task_id,
450 refer,
451 }
452 };
453 event_sender.send(event).ok();
454
455 let first_time = first_result_time.get_or_insert(timestamp);
456 let metrics_key = if result.speech_final {
457 "completed.asr.deepgram"
458 } else {
459 "ttfb.asr.deepgram"
460 };
461 event_sender
462 .send(SessionEvent::Metrics {
463 timestamp,
464 key: metrics_key.to_string(),
465 data: serde_json::json!({ "index": index }),
466 duration: (timestamp - *first_time) as u32,
467 })
468 .ok();
469
470 if result.speech_final {
471 index += 1;
472 final_start_time = None;
473 first_result_time = None;
474 }
475 }
476 Err(e) => {
477 warn!(track_id, "Failed to parse Deepgram ASR response: {} {}", e, text);
478 }
479 }
480 }
481 Some(Ok(Message::Close(frame))) => {
482 info!(track_id, "Deepgram WebSocket connection closed: {:?}", frame);
483 break;
484 }
485 Some(Err(e)) => {
486 return Err(anyhow!("Deepgram WebSocket error: {}", e));
487 }
488 Some(_) => {}
489 None => break,
490 }
491 }
492 _ = token.cancelled() => {
493 break;
494 }
495 }
496 }
497
498 Ok(())
499 }
500}
501
502#[async_trait]
503impl TranscriptionClient for DeepgramAsrClient {
504 fn send_audio(&self, samples: &[Sample], _src_packet: Option<&SourcePacket>) -> Result<()> {
505 if self.is_closed.load(Ordering::SeqCst) {
506 return Ok(());
507 }
508 self.audio_tx
509 .send(samples_to_bytes(samples))
510 .map_err(|_| {
511 self.is_closed.store(true, Ordering::SeqCst);
512 })
513 .ok();
514 Ok(())
515 }
516}