ceylon_next/memory/
in_memory.rs

1//! In-memory implementation of the Memory trait.
2//!
3//! This module provides a fast, volatile storage backend suitable for
4//! development and testing. Data is lost when the program exits.
5
6use super::{Memory, MemoryEntry};
7use async_trait::async_trait;
8use std::collections::HashMap;
9use std::sync::{Arc, RwLock};
10
11/// An in-memory implementation of the [`Memory`] trait.
12///
13/// This is the default memory storage used by agents. It stores all
14/// conversations in RAM and is suitable for development and testing.
15/// For production use with persistence, use [`SqliteStore`](super::SqliteStore) or [`FileStore`](super::FileStore).
16///
17/// # Examples
18///
19/// ```rust
20/// use ceylon_next::memory::InMemoryStore;
21/// use std::sync::Arc;
22///
23/// let memory = Arc::new(InMemoryStore::new());
24/// ```
25pub struct InMemoryStore {
26    // Store entries by ID
27    entries: Arc<RwLock<HashMap<String, MemoryEntry>>>,
28
29    // Index by agent_id for fast lookup
30    agent_index: Arc<RwLock<HashMap<String, Vec<String>>>>,
31}
32
33impl InMemoryStore {
34    /// Creates a new in-memory store.
35    ///
36    /// # Examples
37    ///
38    /// ```rust
39    /// use ceylon_next::memory::InMemoryStore;
40    ///
41    /// let store = InMemoryStore::new();
42    /// ```
43    pub fn new() -> Self {
44        Self {
45            entries: Arc::new(RwLock::new(HashMap::new())),
46            agent_index: Arc::new(RwLock::new(HashMap::new())),
47        }
48    }
49}
50
51#[async_trait]
52impl Memory for InMemoryStore {
53    async fn store(&self, entry: MemoryEntry) -> Result<String, String> {
54        let entry_id = entry.id.clone();
55        let agent_id = entry.agent_id.clone();
56
57        // Store the entry
58        {
59            let mut entries = self.entries.write().unwrap();
60            entries.insert(entry_id.clone(), entry);
61        }
62
63        // Update agent index
64        {
65            let mut index = self.agent_index.write().unwrap();
66            index
67                .entry(agent_id)
68                .or_insert_with(Vec::new)
69                .push(entry_id.clone());
70        }
71
72        Ok(entry_id)
73    }
74
75    async fn get(&self, id: &str) -> Result<Option<MemoryEntry>, String> {
76        let entries = self.entries.read().unwrap();
77        Ok(entries.get(id).cloned())
78    }
79
80    async fn get_agent_history(&self, agent_id: &str) -> Result<Vec<MemoryEntry>, String> {
81        let index = self.agent_index.read().unwrap();
82        let entries = self.entries.read().unwrap();
83
84        // Get all entry IDs for this agent
85        let entry_ids = index.get(agent_id);
86
87        if entry_ids.is_none() {
88            return Ok(Vec::new());
89        }
90
91        // Collect all entries
92        let mut result = Vec::new();
93        for entry_id in entry_ids.unwrap() {
94            if let Some(entry) = entries.get(entry_id) {
95                result.push(entry.clone());
96            }
97        }
98
99        // Sort by timestamp (newest first)
100        result.sort_by(|a, b| b.created_at.cmp(&a.created_at));
101
102        Ok(result)
103    }
104
105    async fn get_recent(&self, agent_id: &str, limit: usize) -> Result<Vec<MemoryEntry>, String> {
106        let mut history = self.get_agent_history(agent_id).await?;
107
108        // Take only the recent ones
109        history.truncate(limit);
110
111        Ok(history)
112    }
113
114    async fn search(&self, agent_id: &str, query: &str) -> Result<Vec<MemoryEntry>, String> {
115        let history = self.get_agent_history(agent_id).await?;
116        let query_lower = query.to_lowercase();
117
118        // Simple text search in messages
119        let results: Vec<MemoryEntry> = history
120            .into_iter()
121            .filter(|entry| {
122                entry
123                    .messages
124                    .iter()
125                    .any(|msg| msg.content.to_lowercase().contains(&query_lower))
126            })
127            .collect();
128
129        Ok(results)
130    }
131
132    async fn clear_agent_memory(&self, agent_id: &str) -> Result<(), String> {
133        let index = self.agent_index.read().unwrap();
134        let entry_ids = index.get(agent_id);
135
136        if let Some(ids) = entry_ids {
137            let mut entries = self.entries.write().unwrap();
138            for id in ids {
139                entries.remove(id);
140            }
141        }
142
143        let mut index = self.agent_index.write().unwrap();
144        index.remove(agent_id);
145
146        Ok(())
147    }
148}
149
150impl Default for InMemoryStore {
151    fn default() -> Self {
152        Self::new()
153    }
154}