spring_ai_rs/ai_interface/
memory.rs

1use std::{collections::HashMap, sync::RwLock};
2
3use lazy_static::lazy_static;
4use serde::{de::DeserializeOwned, Serialize};
5
6use crate::ai_interface::AIInterface;
7
8lazy_static! {
9    static ref AI_STORAGE: RwLock<HashMap<i32, HashMap<String, Vec<u8>>>> =
10        RwLock::new(HashMap::new());
11}
12
13impl AIInterface {
14    pub fn store_memory<T: Serialize + DeserializeOwned>(&self, name: &str, data: &T) {
15        let ai_id = self.ai_id;
16
17        {
18            let mut storage = AI_STORAGE.try_write().unwrap();
19            storage
20                .entry(ai_id)
21                .or_insert_with(Default::default)
22                .insert(name.to_string(), bincode::serialize(data).unwrap());
23        }
24    }
25
26    pub fn load_memory<T: Serialize + DeserializeOwned>(&self, name: &str) -> Option<T> {
27        let ai_id = self.ai_id;
28
29        {
30            let mut storage = AI_STORAGE.try_write().unwrap();
31            storage
32                .entry(ai_id)
33                .or_insert_with(Default::default)
34                .get(&name.to_string())
35                .map(|d| bincode::deserialize(d).unwrap())
36        }
37    }
38}