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");
}
}