rskit_agent/memory/
store.rs1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use parking_lot::RwLock;
5use rskit_errors::AppError;
6use rskit_llm::types::Message;
7
8use super::Memory;
9
10#[derive(Debug, Default)]
14pub struct InMemoryStore {
15 store: RwLock<HashMap<String, Vec<Message>>>,
16}
17
18impl InMemoryStore {
19 pub fn new() -> Self {
21 Self::default()
22 }
23}
24
25#[async_trait]
26impl Memory for InMemoryStore {
27 async fn load(&self, session_id: &str) -> Result<Vec<Message>, AppError> {
28 let guard = self.store.read();
29 Ok(guard.get(session_id).cloned().unwrap_or_default())
30 }
31
32 async fn save(&self, session_id: &str, messages: &[Message]) -> Result<(), AppError> {
33 let mut guard = self.store.write();
34 guard.insert(session_id.to_string(), messages.to_vec());
35 drop(guard);
36 Ok(())
37 }
38
39 async fn append(&self, session_id: &str, messages: &[Message]) -> Result<(), AppError> {
40 let mut guard = self.store.write();
41 guard
42 .entry(session_id.to_string())
43 .or_default()
44 .extend_from_slice(messages);
45 drop(guard);
46 Ok(())
47 }
48
49 async fn clear(&self, session_id: &str) -> Result<(), AppError> {
50 let mut guard = self.store.write();
51 guard.remove(session_id);
52 drop(guard);
53 Ok(())
54 }
55}