heartbit_core/memory/
pruning.rs

1//! Memory entry pruning — removes weak or stale entries to keep the store compact.
2
3use std::sync::Arc;
4
5use crate::auth::TenantScope;
6use crate::error::Error;
7
8use super::Memory;
9
10/// Prune memory entries whose strength has decayed below a threshold
11/// and are older than a minimum age.
12///
13/// Returns the number of entries removed.
14///
15/// This is a convenience wrapper around `Memory::prune()` with default
16/// parameters suitable for periodic maintenance.
17pub async fn prune_weak_entries(
18    memory: &Arc<dyn Memory>,
19    scope: &TenantScope,
20    min_strength: f64,
21    min_age: chrono::Duration,
22) -> Result<usize, Error> {
23    memory.prune(scope, min_strength, min_age, None).await
24}
25
26/// Default minimum strength below which entries are prunable.
27pub const DEFAULT_MIN_STRENGTH: f64 = 0.1;
28
29/// Default minimum age before an entry can be pruned (24 hours).
30pub fn default_min_age() -> chrono::Duration {
31    chrono::Duration::hours(24)
32}
33
34#[cfg(test)]
35mod tests {
36    use super::*;
37    use crate::auth::TenantScope;
38    use crate::memory::in_memory::InMemoryStore;
39    use crate::memory::{Confidentiality, MemoryEntry, MemoryQuery, MemoryType};
40    use chrono::Utc;
41
42    fn make_entry(id: &str, strength: f64, hours_ago: i64) -> MemoryEntry {
43        let now = Utc::now();
44        MemoryEntry {
45            id: id.into(),
46            agent: "test".into(),
47            content: format!("content {id}"),
48            category: "fact".into(),
49            tags: vec![],
50            created_at: now - chrono::Duration::hours(hours_ago),
51            last_accessed: now,
52            access_count: 0,
53            importance: 5,
54            memory_type: MemoryType::default(),
55            keywords: vec![],
56            summary: None,
57            strength,
58            related_ids: vec![],
59            source_ids: vec![],
60            embedding: None,
61            confidentiality: Confidentiality::default(),
62            author_user_id: None,
63            author_tenant_id: None,
64        }
65    }
66
67    #[tokio::test]
68    async fn prune_removes_below_threshold() {
69        let store: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
70        let scope = TenantScope::default();
71        store
72            .store(&scope, make_entry("m1", 0.05, 48))
73            .await
74            .unwrap(); // weak + old
75        store
76            .store(&scope, make_entry("m2", 0.8, 48))
77            .await
78            .unwrap(); // strong + old
79        store
80            .store(&scope, make_entry("m3", 0.05, 0))
81            .await
82            .unwrap(); // weak + recent
83
84        let removed = prune_weak_entries(&store, &scope, 0.1, chrono::Duration::hours(24))
85            .await
86            .unwrap();
87        assert_eq!(removed, 1, "only m1 should be pruned (weak + old)");
88
89        let remaining = store
90            .recall(
91                &scope,
92                MemoryQuery {
93                    limit: 0,
94                    ..Default::default()
95                },
96            )
97            .await
98            .unwrap();
99        assert_eq!(remaining.len(), 2);
100        let ids: Vec<&str> = remaining.iter().map(|e| e.id.as_str()).collect();
101        assert!(ids.contains(&"m2"));
102        assert!(ids.contains(&"m3"));
103    }
104
105    #[tokio::test]
106    async fn prune_preserves_strong_entries() {
107        let store: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
108        let scope = TenantScope::default();
109        store
110            .store(&scope, make_entry("m1", 0.9, 100))
111            .await
112            .unwrap();
113        store
114            .store(&scope, make_entry("m2", 0.5, 100))
115            .await
116            .unwrap();
117
118        let removed = prune_weak_entries(&store, &scope, 0.1, chrono::Duration::hours(24))
119            .await
120            .unwrap();
121        assert_eq!(removed, 0);
122    }
123
124    #[tokio::test]
125    async fn prune_respects_min_age() {
126        let store: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
127        let scope = TenantScope::default();
128        store
129            .store(&scope, make_entry("m1", 0.01, 1))
130            .await
131            .unwrap(); // weak but only 1h old
132
133        let removed = prune_weak_entries(&store, &scope, 0.1, chrono::Duration::hours(24))
134            .await
135            .unwrap();
136        assert_eq!(removed, 0, "entry too recent to prune");
137    }
138
139    #[tokio::test]
140    async fn prune_empty_store() {
141        let store: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
142        let scope = TenantScope::default();
143        let removed = prune_weak_entries(&store, &scope, 0.1, chrono::Duration::hours(24))
144            .await
145            .unwrap();
146        assert_eq!(removed, 0);
147    }
148}
heartbit_core/memory/pruning.rs

heartbit_core/memory/
pruning.rs
