xz_memory_engine/layered/

summary.rs

//! [`SummaryMemory`] implementation for [`LayeredMemory`].

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::SummaryMemory;

impl<S: EntryStore, I: IndexSearcher> super::default::LayeredMemory<S, I> {
    fn summary_partition(scope: &str) -> String {
        format!("summary:{}", scope)
    }
}

fn pack_summary(source: &str, summary: &str) -> String {
    format!("{}\n{}", source, summary)
}

fn unpack_summary(body: &str) -> (String, String) {
    if let Some((source, summary)) = body.split_once('\n') {
        (source.to_string(), summary.to_string())
    } else {
        (String::new(), body.to_string())
    }
}

#[async_trait]
impl<S: EntryStore, I: IndexSearcher> SummaryMemory for super::default::LayeredMemory<S, I> {
    async fn get_latest(&self, scope: &str) -> Result<Option<(String, String)>, StoreError> {
        let partition = Self::summary_partition(scope);
        let opts = QueryOptions { limit: 1, sort: SortOrder::Descending };
        let range = TimeRange { start: None, end: None };
        let mut entries = self.store.query(&partition, &range, &opts).await?;
        Ok(entries.pop().map(|e| {
            let (source, summary) = unpack_summary(&e.body);
            (summary, source)
        }))
    }

    async fn store(&self, scope: &str, summary: &str, source: &str) -> Result<(), StoreError> {
        let entry = Entry {
            id: Uuid::new_v4().to_string(),
            partition: Self::summary_partition(scope),
            body: pack_summary(source, summary),
            recorded_at: Utc::now().timestamp_millis() as u64,
        };
        self.store.append(entry).await
    }

    async fn history(
        &self,
        scope: &str,
        limit: usize,
    ) -> Result<Vec<(String, String, u64)>, StoreError> {
        let partition = Self::summary_partition(scope);
        let opts = QueryOptions { limit, sort: SortOrder::Descending };
        let range = TimeRange { start: None, end: None };
        let entries = self.store.query(&partition, &range, &opts).await?;
        Ok(entries
            .into_iter()
            .map(|e| {
                let (source, summary) = unpack_summary(&e.body);
                (summary, source, e.recorded_at)
            })
            .collect())
    }
}

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

    use crate::backends::InMemoryEntryStore;

    struct MockSearcher;

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

    fn setup() -> Arc<super::super::default::LayeredMemory<InMemoryEntryStore, MockSearcher>> {
        Arc::new(super::super::default::LayeredMemory::new(
            Arc::new(InMemoryEntryStore::new()),
            Arc::new(MockSearcher),
        ))
    }

    #[tokio::test]
    async fn test_store_and_get_latest() {
        let memory = setup();

        memory.store("sess-1", "User asked about weather.", "raw-sess-1").await.unwrap();
        sleep(Duration::from_millis(2)).await;
        memory.store("sess-1", "User asked about travel.", "raw-sess-1").await.unwrap();

        let latest = memory.get_latest("sess-1").await.unwrap();
        assert!(latest.is_some());
        let (summary, source) = latest.unwrap();
        assert_eq!(summary, "User asked about travel.");
        assert_eq!(source, "raw-sess-1");
    }

    #[tokio::test]
    async fn test_get_latest_empty() {
        let memory = setup();
        let result = memory.get_latest("nonexistent").await.unwrap();
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn test_history() {
        let memory = setup();

        memory.store("sess-1", "First summary", "src-1").await.unwrap();
        sleep(Duration::from_millis(2)).await;
        memory.store("sess-1", "Second summary", "src-2").await.unwrap();
        sleep(Duration::from_millis(2)).await;
        memory.store("sess-1", "Third summary", "src-3").await.unwrap();

        let history = memory.history("sess-1", 10).await.unwrap();
        assert_eq!(history.len(), 3);
        assert_eq!(history[0].0, "Third summary");
        assert_eq!(history[1].0, "Second summary");
        assert_eq!(history[2].0, "First summary");
        assert_eq!(history[0].1, "src-3");
        assert_eq!(history[1].1, "src-2");
        assert_eq!(history[2].1, "src-1");
    }

    #[tokio::test]
    async fn test_history_with_limit() {
        let memory = setup();

        for i in 0..5 {
            memory.store("sess-1", &format!("Summary {}", i), "src").await.unwrap();
            sleep(Duration::from_millis(2)).await;
        }

        let history = memory.history("sess-1", 2).await.unwrap();
        assert_eq!(history.len(), 2);
        assert_eq!(history[0].0, "Summary 4");
        assert_eq!(history[1].0, "Summary 3");
    }

    #[tokio::test]
    async fn test_isolated_scopes() {
        let memory = setup();

        memory.store("scope-a", "Summary A", "src-a").await.unwrap();
        memory.store("scope-b", "Summary B", "src-b").await.unwrap();

        let latest_a = memory.get_latest("scope-a").await.unwrap().unwrap();
        let latest_b = memory.get_latest("scope-b").await.unwrap().unwrap();
        assert_eq!(latest_a.0, "Summary A");
        assert_eq!(latest_b.0, "Summary B");

        let history_a = memory.history("scope-a", 10).await.unwrap();
        assert_eq!(history_a.len(), 1);
        let history_b = memory.history("scope-b", 10).await.unwrap();
        assert_eq!(history_b.len(), 1);
    }

    #[tokio::test]
    async fn test_pack_unpack_roundtrip() {
        let body = pack_summary("src", "my summary");
        let (source, summary) = unpack_summary(&body);
        assert_eq!(source, "src");
        assert_eq!(summary, "my summary");
    }

    #[tokio::test]
    async fn test_unpack_no_newline() {
        let (source, summary) = unpack_summary("just a summary");
        assert_eq!(source, "");
        assert_eq!(summary, "just a summary");
    }
}