xz_memory_engine/layered/
fact.rs

1//! [`FactMemory`] implementation for [`LayeredMemory`].
2
3use async_trait::async_trait;
4use chrono::Utc;
5use uuid::Uuid;
6use xz_memory_core::{Entry, EntryStore, IndexSearcher, SearchOptions, StoreError};
7
8use super::traits::FactMemory;
9
10impl<S: EntryStore, I: IndexSearcher> super::default::LayeredMemory<S, I> {
11    fn facts_partition() -> String {
12        "facts".to_string()
13    }
14}
15
16fn pack_body(tags: &[String], fact_text: &str) -> String {
17    let tags_json = serde_json::to_string(tags).unwrap_or_else(|_| "[]".to_string());
18    format!("{}\n{}", tags_json, fact_text)
19}
20
21fn unpack_body(body: &str) -> (String, Vec<String>) {
22    if let Some((tags_json, fact_text)) = body.split_once('\n') {
23        let tags: Vec<String> = serde_json::from_str(tags_json).unwrap_or_default();
24        (fact_text.to_string(), tags)
25    } else {
26        (body.to_string(), vec![])
27    }
28}
29
30#[async_trait]
31impl<S: EntryStore, I: IndexSearcher> FactMemory for super::default::LayeredMemory<S, I> {
32    async fn remember(&self, fact: &str, tags: &[String]) -> Result<String, StoreError> {
33        let id = Uuid::new_v4().to_string();
34        let entry = Entry {
35            id: id.clone(),
36            partition: Self::facts_partition(),
37            body: pack_body(tags, fact),
38            recorded_at: Utc::now().timestamp_millis() as u64,
39        };
40        self.store.append(entry).await?;
41        Ok(id)
42    }
43
44    async fn recall(
45        &self,
46        query: &str,
47        limit: usize,
48    ) -> Result<Vec<(String, f32, Vec<String>)>, StoreError> {
49        let opts = SearchOptions { limit, min_relevance: None };
50        let partition = Self::facts_partition();
51        let results = self.searcher.search(&[partition], query, &opts).await?;
52        Ok(results
53            .into_iter()
54            .map(|se| {
55                let (text, tags) = unpack_body(&se.entry.body);
56                (text, se.relevance, tags)
57            })
58            .collect())
59    }
60
61    async fn forget(&self, id: &str) -> Result<(), StoreError> {
62        let _ = self.store.delete(id).await;
63        Ok(())
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use std::sync::Arc;
71    use xz_memory_core::ScoredEntry;
72
73    use crate::backends::InMemoryEntryStore;
74
75    struct MockSearcher {
76        fact: String,
77    }
78
79    #[async_trait]
80    impl IndexSearcher for MockSearcher {
81        async fn search(
82            &self,
83            _partitions: &[String],
84            _query: &str,
85            _opts: &SearchOptions,
86        ) -> Result<Vec<ScoredEntry>, StoreError> {
87            Ok(vec![ScoredEntry {
88                entry: Entry {
89                    id: "fact-1".into(),
90                    partition: "facts".into(),
91                    body: self.fact.clone(),
92                    recorded_at: 1000,
93                },
94                relevance: 0.95,
95            }])
96        }
97    }
98
99    fn setup() -> Arc<super::super::default::LayeredMemory<InMemoryEntryStore, MockSearcher>> {
100        let fact = pack_body(&["geography".to_string()], "Tokyo is the capital of Japan");
101        Arc::new(super::super::default::LayeredMemory::new(
102            Arc::new(InMemoryEntryStore::new()),
103            Arc::new(MockSearcher { fact }),
104        ))
105    }
106
107    #[tokio::test]
108    async fn test_remember_and_recall() {
109        let memory = setup();
110
111        let id =
112            memory.remember("Tokyo is the capital of Japan", &["geography".into()]).await.unwrap();
113        assert!(!id.is_empty());
114
115        let results = memory.recall("capital", 5).await.unwrap();
116        assert_eq!(results.len(), 1);
117        assert_eq!(results[0].0, "Tokyo is the capital of Japan");
118        assert!((results[0].1 - 0.95).abs() < 0.01);
119        assert_eq!(results[0].2, vec!["geography"]);
120    }
121
122    #[tokio::test]
123    async fn test_remember_without_tags() {
124        let memory = setup();
125
126        let id = memory.remember("Some fact", &[]).await.unwrap();
127        assert!(!id.is_empty());
128    }
129
130    #[tokio::test]
131    async fn test_forget() {
132        let store = Arc::new(InMemoryEntryStore::new());
133        let searcher = Arc::new(MockSearcher { fact: pack_body(&[], "test") });
134        let memory = super::super::default::LayeredMemory::new(store.clone(), searcher);
135
136        let id = memory.remember("test fact", &[]).await.unwrap();
137        memory.forget(&id).await.unwrap();
138
139        let remaining = store
140            .query(
141                "facts",
142                &xz_memory_core::TimeRange { start: None, end: None },
143                &xz_memory_core::QueryOptions {
144                    limit: usize::MAX,
145                    sort: xz_memory_core::SortOrder::Ascending,
146                },
147            )
148            .await
149            .unwrap();
150        assert!(remaining.iter().all(|e| e.id != id));
151    }
152
153    #[tokio::test]
154    async fn test_forget_nonexistent_is_noop() {
155        let memory = setup();
156        let result = memory.forget("nonexistent-id").await;
157        assert!(result.is_ok());
158    }
159
160    #[tokio::test]
161    async fn test_pack_unpack_roundtrip() {
162        let tags = vec!["t1".to_string(), "t2".to_string()];
163        let fact = "Hello world";
164        let body = pack_body(&tags, fact);
165        let (text, parsed_tags) = unpack_body(&body);
166        assert_eq!(text, fact);
167        assert_eq!(parsed_tags, tags);
168    }
169
170    #[tokio::test]
171    async fn test_unpack_no_tags() {
172        let (text, tags) = unpack_body("plain fact with no newlines");
173        assert_eq!(text, "plain fact with no newlines");
174        assert!(tags.is_empty());
175    }
176}
xz_memory_engine/layered/fact.rs

xz_memory_engine/layered/
fact.rs
