xz-memory-engine 0.2.0

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

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

use super::traits::FactMemory;

impl<S: EntryStore, I: IndexSearcher> super::default::LayeredMemory<S, I> {
    fn facts_partition() -> String {
        "facts".to_string()
    }
}

fn pack_body(tags: &[String], fact_text: &str) -> String {
    let tags_json = serde_json::to_string(tags).unwrap_or_else(|_| "[]".to_string());
    format!("{}\n{}", tags_json, fact_text)
}

fn unpack_body(body: &str) -> (String, Vec<String>) {
    if let Some((tags_json, fact_text)) = body.split_once('\n') {
        let tags: Vec<String> = serde_json::from_str(tags_json).unwrap_or_default();
        (fact_text.to_string(), tags)
    } else {
        (body.to_string(), vec![])
    }
}

#[async_trait]
impl<S: EntryStore, I: IndexSearcher> FactMemory for super::default::LayeredMemory<S, I> {
    async fn remember(&self, fact: &str, tags: &[String]) -> Result<String, StoreError> {
        let id = Uuid::new_v4().to_string();
        let entry = Entry {
            id: id.clone(),
            partition: Self::facts_partition(),
            body: pack_body(tags, fact),
            recorded_at: Utc::now().timestamp_millis() as u64,
        };
        self.store.append(entry).await?;
        Ok(id)
    }

    async fn recall(
        &self,
        query: &str,
        limit: usize,
    ) -> Result<Vec<(String, f32, Vec<String>)>, StoreError> {
        let opts = SearchOptions { limit, min_relevance: None };
        let partition = Self::facts_partition();
        let results = self.searcher.search(&[partition], query, &opts).await?;
        Ok(results
            .into_iter()
            .map(|se| {
                let (text, tags) = unpack_body(&se.entry.body);
                (text, se.relevance, tags)
            })
            .collect())
    }

    async fn forget(&self, id: &str) -> Result<(), StoreError> {
        let _ = self.store.delete(id).await;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use xz_memory_core::ScoredEntry;

    use crate::backends::InMemoryEntryStore;

    struct MockSearcher {
        fact: String,
    }

    #[async_trait]
    impl IndexSearcher for MockSearcher {
        async fn search(
            &self,
            _partitions: &[String],
            _query: &str,
            _opts: &SearchOptions,
        ) -> Result<Vec<ScoredEntry>, StoreError> {
            Ok(vec![ScoredEntry {
                entry: Entry {
                    id: "fact-1".into(),
                    partition: "facts".into(),
                    body: self.fact.clone(),
                    recorded_at: 1000,
                },
                relevance: 0.95,
            }])
        }
    }

    fn setup() -> Arc<super::super::default::LayeredMemory<InMemoryEntryStore, MockSearcher>> {
        let fact = pack_body(&["geography".to_string()], "Tokyo is the capital of Japan");
        Arc::new(super::super::default::LayeredMemory::new(
            Arc::new(InMemoryEntryStore::new()),
            Arc::new(MockSearcher { fact }),
        ))
    }

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

        let id =
            memory.remember("Tokyo is the capital of Japan", &["geography".into()]).await.unwrap();
        assert!(!id.is_empty());

        let results = memory.recall("capital", 5).await.unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, "Tokyo is the capital of Japan");
        assert!((results[0].1 - 0.95).abs() < 0.01);
        assert_eq!(results[0].2, vec!["geography"]);
    }

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

        let id = memory.remember("Some fact", &[]).await.unwrap();
        assert!(!id.is_empty());
    }

    #[tokio::test]
    async fn test_forget() {
        let store = Arc::new(InMemoryEntryStore::new());
        let searcher = Arc::new(MockSearcher { fact: pack_body(&[], "test") });
        let memory = super::super::default::LayeredMemory::new(store.clone(), searcher);

        let id = memory.remember("test fact", &[]).await.unwrap();
        memory.forget(&id).await.unwrap();

        let remaining = store
            .query(
                "facts",
                &xz_memory_core::TimeRange { start: None, end: None },
                &xz_memory_core::QueryOptions {
                    limit: usize::MAX,
                    sort: xz_memory_core::SortOrder::Ascending,
                },
            )
            .await
            .unwrap();
        assert!(remaining.iter().all(|e| e.id != id));
    }

    #[tokio::test]
    async fn test_forget_nonexistent_is_noop() {
        let memory = setup();
        let result = memory.forget("nonexistent-id").await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_pack_unpack_roundtrip() {
        let tags = vec!["t1".to_string(), "t2".to_string()];
        let fact = "Hello world";
        let body = pack_body(&tags, fact);
        let (text, parsed_tags) = unpack_body(&body);
        assert_eq!(text, fact);
        assert_eq!(parsed_tags, tags);
    }

    #[tokio::test]
    async fn test_unpack_no_tags() {
        let (text, tags) = unpack_body("plain fact with no newlines");
        assert_eq!(text, "plain fact with no newlines");
        assert!(tags.is_empty());
    }
}