use crate::core::shared::models::UserSession; use crate::core::shared::state::AppState; use diesel::prelude::*; use log::{error, trace}; use rhai::{Dynamic, Engine}; use std::sync::Arc; use uuid::Uuid; pub fn register_user_memory_keywords(state: Arc, user: UserSession, engine: &mut Engine) { set_user_memory_keyword(Arc::clone(&state), user.clone(), engine); get_user_memory_keyword(Arc::clone(&state), user.clone(), engine); remember_user_fact_keyword(Arc::clone(&state), user.clone(), engine); get_user_facts_keyword(Arc::clone(&state), user.clone(), engine); clear_user_memory_keyword(state, user, engine); } pub fn set_user_memory_keyword(state: Arc, user: UserSession, engine: &mut Engine) { let state_clone = Arc::clone(&state); let user_id = user.user_id; engine .register_custom_syntax( ["SET", "USER", "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 key_clone = key; let value_clone = value; tokio::spawn(async move { if let Err(e) = set_user_memory( &state_for_spawn, user_id, &key_clone, &value_clone, "preference", ) { error!("Failed to set user memory: {e}"); } else { trace!( "Set user memory for key: {key_clone} with value length: {}", value_clone.len() ); } }); Ok(Dynamic::UNIT) }, ) .expect("Failed to register SET USER MEMORY syntax"); } pub fn get_user_memory_keyword(state: Arc, user: UserSession, engine: &mut Engine) { let state_clone = Arc::clone(&state); let user_id = user.user_id; engine.register_fn("GET USER MEMORY", move |key_param: String| -> String { let state = Arc::clone(&state_clone); let conn_result = state.conn.get(); if let Ok(mut conn) = conn_result { get_user_memory_sync(&mut conn, user_id, &key_param).unwrap_or_default() } else { String::new() } }); } pub fn remember_user_fact_keyword(state: Arc, user: UserSession, engine: &mut Engine) { let state_clone = Arc::clone(&state); let user_id = user.user_id; engine .register_custom_syntax( ["REMEMBER", "USER", "FACT", "$expr$"], false, move |context, inputs| { let fact = context.eval_expression_tree(&inputs[0])?.to_string(); let state_for_spawn = Arc::clone(&state_clone); let fact_clone = fact; tokio::spawn(async move { if let Err(e) = add_user_fact(&state_for_spawn, user_id, &fact_clone) { error!("Failed to remember user fact: {e}"); } else { trace!("Remembered user fact: {fact_clone}"); } }); Ok(Dynamic::UNIT) }, ) .expect("Failed to register REMEMBER USER FACT syntax"); } pub fn get_user_facts_keyword(state: Arc, user: UserSession, engine: &mut Engine) { let state_clone = Arc::clone(&state); let user_id = user.user_id; engine.register_fn("GET USER FACTS", move || -> rhai::Array { let state = Arc::clone(&state_clone); let conn_result = state.conn.get(); if let Ok(mut conn) = conn_result { get_user_facts_sync(&mut conn, user_id) .unwrap_or_default() .into_iter() .map(Dynamic::from) .collect() } else { rhai::Array::new() } }); } pub fn clear_user_memory_keyword(state: Arc, user: UserSession, engine: &mut Engine) { let state_clone = Arc::clone(&state); let user_id = user.user_id; engine .register_custom_syntax( ["CLEAR", "USER", "MEMORY"], false, move |_context, _inputs| { let state_for_spawn = Arc::clone(&state_clone); tokio::spawn(async move { if let Err(e) = clear_user_memory(&state_for_spawn, user_id) { error!("Failed to clear user memory: {e}"); } else { trace!("Cleared all user memory for user: {user_id}"); } }); Ok(Dynamic::UNIT) }, ) .expect("Failed to register CLEAR USER MEMORY syntax"); } fn set_user_memory( state: &AppState, user_id: Uuid, key: &str, value: &str, memory_type: &str, ) -> Result<(), String> { let mut conn = state .conn .get() .map_err(|e| format!("Failed to acquire database connection: {e}"))?; let now = chrono::Utc::now(); let new_id = Uuid::new_v4(); diesel::sql_query( "INSERT INTO user_memories (id, user_id, key, value, memory_type, created_at, updated_at) \ VALUES ($1, $2, $3, $4, $5, $6, $7) \ ON CONFLICT (user_id, key) DO UPDATE SET \ value = EXCLUDED.value, \ memory_type = EXCLUDED.memory_type, \ updated_at = EXCLUDED.updated_at", ) .bind::(new_id) .bind::(user_id) .bind::(key) .bind::(value) .bind::(memory_type) .bind::(now) .bind::(now) .execute(&mut conn) .map_err(|e| format!("Failed to set user memory: {e}"))?; Ok(()) } fn get_user_memory_sync( conn: &mut diesel::PgConnection, user_id: Uuid, key: &str, ) -> Result { #[derive(QueryableByName)] struct MemoryValue { #[diesel(sql_type = diesel::sql_types::Text)] value: String, } let result: Option = diesel::sql_query( "SELECT value FROM user_memories WHERE user_id = $1 AND key = $2 LIMIT 1", ) .bind::(user_id) .bind::(key) .get_result(conn) .optional() .map_err(|e| format!("Failed to get user memory: {e}"))?; Ok(result.map(|r| r.value).unwrap_or_default()) } fn add_user_fact(state: &AppState, user_id: Uuid, fact: &str) -> Result<(), String> { let mut conn = state .conn .get() .map_err(|e| format!("Failed to acquire database connection: {e}"))?; let now = chrono::Utc::now(); let new_id = Uuid::new_v4(); let fact_key = format!("fact_{}", Uuid::new_v4()); diesel::sql_query( "INSERT INTO user_memories (id, user_id, key, value, memory_type, created_at, updated_at) \ VALUES ($1, $2, $3, $4, 'fact', $5, $6)", ) .bind::(new_id) .bind::(user_id) .bind::(&fact_key) .bind::(fact) .bind::(now) .bind::(now) .execute(&mut conn) .map_err(|e| format!("Failed to add user fact: {e}"))?; Ok(()) } fn get_user_facts_sync( conn: &mut diesel::PgConnection, user_id: Uuid, ) -> Result, String> { #[derive(QueryableByName)] struct FactValue { #[diesel(sql_type = diesel::sql_types::Text)] value: String, } let results: Vec = diesel::sql_query( "SELECT value FROM user_memories WHERE user_id = $1 AND memory_type = 'fact' ORDER BY created_at DESC", ) .bind::(user_id) .load(conn) .map_err(|e| format!("Failed to get user facts: {e}"))?; Ok(results.into_iter().map(|r| r.value).collect()) } fn clear_user_memory(state: &AppState, user_id: Uuid) -> Result<(), String> { let mut conn = state .conn .get() .map_err(|e| format!("Failed to acquire database connection: {e}"))?; diesel::sql_query("DELETE FROM user_memories WHERE user_id = $1") .bind::(user_id) .execute(&mut conn) .map_err(|e| format!("Failed to clear user memory: {e}"))?; Ok(()) }