botserver/src/basic/keywords/bot_memory.rs
Rodrigo Rodriguez (Pragmatismo) 415448088b feat: refactor prompt compaction and clean up test files
- Renamed `execute_compact_prompt` to `compact_prompt_for_bots` and simplified logic
- Removed redundant comments and empty lines in test files
- Consolidated prompt compaction threshold handling
- Cleaned up UI logging implementation by removing unnecessary whitespace
- Improved code organization in ui_tree module

The changes focus on code quality improvements, removing clutter, and making the prompt compaction logic more straightforward. Test files were cleaned up to be more concise.
2025-11-11 18:32:52 -03:00

100 lines
4.4 KiB
Rust

use crate::shared::models::UserSession;
use crate::shared::state::AppState;
use diesel::prelude::*;
use log::{error, trace};
use rhai::{Dynamic, Engine};
use std::sync::Arc;
use uuid::Uuid;
pub fn set_bot_memory_keyword(state: Arc<AppState>, user: UserSession, engine: &mut Engine) {
let state_clone = Arc::clone(&state);
let user_clone = user.clone();
engine
.register_custom_syntax(&["SET_BOT_MEMORY", "$expr$", ",", "$expr$"], false, move |context, inputs| {
let key = context.eval_expression_tree(&inputs[0])?.to_string();
let value = context.eval_expression_tree(&inputs[1])?.to_string();
let state_for_spawn = Arc::clone(&state_clone);
let user_clone_spawn = user_clone.clone();
let key_clone = key.clone();
let value_clone = value.clone();
tokio::spawn(async move {
use crate::shared::models::bot_memories;
let mut conn = match state_for_spawn.conn.get() {
Ok(conn) => conn,
Err(e) => {
error!("Failed to acquire database connection for SET BOT MEMORY: {}", e);
return;
}
};
let bot_uuid = match Uuid::parse_str(&user_clone_spawn.bot_id.to_string()) {
Ok(uuid) => uuid,
Err(e) => {
error!("Invalid bot ID format: {}", e);
return;
}
};
let now = chrono::Utc::now();
let existing_memory: Option<Uuid> = bot_memories::table
.filter(bot_memories::bot_id.eq(bot_uuid))
.filter(bot_memories::key.eq(&key_clone))
.select(bot_memories::id)
.first(&mut *conn)
.optional()
.unwrap_or(None);
if let Some(memory_id) = existing_memory {
let update_result = diesel::update(bot_memories::table.filter(bot_memories::id.eq(memory_id)))
.set((bot_memories::value.eq(&value_clone), bot_memories::updated_at.eq(now)))
.execute(&mut *conn);
match update_result {
Ok(_) => {
trace!("Updated bot memory for key: {} with value length: {}", key_clone, value_clone.len());
}
Err(e) => {
error!("Failed to update bot memory: {}", e);
}
}
} else {
let new_memory = crate::shared::models::BotMemory {
id: Uuid::new_v4(),
bot_id: bot_uuid,
key: key_clone.clone(),
value: value_clone.clone(),
created_at: now,
updated_at: now,
};
let insert_result = diesel::insert_into(bot_memories::table).values(&new_memory).execute(&mut *conn);
match insert_result {
Ok(_) => {
trace!("Created new bot memory for key: {} with value length: {}", key_clone, value_clone.len());
}
Err(e) => {
error!("Failed to insert bot memory: {}", e);
}
}
}
});
Ok(Dynamic::UNIT)
})
.unwrap();
}
pub fn get_bot_memory_keyword(state: Arc<AppState>, user: UserSession, engine: &mut Engine) {
let state_clone = Arc::clone(&state);
let user_clone = user.clone();
engine.register_fn("GET_BOT_MEMORY", move |key_param: String| -> String {
use crate::shared::models::bot_memories;
let state = Arc::clone(&state_clone);
let conn_result = state.conn.get();
if let Ok(mut conn) = conn_result {
let bot_uuid = user_clone.bot_id;
let memory_value: Option<String> = bot_memories::table
.filter(bot_memories::bot_id.eq(bot_uuid))
.filter(bot_memories::key.eq(&key_param))
.select(bot_memories::value)
.first(&mut *conn)
.optional()
.unwrap_or(None);
memory_value.unwrap_or_default()
} else {
String::new()
}
});
}