enki_memory/

in_memory.rs

use async_trait::async_trait;
use enki_core::error::Result;
use enki_core::memory::{Memory, MemoryEntry, MemoryQuery};
use dashmap::DashMap;
use std::sync::Arc;
use tracing::trace;

/// In-memory storage backend with TTL support
pub struct InMemoryBackend {
    store: Arc<DashMap<String, MemoryEntry>>,
    max_entries: Option<usize>,
    default_ttl_seconds: Option<i64>,
}

impl InMemoryBackend {
    /// Create a new in-memory backend
    pub fn new() -> Self {
        Self {
            store: Arc::new(DashMap::new()),
            max_entries: None,
            default_ttl_seconds: None,
        }
    }

    /// Create with maximum entry limit
    pub fn with_max_entries(mut self, max: usize) -> Self {
        self.max_entries = Some(max);
        self
    }

    /// Create with default TTL
    pub fn with_ttl_seconds(mut self, seconds: i64) -> Self {
        self.default_ttl_seconds = Some(seconds);
        self
    }

    /// Check if entry is expired
    fn is_expired(&self, entry: &MemoryEntry) -> bool {
        entry.is_expired()
    }

    /// Remove all expired entries
    fn cleanup_expired(&self) {
        self.store.retain(|_, entry| !self.is_expired(entry));
    }

    /// Evict oldest entry to make room
    fn evict_oldest(&self) {
        if let Some(oldest) = self
            .store
            .iter()
            .min_by_key(|e| e.value().created_at)
            .map(|e| e.key().clone())
        {
            self.store.remove(&oldest);
        }
    }
}

impl Default for InMemoryBackend {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Memory for InMemoryBackend {
    async fn store(&self, mut entry: MemoryEntry) -> Result<String> {
        // Apply default TTL if not set
        if entry.expires_at.is_none() {
            if let Some(ttl) = self.default_ttl_seconds {
                entry = entry.with_ttl_seconds(ttl);
            }
        }

        // Cleanup expired entries first
        self.cleanup_expired();

        // Enforce max entries limit
        if let Some(max) = self.max_entries {
            while self.store.len() >= max {
                self.evict_oldest();
            }
        }

        let id = entry.id.clone();
        self.store.insert(id.clone(), entry);
        trace!(id = %id, "stored memory entry");
        Ok(id)
    }

    async fn get(&self, id: &str) -> Result<Option<MemoryEntry>> {
        if let Some(entry) = self.store.get(id) {
            if self.is_expired(&entry) {
                // Remove expired entry
                drop(entry); // Release read lock
                self.store.remove(id);
                trace!(id = %id, "memory miss (expired)");
                Ok(None)
            } else {
                trace!(id = %id, "memory hit");
                Ok(Some(entry.clone()))
            }
        } else {
            trace!(id = %id, "memory miss");
            Ok(None)
        }
    }

    async fn search(&self, query: MemoryQuery) -> Result<Vec<MemoryEntry>> {
        // Cleanup expired entries
        self.cleanup_expired();

        let mut results: Vec<MemoryEntry> = self
            .store
            .iter()
            .map(|e| e.value().clone())
            .filter(|entry| {
                // Apply metadata filters - all filters must match
                let matches_filters = query
                    .filters
                    .iter()
                    .all(|(key, value)| entry.metadata.get(key) == Some(value));

                // Apply semantic query (keyword search) if present
                // This performs keyword-based search by splitting the query into words
                // and matching if any significant keyword is found in the content
                let matches_query = if let Some(ref q) = query.semantic_query {
                    let content_lower = entry.content.to_lowercase();
                    let query_lower = q.to_lowercase();

                    // Stop words to filter out (common words that don't add meaning)
                    let stop_words: std::collections::HashSet<&str> = [
                        "a", "an", "the", "is", "are", "was", "were", "be", "been", "being",
                        "have", "has", "had", "do", "does", "did", "will", "would", "could",
                        "should", "may", "might", "can", "what", "which", "who", "whom", "this",
                        "that", "these", "those", "am", "of", "for", "with", "at", "by", "from",
                        "to", "in", "on", "how", "does", "it", "its", "and", "or", "but", "not",
                        "each", "other", "some", "any", "no", "all", "most", "more",
                    ]
                    .into_iter()
                    .collect();

                    // Extract keywords from query (words with 2+ chars, not stop words)
                    let keywords: Vec<&str> = query_lower
                        .split(|c: char| !c.is_alphanumeric())
                        .filter(|word| word.len() >= 2 && !stop_words.contains(word))
                        .collect();

                    // Match if any keyword is found in the content
                    keywords
                        .iter()
                        .any(|keyword| content_lower.contains(keyword))
                } else {
                    true
                };

                matches_filters && matches_query
            })
            .collect();

        // Sort by creation time (newest first) for consistent ordering
        results.sort_by(|a, b| b.created_at.cmp(&a.created_at));

        // Apply limit
        if let Some(limit) = query.limit {
            results.truncate(limit);
        }

        Ok(results)
    }

    async fn delete(&self, id: &str) -> Result<bool> {
        Ok(self.store.remove(id).is_some())
    }

    async fn clear(&self) -> Result<()> {
        self.store.clear();
        Ok(())
    }

    async fn count(&self) -> Result<usize> {
        self.cleanup_expired();
        Ok(self.store.len())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[tokio::test]
    async fn test_basic_storage() {
        let backend = InMemoryBackend::new();

        let entry = MemoryEntry::new("test content");
        let id = backend.store(entry.clone()).await.unwrap();

        let retrieved = backend.get(&id).await.unwrap();
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().content, "test content");
    }

    #[tokio::test]
    async fn test_ttl_expiration() {
        let backend = InMemoryBackend::new();

        // Create entry that expires in 1 second
        let entry = MemoryEntry::new("expires soon").with_ttl_seconds(1);
        let id = backend.store(entry).await.unwrap();

        // Should exist immediately
        assert!(backend.get(&id).await.unwrap().is_some());

        // Wait for expiration
        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;

        // Should be gone
        assert!(backend.get(&id).await.unwrap().is_none());
    }

    #[tokio::test]
    async fn test_metadata_search() {
        let backend = InMemoryBackend::new();

        let entry1 = MemoryEntry::new("user message")
            .with_metadata("type", json!("user"))
            .with_metadata("user_id", json!("123"));

        let entry2 = MemoryEntry::new("system message").with_metadata("type", json!("system"));

        backend.store(entry1).await.unwrap();
        backend.store(entry2).await.unwrap();

        // Search for user messages
        let query = MemoryQuery::new().with_filter("type", json!("user"));
        let results = backend.search(query).await.unwrap();

        assert_eq!(results.len(), 1);
        assert_eq!(results[0].content, "user message");
    }

    #[tokio::test]
    async fn test_max_entries() {
        let backend = InMemoryBackend::new().with_max_entries(2);

        backend.store(MemoryEntry::new("first")).await.unwrap();
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

        backend.store(MemoryEntry::new("second")).await.unwrap();
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

        backend.store(MemoryEntry::new("third")).await.unwrap();

        // Should only have 2 entries (oldest evicted)
        assert_eq!(backend.count().await.unwrap(), 2);

        // First entry should be evicted
        let results = backend.search(MemoryQuery::new()).await.unwrap();
        assert!(!results.iter().any(|e| e.content == "first"));
    }

    #[tokio::test]
    async fn test_delete() {
        let backend = InMemoryBackend::new();

        let entry = MemoryEntry::new("to delete");
        let id = backend.store(entry).await.unwrap();

        assert!(backend.delete(&id).await.unwrap());
        assert!(backend.get(&id).await.unwrap().is_none());
        assert!(!backend.delete(&id).await.unwrap()); // Already deleted
    }

    #[tokio::test]
    async fn test_clear() {
        let backend = InMemoryBackend::new();

        backend.store(MemoryEntry::new("entry 1")).await.unwrap();
        backend.store(MemoryEntry::new("entry 2")).await.unwrap();

        assert_eq!(backend.count().await.unwrap(), 2);

        backend.clear().await.unwrap();
        assert_eq!(backend.count().await.unwrap(), 0);
    }
}