use anyhow::{Context, Result};
const VAD_INDEX: usize = 2;
pub struct SttClient(gradium::Client);
pub struct SttStreamReceiver(gradium::stt::SttStreamReceiver);
pub struct SttStreamSender(gradium::stt::SttStreamSender);
#[derive(Debug)]
pub enum Msg {
Step {
end_of_turn: bool,
current_s: f64,
inactivity_prob: f64,
},
Text {
text: String,
start_s: f64,
},
}
impl SttClient {
pub fn new(gradium_api_key: Option<&str>, base_url: &str) -> Result<Self> {
let api_key = match gradium_api_key {
Some(key) => key.to_string(),
None => std::env::var("GRADIUM_API_KEY")
.map_err(|_| anyhow::anyhow!("GRADIUM_API_KEY environment variable not set"))?,
};
let client = gradium::Client::new(&api_key)
.with_base_url(base_url)
.context("STT: failed to initialize Gradium client")?;
Ok(Self(client))
}
pub async fn stt_stream(
&self,
lang: crate::Lang,
extra_config: Option<&str>,
) -> Result<(SttStreamSender, SttStreamReceiver)> {
let lang_code = match lang {
crate::Lang::En => "en",
crate::Lang::Fr => "fr",
crate::Lang::Es => "es",
crate::Lang::De => "de",
crate::Lang::Pt => "pt",
};
tracing::info!(language = lang_code, "starting STT stream");
let json_config = {
let mut config = serde_json::Map::new();
config.insert("language".into(), serde_json::json!(lang_code));
if let Some(extra) = extra_config
&& let Ok(serde_json::Value::Object(map)) = serde_json::from_str(extra)
{
config.extend(map);
}
serde_json::Value::Object(config).to_string()
};
let setup = gradium::protocol::stt::Setup {
model_name: "default".to_string(),
input_format: gradium::protocol::AudioFormat::Pcm,
json_config: Some(json_config),
};
let stream = self
.0
.stt_stream(setup)
.await
.context("STT: failed to connect")?;
let (tx, rx) = stream.split();
Ok((SttStreamSender(tx), SttStreamReceiver(rx)))
}
}
impl SttStreamSender {
pub async fn send_audio(&mut self, audio: &[f32]) -> Result<()> {
let audio = audio
.iter()
.flat_map(|s| {
let s = (s.clamp(-1.0, 1.0) * 32768.0) as i16;
s.to_le_bytes()
})
.collect::<Vec<u8>>();
self.0
.send_audio(audio)
.await
.context("STT: failed to send audio")?;
Ok(())
}
}
impl SttStreamReceiver {
pub async fn next_message(&mut self) -> Result<Option<Msg>> {
use gradium::protocol::stt::Response;
while let Some(msg) = self.0.next_message().await? {
match msg {
Response::Vad(vad_event) => {
let inactivity_prob = vad_event
.vad
.get(VAD_INDEX)
.map(|v| v.inactivity_prob)
.unwrap_or(1.0);
let end_of_turn = inactivity_prob > 0.8;
return Ok(Some(Msg::Step {
end_of_turn,
current_s: vad_event.total_duration_s,
inactivity_prob,
}));
}
Response::Error { code, message } => {
anyhow::bail!("STT Error {code:?}: {message}")
}
Response::Ready(_) | Response::EndText(_) => {}
Response::Text(text) => {
return Ok(Some(Msg::Text {
text: text.text,
start_s: text.start_s,
}));
}
Response::EndOfStream => break,
}
}
Ok(None)
}
}