- 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.
100 lines
4.4 KiB
Rust
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()
|
|
}
|
|
});
|
|
}
|