use crate::multimodal::BotModelsClient; use crate::shared::models::UserSession; use crate::shared::state::AppState; use log::{error, trace}; use rhai::{Dynamic, Engine}; use std::sync::Arc; use std::time::Duration; pub fn register_multimodal_keywords(state: Arc, user: UserSession, engine: &mut Engine) { image_keyword(state.clone(), user.clone(), engine); video_keyword(state.clone(), user.clone(), engine); audio_keyword(state.clone(), user.clone(), engine); see_keyword(state.clone(), user.clone(), engine); } pub fn image_keyword(state: Arc, user: UserSession, engine: &mut Engine) { let state_clone = Arc::clone(&state); let user_clone = user.clone(); engine .register_custom_syntax(&["IMAGE", "$expr$"], false, move |context, inputs| { let prompt = context.eval_expression_tree(&inputs[0])?.to_string(); trace!("IMAGE keyword: generating image for prompt: {}", prompt); let state_for_thread = Arc::clone(&state_clone); let bot_id = user_clone.bot_id; let (tx, rx) = std::sync::mpsc::channel(); std::thread::spawn(move || { let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .enable_all() .build(); let send_err = if let Ok(rt) = rt { let result = rt.block_on(async move { execute_image_generation(state_for_thread, bot_id, prompt).await }); tx.send(result).err() } else { tx.send(Err("Failed to build tokio runtime".into())).err() }; if send_err.is_some() { error!("Failed to send IMAGE result"); } }); match rx.recv_timeout(Duration::from_secs(300)) { Ok(Ok(result)) => Ok(Dynamic::from(result)), Ok(Err(e)) => Err(Box::new(rhai::EvalAltResult::ErrorRuntime( e.to_string().into(), rhai::Position::NONE, ))), Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { Err(Box::new(rhai::EvalAltResult::ErrorRuntime( "Image generation timed out".into(), rhai::Position::NONE, ))) } Err(e) => Err(Box::new(rhai::EvalAltResult::ErrorRuntime( format!("IMAGE thread failed: {}", e).into(), rhai::Position::NONE, ))), } }) .unwrap(); } async fn execute_image_generation( state: Arc, bot_id: uuid::Uuid, prompt: String, ) -> Result> { let client = BotModelsClient::from_state(&state, &bot_id); if !client.is_enabled() { return Err("BotModels is not enabled. Set botmodels-enabled=true in config.csv".into()); } client.generate_image(&prompt).await } pub fn video_keyword(state: Arc, user: UserSession, engine: &mut Engine) { let state_clone = Arc::clone(&state); let user_clone = user.clone(); engine .register_custom_syntax(&["VIDEO", "$expr$"], false, move |context, inputs| { let prompt = context.eval_expression_tree(&inputs[0])?.to_string(); trace!("VIDEO keyword: generating video for prompt: {}", prompt); let state_for_thread = Arc::clone(&state_clone); let bot_id = user_clone.bot_id; let (tx, rx) = std::sync::mpsc::channel(); std::thread::spawn(move || { let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .enable_all() .build(); let send_err = if let Ok(rt) = rt { let result = rt.block_on(async move { execute_video_generation(state_for_thread, bot_id, prompt).await }); tx.send(result).err() } else { tx.send(Err("Failed to build tokio runtime".into())).err() }; if send_err.is_some() { error!("Failed to send VIDEO result"); } }); match rx.recv_timeout(Duration::from_secs(600)) { Ok(Ok(result)) => Ok(Dynamic::from(result)), Ok(Err(e)) => Err(Box::new(rhai::EvalAltResult::ErrorRuntime( e.to_string().into(), rhai::Position::NONE, ))), Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { Err(Box::new(rhai::EvalAltResult::ErrorRuntime( "Video generation timed out".into(), rhai::Position::NONE, ))) } Err(e) => Err(Box::new(rhai::EvalAltResult::ErrorRuntime( format!("VIDEO thread failed: {}", e).into(), rhai::Position::NONE, ))), } }) .unwrap(); } async fn execute_video_generation( state: Arc, bot_id: uuid::Uuid, prompt: String, ) -> Result> { let client = BotModelsClient::from_state(&state, &bot_id); if !client.is_enabled() { return Err("BotModels is not enabled. Set botmodels-enabled=true in config.csv".into()); } client.generate_video(&prompt).await } pub fn audio_keyword(state: Arc, user: UserSession, engine: &mut Engine) { let state_clone = Arc::clone(&state); let user_clone = user.clone(); engine .register_custom_syntax(&["AUDIO", "$expr$"], false, move |context, inputs| { let text = context.eval_expression_tree(&inputs[0])?.to_string(); trace!("AUDIO keyword: generating speech for text: {}", text); let state_for_thread = Arc::clone(&state_clone); let bot_id = user_clone.bot_id; let (tx, rx) = std::sync::mpsc::channel(); std::thread::spawn(move || { let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .enable_all() .build(); let send_err = if let Ok(rt) = rt { let result = rt.block_on(async move { execute_audio_generation(state_for_thread, bot_id, text).await }); tx.send(result).err() } else { tx.send(Err("Failed to build tokio runtime".into())).err() }; if send_err.is_some() { error!("Failed to send AUDIO result"); } }); match rx.recv_timeout(Duration::from_secs(120)) { Ok(Ok(result)) => Ok(Dynamic::from(result)), Ok(Err(e)) => Err(Box::new(rhai::EvalAltResult::ErrorRuntime( e.to_string().into(), rhai::Position::NONE, ))), Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { Err(Box::new(rhai::EvalAltResult::ErrorRuntime( "Audio generation timed out".into(), rhai::Position::NONE, ))) } Err(e) => Err(Box::new(rhai::EvalAltResult::ErrorRuntime( format!("AUDIO thread failed: {}", e).into(), rhai::Position::NONE, ))), } }) .unwrap(); } async fn execute_audio_generation( state: Arc, bot_id: uuid::Uuid, text: String, ) -> Result> { let client = BotModelsClient::from_state(&state, &bot_id); if !client.is_enabled() { return Err("BotModels is not enabled. Set botmodels-enabled=true in config.csv".into()); } client.generate_audio(&text, None, None).await } pub fn see_keyword(state: Arc, user: UserSession, engine: &mut Engine) { let state_clone = Arc::clone(&state); let user_clone = user.clone(); engine .register_custom_syntax(&["SEE", "$expr$"], false, move |context, inputs| { let file_path = context.eval_expression_tree(&inputs[0])?.to_string(); trace!("SEE keyword: getting caption for file: {}", file_path); let state_for_thread = Arc::clone(&state_clone); let bot_id = user_clone.bot_id; let (tx, rx) = std::sync::mpsc::channel(); std::thread::spawn(move || { let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .enable_all() .build(); let send_err = if let Ok(rt) = rt { let result = rt.block_on(async move { execute_see_caption(state_for_thread, bot_id, file_path).await }); tx.send(result).err() } else { tx.send(Err("Failed to build tokio runtime".into())).err() }; if send_err.is_some() { error!("Failed to send SEE result"); } }); match rx.recv_timeout(Duration::from_secs(60)) { Ok(Ok(result)) => Ok(Dynamic::from(result)), Ok(Err(e)) => Err(Box::new(rhai::EvalAltResult::ErrorRuntime( e.to_string().into(), rhai::Position::NONE, ))), Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { Err(Box::new(rhai::EvalAltResult::ErrorRuntime( "Vision/caption timed out".into(), rhai::Position::NONE, ))) } Err(e) => Err(Box::new(rhai::EvalAltResult::ErrorRuntime( format!("SEE thread failed: {}", e).into(), rhai::Position::NONE, ))), } }) .unwrap(); } async fn execute_see_caption( state: Arc, bot_id: uuid::Uuid, file_path: String, ) -> Result> { let client = BotModelsClient::from_state(&state, &bot_id); if !client.is_enabled() { return Err("BotModels is not enabled. Set botmodels-enabled=true in config.csv".into()); } let lower_path = file_path.to_lowercase(); if lower_path.ends_with(".mp4") || lower_path.ends_with(".avi") || lower_path.ends_with(".mov") || lower_path.ends_with(".webm") || lower_path.ends_with(".mkv") { client.describe_video(&file_path).await } else { client.describe_image(&file_path).await } }