botserver/src/basic/keywords/add_tool.rs

116 lines
5 KiB
Rust
Raw Normal View History

2025-10-18 18:20:02 -03:00
use crate::shared::models::UserSession;
use crate::shared::state::AppState;
use diesel::prelude::*;
use log::{error, trace, warn};
2025-10-18 18:20:02 -03:00
use rhai::{Dynamic, Engine};
use std::sync::Arc;
use uuid::Uuid;
pub fn add_tool_keyword(state: Arc<AppState>, user: UserSession, engine: &mut Engine) {
let state_clone = Arc::clone(&state);
let user_clone = user.clone();
engine
.register_custom_syntax(&["ADD_TOOL", "$expr$"], false, move |context, inputs| {
let tool_path = context.eval_expression_tree(&inputs[0])?;
let tool_path_str = tool_path.to_string().trim_matches('"').to_string();
trace!("ADD_TOOL command executed: {} for session: {}", tool_path_str, user_clone.id);
let tool_name = tool_path_str.strip_prefix(".gbdialog/").unwrap_or(&tool_path_str).strip_suffix(".bas").unwrap_or(&tool_path_str).to_string();
if tool_name.is_empty() {
return Err(Box::new(rhai::EvalAltResult::ErrorRuntime("Invalid tool name".into(), rhai::Position::NONE)));
}
let state_for_task = Arc::clone(&state_clone);
let user_for_task = user_clone.clone();
let tool_name_for_task = tool_name.clone();
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 {
associate_tool_with_session(&state_for_task, &user_for_task, &tool_name_for_task).await
});
tx.send(result).err()
} else {
tx.send(Err("Failed to build tokio runtime".to_string())).err()
};
if send_err.is_some() {
error!("Failed to send result from thread");
}
});
match rx.recv_timeout(std::time::Duration::from_secs(10)) {
Ok(Ok(message)) => {
Ok(Dynamic::from(message))
}
Ok(Err(e)) => Err(Box::new(rhai::EvalAltResult::ErrorRuntime(e.into(), rhai::Position::NONE))),
Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
Err(Box::new(rhai::EvalAltResult::ErrorRuntime("ADD_TOOL timed out".into(), rhai::Position::NONE)))
}
Err(e) => Err(Box::new(rhai::EvalAltResult::ErrorRuntime(format!("ADD_TOOL failed: {}", e).into(), rhai::Position::NONE))),
}
})
.unwrap();
2025-10-18 18:20:02 -03:00
}
async fn associate_tool_with_session(state: &AppState, user: &UserSession, tool_name: &str) -> Result<String, String> {
use crate::shared::models::schema::{basic_tools, session_tool_associations};
let mut conn = state.conn.get().map_err(|e| format!("DB error: {}", e))?;
let tool_exists: Result<bool, diesel::result::Error> = basic_tools::table
.filter(basic_tools::bot_id.eq(user.bot_id.to_string()))
.filter(basic_tools::tool_name.eq(tool_name))
.filter(basic_tools::is_active.eq(1))
.select(diesel::dsl::count(basic_tools::id))
.first::<i64>(&mut *conn)
.map(|count| count > 0);
match tool_exists {
Ok(true) => {
trace!("Tool '{}' exists and is active for bot '{}'", tool_name, user.bot_id);
}
Ok(false) => {
warn!("Tool '{}' does not exist or is not active for bot '{}'", tool_name, user.bot_id);
return Err(format!("Tool '{}' is not available. Make sure the tool file is compiled and active.", tool_name));
}
Err(e) => {
error!("Failed to check tool existence: {}", e);
return Err(format!("Database error while checking tool: {}", e));
}
}
let association_id = Uuid::new_v4().to_string();
let session_id_str = user.id.to_string();
let added_at = chrono::Utc::now().to_rfc3339();
let insert_result: Result<usize, diesel::result::Error> = diesel::insert_into(session_tool_associations::table)
.values((
session_tool_associations::id.eq(&association_id),
session_tool_associations::session_id.eq(&session_id_str),
session_tool_associations::tool_name.eq(tool_name),
session_tool_associations::added_at.eq(&added_at),
))
.on_conflict((session_tool_associations::session_id, session_tool_associations::tool_name))
.do_nothing()
.execute(&mut *conn);
match insert_result {
Ok(rows_affected) => {
if rows_affected > 0 {
trace!("Tool '{}' newly associated with session '{}' (user: {}, bot: {})", tool_name, user.id, user.user_id, user.bot_id);
Ok(format!("Tool '{}' is now available in this conversation", tool_name))
} else {
trace!("Tool '{}' was already associated with session '{}'", tool_name, user.id);
Ok(format!("Tool '{}' is already available in this conversation", tool_name))
}
}
Err(e) => {
error!("Failed to associate tool '{}' with session '{}': {}", tool_name, user.id, e);
Err(format!("Failed to add tool to session: {}", e))
}
}
2025-10-18 18:20:02 -03:00
}
pub fn get_session_tools(conn: &mut PgConnection, session_id: &Uuid) -> Result<Vec<String>, diesel::result::Error> {
use crate::shared::models::schema::session_tool_associations;
let session_id_str = session_id.to_string();
session_tool_associations::table
.filter(session_tool_associations::session_id.eq(&session_id_str))
.select(session_tool_associations::tool_name)
.load::<String>(conn)
2025-10-18 18:20:02 -03:00
}
pub fn clear_session_tools(conn: &mut PgConnection, session_id: &Uuid) -> Result<usize, diesel::result::Error> {
use crate::shared::models::schema::session_tool_associations;
let session_id_str = session_id.to_string();
diesel::delete(session_tool_associations::table.filter(session_tool_associations::session_id.eq(&session_id_str))).execute(conn)
2025-10-18 18:20:02 -03:00
}