Skip to main content

rskit_agent/memory/
store.rs

1use 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// ── InMemoryStore ───────────────────────────────────────────────────────────
11
12/// A simple in-process memory store backed by a `RwLock<HashMap>`.
13#[derive(Debug, Default)]
14pub struct InMemoryStore {
15    store: RwLock<HashMap<String, Vec<Message>>>,
16}
17
18impl InMemoryStore {
19    /// Create an empty in-memory message store.
20    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}