use crate::{
db_connectors::mongodb::get_db,
encrypt::{decrypt_data, encrypt_data},
Client, ConversationInfo, EngineError, Memory,
};
use bson::{doc, Bson};
fn format_memories(
data: &mut ConversationInfo,
memories: &[Memory],
interaction_order: i32,
) -> Result<Vec<bson::Document>, EngineError> {
let client = bson::to_bson(&data.client)?;
memories
.iter()
.enumerate()
.fold(Ok(vec![]), |vec, (memory_order, var)| {
let time = Bson::UtcDatetime(chrono::Utc::now());
let value = encrypt_data(&var.value)?;
let mut vec = vec?;
vec.push(doc! {
"client": client.clone(),
"interaction_id": &data.interaction_id,
"conversation_id": &data.conversation_id,
"flow_id": &data.context.flow,
"step_id": &data.context.step,
"memory_order": memory_order as i32,
"interaction_order": interaction_order,
"key": &var.key,
"value": value, "expires_at": Bson::Null,
"created_at": time
});
Ok(vec)
})
}
pub fn add_memories(
data: &mut ConversationInfo,
memories: &[Memory],
interaction_order: i32,
) -> Result<(), EngineError> {
if memories.is_empty() {
return Ok(());
}
let mem = format_memories(data, memories, interaction_order)?;
let db = get_db(&data.db)?;
let collection = db.collection("memory");
collection.insert_many(mem, None)?;
Ok(())
}
pub fn get_memories(
client: &Client,
db: &mongodb::Database,
) -> Result<serde_json::Value, EngineError> {
let collection = db.collection("memory");
let filter = doc! {
"client": bson::to_bson(&client)?,
};
let find_options = mongodb::options::FindOptions::builder()
.sort(doc! { "$natural": -1 })
.build();
let cursor = collection.find(filter, find_options)?;
let mut map = serde_json::Map::new();
for elem in cursor {
if let Ok(doc) = elem {
let mem: serde_json::Value = bson::from_bson(bson::Bson::Document(doc))?;
let value: serde_json::Value = decrypt_data(mem["value"].as_str().unwrap().to_owned())?;
if !map.contains_key(mem["key"].as_str().unwrap()) {
map.insert(mem["key"].as_str().unwrap().to_owned(), value);
}
}
}
Ok(serde_json::json!(map))
}