use super::{TranscriptionClient, TranscriptionOption, handle_wait_for_answer_with_audio_drop};
use crate::{
event::{EventSender, SessionEvent},
media::{Sample, SourcePacket, TrackId},
};
use anyhow::{Result, anyhow};
use async_trait::async_trait;
use audio_codec::samples_to_bytes;
use futures::{SinkExt, StreamExt};
use http::StatusCode;
use serde::{Deserialize, Serialize};
use serde_with::skip_serializing_none;
use std::{future::Future, pin::Pin, sync::Arc};
use tokio::{net::TcpStream, sync::mpsc};
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async, tungstenite::Message};
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};
use uuid::Uuid;
struct AliyunAsrClientInner {
audio_tx: mpsc::UnboundedSender<Vec<u8>>,
option: TranscriptionOption,
}
pub struct AliyunAsrClient {
inner: Arc<AliyunAsrClientInner>,
}
#[derive(Debug, Serialize)]
struct RunTaskCommand {
header: CommandHeader,
payload: RunTaskPayload,
}
impl RunTaskCommand {
pub fn new(task_id: String, option: &TranscriptionOption) -> Self {
let model_type = option
.model_type
.clone()
.unwrap_or("paraformer-realtime-v2".to_string());
let format = "pcm";
let sample_rate = option.samplerate.unwrap_or(16000);
let vocabulary_id = option
.extra
.as_ref()
.and_then(|extra| extra.get("vocabulary_id"))
.cloned();
let disfluency_removal_enabled = option
.extra
.as_ref()
.and_then(|extra| extra.get("disfluency_removal_enabled"))
.and_then(|v| v.parse::<bool>().ok());
let mut language_hints = Vec::new();
if let Some(language) = option.language.clone() {
language_hints.push(language);
}
let semantic_punctuation_enabled = option
.extra
.as_ref()
.and_then(|extra| extra.get("semantic_punctuation_enabled"))
.and_then(|v| v.parse::<bool>().ok());
let max_sentence_silence = option
.extra
.as_ref()
.and_then(|extra| extra.get("max_sentence_silence"))
.and_then(|v| v.parse::<u32>().ok());
let multi_threshold_mode_enabled = option
.extra
.as_ref()
.and_then(|extra| extra.get("multi_threshold_mode_enabled"))
.and_then(|v| v.parse::<bool>().ok());
let punctuation_prediction_enabled = option
.extra
.as_ref()
.and_then(|extra| extra.get("punctuation_prediction_enabled"))
.and_then(|v| v.parse::<bool>().ok());
let heartbeat = option
.extra
.as_ref()
.and_then(|extra| extra.get("heartbeat"))
.and_then(|v| v.parse::<bool>().ok());
let inverse_text_normalization_enabled = option
.extra
.as_ref()
.and_then(|extra| extra.get("inverse_text_normalization_enabled"))
.and_then(|v| v.parse::<bool>().ok());
Self {
header: CommandHeader {
action: "run-task".to_string(),
task_id,
streaming: "duplex".to_string(),
},
payload: RunTaskPayload {
task_group: "audio".to_string(),
task: "asr".to_string(),
function: "recognition".to_string(),
model: model_type,
input: CommandInput {},
parameters: RunTaskCommandParameters {
format: format.to_string(),
sample_rate,
vocabulary_id,
disfluency_removal_enabled,
language_hints,
semantic_punctuation_enabled,
max_sentence_silence,
multi_threshold_mode_enabled,
punctuation_prediction_enabled,
heartbeat,
inverse_text_normalization_enabled,
},
},
}
}
}
#[derive(Debug, Serialize)]
pub struct FinishTaskCommand {
pub header: CommandHeader,
pub payload: FinishTaskPayload,
}
impl FinishTaskCommand {
fn new(task_id: String) -> Self {
Self {
header: CommandHeader {
action: "finish-task".to_string(),
task_id,
streaming: "duplex".to_string(),
},
payload: FinishTaskPayload {
input: CommandInput {},
},
}
}
}
#[derive(Debug, Serialize)]
pub struct CommandHeader {
pub action: String,
pub task_id: String,
pub streaming: String,
}
#[derive(Debug, Serialize)]
struct RunTaskPayload {
task_group: String,
task: String,
function: String,
model: String,
input: CommandInput,
parameters: RunTaskCommandParameters,
}
#[skip_serializing_none]
#[derive(Debug, Serialize)]
struct RunTaskCommandParameters {
format: String,
sample_rate: u32,
vocabulary_id: Option<String>,
disfluency_removal_enabled: Option<bool>,
#[serde(skip_serializing_if = "Vec::is_empty")]
language_hints: Vec<String>,
semantic_punctuation_enabled: Option<bool>,
max_sentence_silence: Option<u32>, multi_threshold_mode_enabled: Option<bool>, punctuation_prediction_enabled: Option<bool>,
heartbeat: Option<bool>,
inverse_text_normalization_enabled: Option<bool>,
}
#[derive(Debug, Serialize)]
pub struct CommandInput {}
#[derive(Debug, Serialize)]
pub struct FinishTaskPayload {
pub input: CommandInput,
}
#[derive(Debug, Deserialize)]
pub struct AsrEvent {
pub header: EventHeader,
pub payload: Option<EventPayload>,
}
#[derive(Debug, Deserialize)]
pub struct EventHeader {
pub task_id: String,
pub event: String,
pub error_code: Option<String>,
pub error_message: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct EventPayload {
pub output: Option<EventOutput>,
}
#[derive(Debug, Deserialize)]
pub struct EventOutput {
pub sentence: OutputSentence,
}
#[allow(dead_code)]
#[derive(Debug, Deserialize)]
pub struct OutputSentence {
pub sentence_id: u32,
pub begin_time: u32,
pub end_time: Option<u32>,
pub text: String,
pub words: Vec<OutputWord>,
pub heartbeat: Option<bool>,
pub sentence_end: bool,
}
#[allow(dead_code)]
#[derive(Debug, Deserialize)]
pub struct OutputWord {
pub begin_time: u32,
pub end_time: u32,
pub text: String,
pub punctuation: String,
}
impl AliyunAsrClientInner {
async fn connect_websocket(
&self,
voice_id: &String,
) -> Result<WebSocketStream<MaybeTlsStream<TcpStream>>> {
let api_key = self
.option
.secret_key
.clone()
.or_else(|| std::env::var("DASHSCOPE_API_KEY").ok())
.ok_or_else(|| anyhow!("No DASHSCOPE_API_KEY provided"))?;
let ws_url = self
.option
.endpoint
.as_deref()
.unwrap_or("wss://dashscope.aliyuncs.com/api-ws/v1/inference");
let mut request = ws_url.into_client_request()?;
let headers = request.headers_mut();
headers.insert("Authorization", format!("Bearer {}", api_key).parse()?);
headers.insert("X-DashScope-DataInspection", "enable".parse()?);
let (ws_stream, response) = connect_async(request).await?;
debug!(
voice_id,
"WebSocket connection established. Response: {}",
response.status()
);
match response.status() {
StatusCode::SWITCHING_PROTOCOLS => Ok(ws_stream),
_ => Err(anyhow!("Failed to connect to WebSocket: {:?}", response)),
}
}
}
struct WebSocketContext {
track_id: TrackId,
option: TranscriptionOption,
}
impl AliyunAsrClient {
async fn handle_websocket_message(
ctx: WebSocketContext,
ws_stream: WebSocketStream<MaybeTlsStream<TcpStream>>,
mut audio_rx: mpsc::UnboundedReceiver<Vec<u8>>,
event_sender: EventSender,
token: CancellationToken,
) -> Result<()> {
let (mut ws_sender, mut ws_receiver) = ws_stream.split();
let begin_time = crate::media::get_timestamp();
let aliyun_task_id = Uuid::new_v4().to_string();
let start_msg = RunTaskCommand::new(aliyun_task_id.clone(), &ctx.option);
if let Ok(msg_json) = serde_json::to_string(&start_msg) {
if let Err(e) = ws_sender.send(Message::Text(msg_json.into())).await {
warn!(ctx.track_id, "Failed to send start message: {}", e);
return Err(anyhow!("Failed to send start message: {}", e));
}
}
let track_id_for_recv = ctx.track_id.clone();
let track_id_for_send = ctx.track_id.clone();
let aliyun_task_id_for_finish = aliyun_task_id.clone();
let recv_loop = async move {
let track_id = track_id_for_recv;
while let Some(msg) = ws_receiver.next().await {
match msg {
Ok(Message::Text(text)) => {
debug!(track_id = %track_id, "Received ASR text message: {}", text);
match serde_json::from_str::<AsrEvent>(&text) {
Ok(response) => {
let task_id = response.header.task_id.clone();
debug!(track_id, task_id, "Task: {}", response.header.event);
if response.header.event == "task-failed" {
let code = response.header.error_code.expect("mising_code");
let message = response
.header
.error_message
.expect("mising error message");
let error =
format!("Error from ASR service: {} ({})", code, message);
warn!("{}", error);
return Err(anyhow!(error));
}
if response.header.event == "task-finished" {
break;
}
if response.header.event == "task-started" {
continue;
}
let payload = response.payload.ok_or(anyhow!("missing payload"))?;
let output = payload.output.ok_or(anyhow!("missing output"))?;
if output.sentence.heartbeat.unwrap_or(false) {
continue;
}
let sentence = output.sentence;
let words = sentence.words;
let sentence_start_time = begin_time + sentence.begin_time as u64;
let sentence_end_time = begin_time
+ sentence.end_time.unwrap_or_else(|| {
words
.last()
.map(|w| w.end_time)
.unwrap_or(sentence.begin_time)
}) as u64;
let text = sentence.text;
let event = if sentence.sentence_end {
SessionEvent::AsrFinal {
track_id: track_id.clone(),
index: sentence.sentence_id,
text,
timestamp: crate::media::get_timestamp(),
start_time: Some(sentence_start_time),
end_time: Some(sentence_end_time),
is_filler: None,
confidence: None,
task_id: None,
}
} else {
SessionEvent::AsrDelta {
track_id: track_id.clone(),
index: sentence.sentence_id,
text,
timestamp: crate::media::get_timestamp(),
start_time: Some(sentence_start_time),
end_time: Some(sentence_end_time),
is_filler: None,
confidence: None,
task_id: None,
}
};
event_sender.send(event).ok();
let diff_time = (crate::media::get_timestamp() - begin_time) as u32;
let metrics_event = if sentence.sentence_end {
SessionEvent::Metrics {
timestamp: crate::media::get_timestamp(),
key: "completed.asr.aliyun".to_string(),
data: serde_json::json!({
"sentence_id": sentence.sentence_id,
}),
duration: diff_time,
}
} else {
SessionEvent::Metrics {
timestamp: crate::media::get_timestamp(),
key: "ttfb.asr.aliyun".to_string(),
data: serde_json::json!({
"sentence_id": sentence.sentence_id,
}),
duration: diff_time,
}
};
event_sender.send(metrics_event).ok();
}
Err(e) => {
warn!(track_id, "Failed to parse ASR response: {}", e);
}
}
}
Ok(Message::Close(_)) => {
info!(track_id = %track_id, "WebSocket connection closed by server");
break;
}
Err(e) => {
warn!(track_id = %track_id, "WebSocket error: {}", e);
return Err(anyhow!("WebSocket error: {}", e));
}
_ => {
debug!("Received non-text message");
}
}
}
Ok(())
};
let token_clone = token.clone();
let send_loop = async move {
let mut total_bytes = 0;
while let Some(audio_data) = audio_rx.recv().await {
if token_clone.is_cancelled() {
break;
}
let data_len = audio_data.len();
if let Err(e) = ws_sender.send(Message::Binary(audio_data.into())).await {
warn!("Failed to send audio data: {}", e);
break;
}
total_bytes += data_len;
if total_bytes % 32000 == 0 {
debug!(track_id = %track_id_for_send, "Sent {} bytes of audio to Aliyun", total_bytes);
}
}
let end_msg = FinishTaskCommand::new(aliyun_task_id_for_finish);
if let Ok(msg_json) = serde_json::to_string(&end_msg) {
if let Err(e) = ws_sender.send(Message::Text(msg_json.into())).await {
warn!("Failed to send end message: {}", e);
}
}
Ok::<(), anyhow::Error>(())
};
crate::spawn(async move {
if let Err(e) = send_loop.await {
debug!("Aliyun ASR send_loop finished with error: {:?}", e);
}
});
tokio::select! {
result = recv_loop => result,
_ = token.cancelled() => {
Ok::<(), anyhow::Error>(())
}
}
}
}
pub struct AliyunAsrClientBuilder {
option: TranscriptionOption,
track_id: Option<String>,
token: Option<CancellationToken>,
event_sender: EventSender,
}
type TranscriptionClientFuture =
Pin<Box<dyn Future<Output = Result<Box<dyn TranscriptionClient>>> + Send>>;
impl AliyunAsrClientBuilder {
pub fn create(
track_id: TrackId,
token: CancellationToken,
option: TranscriptionOption,
event_sender: EventSender,
) -> TranscriptionClientFuture {
Box::pin(async move {
let builder = Self::new(option, event_sender);
builder
.with_cancel_token(token)
.with_track_id(track_id)
.build()
.await
.map(|client| Box::new(client) as Box<dyn TranscriptionClient>)
})
}
pub fn new(option: TranscriptionOption, event_sender: EventSender) -> Self {
Self {
option,
token: None,
track_id: None,
event_sender,
}
}
pub fn with_cancel_token(mut self, cancellation_token: CancellationToken) -> Self {
self.token = Some(cancellation_token);
self
}
pub fn with_secret_key(mut self, secret_key: String) -> Self {
self.option.secret_key = Some(secret_key);
self
}
pub fn with_model_type(mut self, model_type: String) -> Self {
self.option.model_type = Some(model_type);
self
}
pub fn with_track_id(mut self, track_id: String) -> Self {
self.track_id = Some(track_id);
self
}
pub async fn build(self) -> Result<AliyunAsrClient> {
let (audio_tx, mut audio_rx) = mpsc::unbounded_channel();
let event_sender_rx = match self.option.start_when_answer {
Some(true) => Some(self.event_sender.subscribe()),
_ => None,
};
let inner = Arc::new(AliyunAsrClientInner {
audio_tx,
option: self.option.clone(),
});
let client = AliyunAsrClient {
inner: inner.clone(),
};
let track_id = self.track_id.unwrap_or_else(|| Uuid::new_v4().to_string());
let token = self.token.unwrap_or_default();
let event_sender = self.event_sender;
crate::spawn(async move {
let res = async move {
if event_sender_rx.is_some() {
handle_wait_for_answer_with_audio_drop(event_sender_rx, &mut audio_rx, &token)
.await;
if token.is_cancelled() {
debug!("Cancelled during wait for answer");
return Ok::<(), anyhow::Error>(());
}
}
let ws_stream = match inner.connect_websocket(&track_id).await {
Ok(stream) => stream,
Err(e) => {
warn!(track_id, "Failed to connect to Aliyun ASR WebSocket: {}", e);
let _ = event_sender.send(SessionEvent::Error {
timestamp: crate::media::get_timestamp(),
track_id,
sender: "AliyunAsrClient".to_string(),
error: format!("Failed to connect to Aliyun ASR WebSocket: {}", e),
code: Some(500),
});
return Err(e);
}
};
info!(%track_id, option=?inner.option, "Starting Aliyun ASR client");
let ctx = WebSocketContext {
track_id: track_id.clone(),
option: self.option,
};
match AliyunAsrClient::handle_websocket_message(
ctx,
ws_stream,
audio_rx,
event_sender.clone(),
token,
)
.await
{
Ok(_) => {
debug!("WebSocket message handling completed");
}
Err(e) => {
info!("Error in handle_websocket_message: {}", e);
event_sender
.send(SessionEvent::Error {
track_id,
timestamp: crate::media::get_timestamp(),
sender: "aliyun_asr".to_string(),
error: e.to_string(),
code: None,
})
.ok();
}
}
Ok::<(), anyhow::Error>(())
}
.await;
if let Err(e) = res {
debug!("Aliyun ASR task finished with error: {:?}", e);
}
});
Ok(client)
}
}
#[async_trait]
impl TranscriptionClient for AliyunAsrClient {
fn send_audio(&self, samples: &[Sample], _src_packet: Option<&SourcePacket>) -> Result<()> {
let audio_data = samples_to_bytes(samples);
self.inner
.audio_tx
.send(audio_data)
.map_err(|_| anyhow!("Failed to send audio data"))?;
Ok(())
}
}