//! Memory Cache //! //! Provides caching for frequently accessed memories to improve //! retrieval performance. use crate::types::{MemoryEntry, MemoryType}; use std::collections::HashMap; use std::time::{Duration, Instant}; use tokio::sync::RwLock; /// Cache entry with metadata struct CacheEntry { /// The memory entry entry: MemoryEntry, /// Last access time last_accessed: Instant, /// Access count access_count: u32, } /// Cache key for efficient lookups #[derive(Debug, Clone, Hash, Eq, PartialEq)] struct CacheKey { agent_id: String, memory_type: MemoryType, category: String, } impl From<&MemoryEntry> for CacheKey { fn from(entry: &MemoryEntry) -> Self { // Parse URI to extract components let parts: Vec<&str> = entry.uri.trim_start_matches("agent://").split('/').collect(); Self { agent_id: parts.first().unwrap_or(&"").to_string(), memory_type: entry.memory_type, category: parts.get(2).unwrap_or(&"").to_string(), } } } /// Memory cache configuration #[derive(Debug, Clone)] pub struct CacheConfig { /// Maximum number of entries pub max_entries: usize, /// Time-to-live for entries pub ttl: Duration, /// Enable/disable caching pub enabled: bool, } impl Default for CacheConfig { fn default() -> Self { Self { max_entries: 1000, ttl: Duration::from_secs(3600), // 1 hour enabled: true, } } } /// Memory cache for hot memories pub struct MemoryCache { /// Cache storage cache: RwLock>, /// Configuration config: CacheConfig, /// Cache statistics stats: RwLock, } /// Cache statistics #[derive(Debug, Clone, Default)] pub struct CacheStats { /// Total cache hits pub hits: u64, /// Total cache misses pub misses: u64, /// Total entries evicted pub evictions: u64, } impl MemoryCache { /// Create a new memory cache pub fn new(config: CacheConfig) -> Self { Self { cache: RwLock::new(HashMap::new()), config, stats: RwLock::new(CacheStats::default()), } } /// Create with default configuration pub fn default_config() -> Self { Self::new(CacheConfig::default()) } /// Get a memory from cache pub async fn get(&self, uri: &str) -> Option { if !self.config.enabled { return None; } let mut cache = self.cache.write().await; if let Some(cached) = cache.get_mut(uri) { // Check TTL if cached.last_accessed.elapsed() > self.config.ttl { cache.remove(uri); return None; } // Update access metadata cached.last_accessed = Instant::now(); cached.access_count += 1; // Update stats let mut stats = self.stats.write().await; stats.hits += 1; return Some(cached.entry.clone()); } // Update stats let mut stats = self.stats.write().await; stats.misses += 1; None } /// Put a memory into cache pub async fn put(&self, entry: MemoryEntry) { if !self.config.enabled { return; } let mut cache = self.cache.write().await; // Check capacity and evict if necessary if cache.len() >= self.config.max_entries { self.evict_lru(&mut cache).await; } cache.insert( entry.uri.clone(), CacheEntry { entry, last_accessed: Instant::now(), access_count: 0, }, ); } /// Remove a memory from cache pub async fn remove(&self, uri: &str) { let mut cache = self.cache.write().await; cache.remove(uri); } /// Clear the cache pub async fn clear(&self) { let mut cache = self.cache.write().await; cache.clear(); } /// Evict least recently used entries async fn evict_lru(&self, cache: &mut HashMap) { // Find LRU entry let lru_key = cache .iter() .min_by_key(|(_, v)| (v.access_count, v.last_accessed)) .map(|(k, _)| k.clone()); if let Some(key) = lru_key { cache.remove(&key); let mut stats = self.stats.write().await; stats.evictions += 1; } } /// Get cache statistics pub async fn stats(&self) -> CacheStats { self.stats.read().await.clone() } /// Get cache hit rate pub async fn hit_rate(&self) -> f32 { let stats = self.stats.read().await; let total = stats.hits + stats.misses; if total == 0 { return 0.0; } stats.hits as f32 / total as f32 } /// Get cache size pub async fn size(&self) -> usize { self.cache.read().await.len() } /// Warm up cache with frequently accessed entries pub async fn warmup(&self, entries: Vec) { for entry in entries { self.put(entry).await; } } /// Get top accessed entries (for preloading) pub async fn get_hot_entries(&self, limit: usize) -> Vec { let cache = self.cache.read().await; let mut entries: Vec<_> = cache .values() .map(|c| (c.access_count, c.entry.clone())) .collect(); entries.sort_by(|a, b| b.0.cmp(&a.0)); entries.truncate(limit); entries.into_iter().map(|(_, e)| e).collect() } } #[cfg(test)] mod tests { use super::*; use crate::types::MemoryType; #[tokio::test] async fn test_cache_put_and_get() { let cache = MemoryCache::default_config(); let entry = MemoryEntry::new( "test-agent", MemoryType::Preference, "style", "User prefers concise responses".to_string(), ); cache.put(entry.clone()).await; let retrieved = cache.get(&entry.uri).await; assert!(retrieved.is_some()); assert_eq!(retrieved.unwrap().content, "User prefers concise responses"); } #[tokio::test] async fn test_cache_miss() { let cache = MemoryCache::default_config(); let retrieved = cache.get("nonexistent").await; assert!(retrieved.is_none()); let stats = cache.stats().await; assert_eq!(stats.misses, 1); } #[tokio::test] async fn test_cache_remove() { let cache = MemoryCache::default_config(); let entry = MemoryEntry::new( "test-agent", MemoryType::Preference, "style", "test".to_string(), ); cache.put(entry.clone()).await; cache.remove(&entry.uri).await; let retrieved = cache.get(&entry.uri).await; assert!(retrieved.is_none()); } #[tokio::test] async fn test_cache_clear() { let cache = MemoryCache::default_config(); let entry = MemoryEntry::new( "test-agent", MemoryType::Preference, "style", "test".to_string(), ); cache.put(entry).await; cache.clear().await; let size = cache.size().await; assert_eq!(size, 0); } #[tokio::test] async fn test_cache_stats() { let cache = MemoryCache::default_config(); let entry = MemoryEntry::new( "test-agent", MemoryType::Preference, "style", "test".to_string(), ); cache.put(entry.clone()).await; // Hit cache.get(&entry.uri).await; // Miss cache.get("nonexistent").await; let stats = cache.stats().await; assert_eq!(stats.hits, 1); assert_eq!(stats.misses, 1); let hit_rate = cache.hit_rate().await; assert!((hit_rate - 0.5).abs() < 0.001); } #[tokio::test] async fn test_cache_eviction() { let config = CacheConfig { max_entries: 2, ttl: Duration::from_secs(3600), enabled: true, }; let cache = MemoryCache::new(config); let entry1 = MemoryEntry::new("test", MemoryType::Preference, "1", "1".to_string()); let entry2 = MemoryEntry::new("test", MemoryType::Preference, "2", "2".to_string()); let entry3 = MemoryEntry::new("test", MemoryType::Preference, "3", "3".to_string()); cache.put(entry1.clone()).await; cache.put(entry2.clone()).await; // Access entry1 to make it hot cache.get(&entry1.uri).await; // Add entry3, should evict entry2 (LRU) cache.put(entry3).await; let size = cache.size().await; assert_eq!(size, 2); let stats = cache.stats().await; assert_eq!(stats.evictions, 1); } #[tokio::test] async fn test_get_hot_entries() { let cache = MemoryCache::default_config(); let entry1 = MemoryEntry::new("test", MemoryType::Preference, "1", "1".to_string()); let entry2 = MemoryEntry::new("test", MemoryType::Preference, "2", "2".to_string()); cache.put(entry1.clone()).await; cache.put(entry2.clone()).await; // Access entry1 multiple times cache.get(&entry1.uri).await; cache.get(&entry1.uri).await; let hot = cache.get_hot_entries(10).await; assert_eq!(hot.len(), 2); // entry1 should be first (more accesses) assert_eq!(hot[0].uri, entry1.uri); } }