xz-memory-engine 0.2.0

//! Default [`LayeredMemory`] implementation.
//!
//! Combines an [`EntryStore`] for append/query/evict operations with an
//! [`IndexSearcher`] for future search capabilities.  Implements the
//! [`ConversationMemory`] trait using a `conv:{session_id}` partition scheme.

use std::sync::Arc;

use async_trait::async_trait;
use chrono::Utc;
use uuid::Uuid;
use xz_memory_core::{
    Entry, EntryStore, IndexSearcher, QueryOptions, SortOrder, StoreError, TimeRange,
};

use super::traits::ConversationMemory;

/// Default layered memory backed by an entry store and index searcher.
pub struct LayeredMemory<S: EntryStore, I: IndexSearcher> {
    pub(crate) store: Arc<S>,
    pub(crate) searcher: Arc<I>,
}

impl<S: EntryStore, I: IndexSearcher> LayeredMemory<S, I> {
    /// Create a new `LayeredMemory`.
    pub fn new(store: Arc<S>, searcher: Arc<I>) -> Self {
        Self { store, searcher }
    }
}

#[async_trait]
impl<S: EntryStore, I: IndexSearcher> ConversationMemory for LayeredMemory<S, I> {
    async fn append(&self, session_id: &str, message: &str) -> Result<(), StoreError> {
        let partition = format!("conv:{}", session_id);
        let entry = Entry {
            id: Uuid::new_v4().to_string(),
            partition,
            body: message.to_string(),
            recorded_at: Utc::now().timestamp_millis() as u64,
        };
        self.store.append(entry).await
    }

    async fn recent(&self, session_id: &str, n: usize) -> Result<Vec<String>, StoreError> {
        let partition = format!("conv:{}", session_id);
        let limit = if n == 0 { usize::MAX } else { n };
        let opts = QueryOptions { limit, sort: SortOrder::Descending };
        let range = TimeRange { start: None, end: None };
        let mut entries = self.store.query(&partition, &range, &opts).await?;
        entries.reverse();
        Ok(entries.into_iter().map(|e| e.body).collect())
    }

    async fn evict(&self, session_id: &str, keep: usize) -> Result<usize, StoreError> {
        let partition = format!("conv:{}", session_id);
        self.store.evict(&partition, keep).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use tokio::time::{Duration, sleep};
    use xz_memory_core::{ScoredEntry, SearchOptions};

    use crate::backends::InMemoryEntryStore;

    struct MockSearcher;

    #[async_trait]
    impl IndexSearcher for MockSearcher {
        async fn search(
            &self,
            _partitions: &[String],
            _query: &str,
            _opts: &SearchOptions,
        ) -> Result<Vec<ScoredEntry>, StoreError> {
            Ok(vec![])
        }
    }

    #[tokio::test]
    async fn test_append_and_recent() {
        let store = Arc::new(InMemoryEntryStore::new());
        let searcher = Arc::new(MockSearcher);
        let memory = LayeredMemory::new(store, searcher);

        memory.append("sess-1", "Hello").await.unwrap();
        sleep(Duration::from_millis(2)).await;
        memory.append("sess-1", "World").await.unwrap();

        let messages = memory.recent("sess-1", 10).await.unwrap();
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[0], "Hello");
        assert_eq!(messages[1], "World");
    }

    #[tokio::test]
    async fn test_evict() {
        let store = Arc::new(InMemoryEntryStore::new());
        let searcher = Arc::new(MockSearcher);
        let memory = LayeredMemory::new(store, searcher);

        for msg in ["a", "b", "c", "d", "e"] {
            memory.append("sess-1", msg).await.unwrap();
            sleep(Duration::from_millis(2)).await;
        }

        let evicted = memory.evict("sess-1", 2).await.unwrap();
        assert_eq!(evicted, 3);

        let remaining = memory.recent("sess-1", 10).await.unwrap();
        assert_eq!(remaining.len(), 2);
        assert_eq!(remaining[0], "d");
        assert_eq!(remaining[1], "e");
    }

    #[tokio::test]
    async fn test_empty_recent() {
        let store = Arc::new(InMemoryEntryStore::new());
        let searcher = Arc::new(MockSearcher);
        let memory = LayeredMemory::new(store, searcher);

        let messages = memory.recent("nonexistent", 10).await.unwrap();
        assert!(messages.is_empty());
    }
}