use std::sync::Arc;
use async_trait::async_trait;
use futures::StreamExt;
use serde_json::json;
use tokio::sync::Mutex;
use crate::error::Result;
use crate::processor::frame::{Frame, LlmContext, ServiceKind, StartParams};
use crate::processor::{Envelope, FrameProcessor, Link, ProcessorSetup};
use crate::service::{LlmService, SttService, TtsService};
pub struct SttProcessor<S: SttService> {
svc: Arc<Mutex<S>>,
}
impl<S: SttService> SttProcessor<S> {
pub fn new(svc: S) -> Self {
Self {
svc: Arc::new(Mutex::new(svc)),
}
}
}
#[async_trait]
impl<S: SttService + 'static> FrameProcessor for SttProcessor<S> {
fn name(&self) -> &str {
"Stt"
}
async fn start(&mut self, _s: &ProcessorSetup, p: &StartParams) -> Result<()> {
self.svc.lock().await.start(p).await
}
fn can_generate_metrics(&self) -> bool {
true
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
match &env.frame {
Frame::InputAudio(audio) => {
let frames = self.svc.lock().await.run_stt(audio.clone()).await;
match frames {
Ok(frames) => {
for f in frames {
link.push_down(f).await;
}
}
Err(e) => link.push_error(format!("stt: {e}"), false).await,
}
}
Frame::SttMute(muted) => {
self.svc.lock().await.set_muted(*muted).await;
link.push(env.meta, env.frame, env.direction).await;
}
_ => link.push(env.meta, env.frame, env.direction).await,
}
Ok(())
}
}
pub struct LlmProcessor<L: LlmService> {
svc: Arc<Mutex<L>>,
}
impl<L: LlmService> LlmProcessor<L> {
pub fn new(svc: L) -> Self {
Self {
svc: Arc::new(Mutex::new(svc)),
}
}
async fn run(&self, ctx: &LlmContext, link: &Link) {
tracing::debug!(messages = ctx.messages.len(), "cascaded LLM run");
let mut guard = self.svc.lock().await;
let mut pushed = 0usize;
let err = match guard.run_llm(ctx).await {
Ok(mut stream) => {
while let Some(f) = stream.next().await {
pushed += 1;
link.push_down(f).await;
}
None
}
Err(e) => Some(format!("llm: {e}")),
};
drop(guard);
match &err {
Some(msg) => tracing::warn!(error = %msg, "cascaded LLM error"),
None => tracing::debug!(frames = pushed, "cascaded LLM produced frames"),
}
if let Some(msg) = err {
link.push_error(msg, false).await;
}
}
}
#[async_trait]
impl<L: LlmService + 'static> FrameProcessor for LlmProcessor<L> {
fn name(&self) -> &str {
"Llm"
}
async fn start(&mut self, _s: &ProcessorSetup, p: &StartParams) -> Result<()> {
self.svc.lock().await.start(p).await
}
fn can_generate_metrics(&self) -> bool {
true
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
match &env.frame {
Frame::LlmContext(ctx) => {
self.run(ctx, link).await;
}
Frame::Transcription {
text, final_: true, ..
} => {
let ctx = LlmContext {
messages: vec![json!({"role": "user", "content": text})],
tools: vec![],
};
self.run(&ctx, link).await;
}
Frame::UpdateSettings {
target: ServiceKind::Llm,
settings,
} => {
if let Some(tools) = settings.get("tools").and_then(|t| t.as_array()) {
let tools = tools
.iter()
.filter_map(|t| serde_json::from_value(t.clone()).ok())
.collect();
self.svc.lock().await.set_tools(tools);
}
link.push(env.meta, env.frame, env.direction).await;
}
_ => link.push(env.meta, env.frame, env.direction).await,
}
Ok(())
}
}
pub struct TtsProcessor<T: TtsService> {
svc: Arc<Mutex<T>>,
}
impl<T: TtsService> TtsProcessor<T> {
pub fn new(svc: T) -> Self {
Self {
svc: Arc::new(Mutex::new(svc)),
}
}
async fn run(&self, text: &str, link: &Link) {
tracing::debug!(chars = text.len(), "cascaded TTS run");
let frames = self.svc.lock().await.run_tts(text).await;
match frames {
Ok(frames) => {
tracing::debug!(frames = frames.len(), "cascaded TTS produced frames");
for f in frames {
let f = match f {
Frame::TtsAudio { audio, .. } => Frame::OutputAudio(audio),
other => other,
};
link.push_down(f).await;
}
}
Err(e) => {
tracing::warn!(error = %e, "cascaded TTS error");
link.push_error(format!("tts: {e}"), false).await;
}
}
}
}
#[async_trait]
impl<T: TtsService + 'static> FrameProcessor for TtsProcessor<T> {
fn name(&self) -> &str {
"Tts"
}
async fn start(&mut self, _s: &ProcessorSetup, p: &StartParams) -> Result<()> {
self.svc.lock().await.start(p).await
}
fn can_generate_metrics(&self) -> bool {
true
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
match &env.frame {
Frame::TtsSpeak { text, .. } => {
self.run(text, link).await;
}
Frame::LlmText(text) => {
self.run(text, link).await;
}
_ => link.push(env.meta, env.frame, env.direction).await,
}
Ok(())
}
}