heartbit_core/memory/
mod.rs

1//! Agent memory system — `Memory` trait, in-memory and PostgreSQL stores, BM25 and vector recall, Ebbinghaus decay, reflection, and consolidation.
2
3pub mod bm25;
4pub mod consolidation;
5pub mod embedding;
6pub mod hybrid;
7pub mod in_memory;
8pub mod namespaced;
9pub mod pruning;
10pub mod reflection;
11pub mod scoring;
12pub mod shared_tools;
13pub mod tools;
14
15use std::future::Future;
16use std::pin::Pin;
17
18use chrono::{DateTime, Utc};
19use serde::{Deserialize, Serialize};
20
21use crate::auth::TenantScope;
22use crate::error::Error;
23
24/// Classification of a memory entry's origin and purpose.
25#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
26#[serde(rename_all = "snake_case")]
27pub enum MemoryType {
28    /// Direct experience or observation from an agent run.
29    #[default]
30    Episodic,
31    /// Generalized knowledge derived from consolidation or reflection.
32    Semantic,
33    /// Higher-order insight generated by reflecting on episodic memories.
34    Reflection,
35}
36
37/// Access classification for memory entries.
38///
39/// Ordered from least to most sensitive — `PartialOrd`/`Ord` derives use
40/// variant declaration order, so `Public < Internal < Confidential < Restricted`.
41#[derive(
42    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default,
43)]
44#[serde(rename_all = "snake_case")]
45pub enum Confidentiality {
46    /// Shareable with anyone (public facts, general knowledge).
47    #[default]
48    Public,
49    /// Internal context (work items, project details). Shareable with Verified+ senders.
50    Internal,
51    /// Personal/sensitive (expenses, health, private conversations). Owner only.
52    Confidential,
53    /// Secrets (API keys, passwords, tokens). Never included in LLM context.
54    Restricted,
55}
56
57/// A single memory entry stored by an agent.
58#[derive(Debug, Clone, Serialize, Deserialize)]
59pub struct MemoryEntry {
60    /// Unique entry ID (UUID).
61    pub id: String,
62    /// Owning agent / namespace.
63    pub agent: String,
64    /// Stored content (free text or serialized fact).
65    pub content: String,
66    /// User-supplied category (e.g. `fact`, `preference`).
67    pub category: String,
68    /// Free-form tags for filtering.
69    pub tags: Vec<String>,
70    /// Insertion timestamp.
71    pub created_at: DateTime<Utc>,
72    /// Timestamp of the most recent recall hit.
73    pub last_accessed: DateTime<Utc>,
74    /// Number of times this entry has been recalled.
75    pub access_count: u32,
76    /// Importance score (1-10). Default: 5. Set by agent at store time.
77    #[serde(default = "default_importance")]
78    pub importance: u8,
79    /// Classification of memory origin (episodic, semantic, reflection).
80    #[serde(default)]
81    pub memory_type: MemoryType,
82    /// LLM-generated keywords for improved retrieval.
83    #[serde(default)]
84    pub keywords: Vec<String>,
85    /// One-sentence summary providing context for the memory content.
86    #[serde(default)]
87    pub summary: Option<String>,
88    /// Ebbinghaus strength score. Starts at 1.0, decays over time,
89    /// reinforced on access. Entries with low strength may be pruned.
90    ///
91    /// `Memory::store` preserves whatever value the caller supplies — no
92    /// normalisation or clamping at insert time. `Memory::recall` reinforces
93    /// the stored value by `+0.2` per access (capped at 1.0) unless the
94    /// caller opts out via [`MemoryQuery::reinforce`] = `false`. Decay is
95    /// applied lazily at read time via `effective_strength`.
96    #[serde(default = "default_strength")]
97    pub strength: f64,
98    /// Bidirectional links to related memory entries.
99    #[serde(default)]
100    pub related_ids: Vec<String>,
101    /// IDs of source entries that were consolidated into this one.
102    #[serde(default)]
103    pub source_ids: Vec<String>,
104    /// Optional vector embedding for semantic search (hybrid retrieval).
105    #[serde(default, skip_serializing_if = "Option::is_none")]
106    pub embedding: Option<Vec<f32>>,
107    /// Access classification — controls which trust levels may read this entry.
108    #[serde(default)]
109    pub confidentiality: Confidentiality,
110    /// User ID of the agent/user who authored this entry (multi-tenant authorship).
111    #[serde(default, skip_serializing_if = "Option::is_none")]
112    pub author_user_id: Option<String>,
113    /// Tenant ID of the agent/user who authored this entry (multi-tenant authorship).
114    #[serde(default, skip_serializing_if = "Option::is_none")]
115    pub author_tenant_id: Option<String>,
116}
117
118pub(crate) fn default_importance() -> u8 {
119    5
120}
121
122pub(crate) fn default_strength() -> f64 {
123    1.0
124}
125
126pub(crate) fn default_category() -> String {
127    "fact".into()
128}
129
130pub(crate) fn default_recall_limit() -> usize {
131    10
132}
133
134/// Query parameters for recalling memories.
135///
136/// `limit` controls the maximum number of results returned. A value of `0`
137/// means no limit (return all matching entries). This is the default.
138#[derive(Debug, Clone)]
139pub struct MemoryQuery {
140    /// Optional full-text query (BM25 scored against `content` + `keywords`).
141    pub text: Option<String>,
142    /// Restrict to a single category.
143    pub category: Option<String>,
144    /// Restrict to entries that include all of these tags.
145    pub tags: Vec<String>,
146    /// Restrict to a single agent / namespace.
147    pub agent: Option<String>,
148    /// Filter entries whose `agent` field starts with this prefix.
149    /// Useful for cross-agent recall within a user namespace (e.g. `"tg:123"`
150    /// matches `"tg:123:assistant"`, `"tg:123:researcher"`, etc.).
151    /// Mutually exclusive with `agent` — if both are set, `agent` takes precedence.
152    pub agent_prefix: Option<String>,
153    /// Maximum number of results. `0` means unlimited.
154    pub limit: usize,
155    /// Filter by memory type.
156    pub memory_type: Option<MemoryType>,
157    /// Minimum strength threshold. Entries below this are excluded.
158    pub min_strength: Option<f64>,
159    /// Optional query embedding for hybrid (BM25 + vector) retrieval.
160    /// When present and entries have stored embeddings, cosine similarity
161    /// is computed and fused with BM25 via Reciprocal Rank Fusion.
162    /// Populated automatically by `EmbeddingMemory::recall()`.
163    pub query_embedding: Option<Vec<f32>>,
164    /// When set, recall excludes entries with confidentiality above this level.
165    /// `None` means no restriction (all levels returned).
166    pub max_confidentiality: Option<Confidentiality>,
167    /// Whether to reinforce the `strength` of returned entries on this read
168    /// (Ebbinghaus reinforcement, +0.2 per access, capped at 1.0).
169    ///
170    /// Defaults to `true` to preserve historical recall semantics. Set to
171    /// `false` for a pure read — useful when surfacing strength to a UI,
172    /// driving deterministic decay tests, or letting `prune_weak_entries`
173    /// observe a freshly-stored low-strength entry without first promoting
174    /// it above the prune threshold. `last_accessed` and `access_count`
175    /// are still updated regardless.
176    pub reinforce: bool,
177
178    /// Opt in to **exact-word** text matching instead of the default
179    /// substring (`word.contains(token)`) semantics.
180    ///
181    /// When `true`, `InMemoryStore::recall` short-circuits to entries
182    /// whose lowercased content / keyword tokens **exactly** equal at
183    /// least one query token, looked up via the in-memory inverted
184    /// index built at store time. Estimated gain at N=10k entries:
185    /// 12.69 ms → 1–3 ms text-query recall (Phase 8 in
186    /// `tasks/perf-audit-v2-2026-05-07.md`).
187    ///
188    /// Trade-off: queries whose tokens are *prefixes / substrings*
189    /// of indexed words ("perf" matching "performance") will no
190    /// longer match. Default is `false` — substring semantics
191    /// preserved; opt-in when callers know their queries are full
192    /// words. The Postgres path ignores this flag (it doesn't
193    /// implement substring matching the same way).
194    pub exact_words: bool,
195}
196
197impl Default for MemoryQuery {
198    fn default() -> Self {
199        Self {
200            text: None,
201            category: None,
202            tags: Vec::new(),
203            agent: None,
204            agent_prefix: None,
205            limit: 0,
206            memory_type: None,
207            min_strength: None,
208            query_embedding: None,
209            max_confidentiality: None,
210            reinforce: true,
211            exact_words: false,
212        }
213    }
214}
215
216/// Trait for persistent memory stores.
217///
218/// Every method requires a `&TenantScope` as the first parameter so the
219/// compiler rejects code that accidentally drops tenant context. Single-tenant
220/// deployments pass `&TenantScope::default()` (the empty-string sentinel).
221///
222/// Uses `Pin<Box<dyn Future>>` for dyn-compatibility, matching the `Tool` trait pattern.
223///
224/// # Example
225///
226/// Recalling memory entries with the in-memory backend:
227///
228/// ```rust,no_run
229/// use heartbit_core::auth::TenantScope;
230/// use heartbit_core::{InMemoryStore, Memory, MemoryQuery};
231///
232/// # async fn run() -> Result<(), heartbit_core::Error> {
233/// let store = InMemoryStore::new();
234/// let scope = TenantScope::default();
235/// let hits = store
236///     .recall(&scope, MemoryQuery {
237///         agent: Some("assistant".into()),
238///         text: Some("preferences".into()),
239///         ..MemoryQuery::default()
240///     })
241///     .await?;
242/// for entry in hits {
243///     println!("{}: {}", entry.id, entry.content);
244/// }
245/// # Ok(()) }
246/// ```
247pub trait Memory: Send + Sync {
248    /// Persist `entry` under `scope`.
249    ///
250    /// The caller-supplied [`MemoryEntry::strength`] is preserved verbatim;
251    /// implementations must not normalise or clamp it at insert time.
252    /// Reinforcement happens on read via [`Memory::recall`], not on write.
253    fn store(
254        &self,
255        scope: &TenantScope,
256        entry: MemoryEntry,
257    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>>;
258
259    /// Recall entries matching `query` from `scope`.
260    ///
261    /// By default, recall reinforces the `strength` of returned entries
262    /// (Ebbinghaus reinforcement, +0.2 per access, capped at 1.0). To
263    /// observe strength without modifying it, set
264    /// [`MemoryQuery::reinforce`] to `false`.
265    fn recall(
266        &self,
267        scope: &TenantScope,
268        query: MemoryQuery,
269    ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryEntry>, Error>> + Send + '_>>;
270
271    /// Replace `content` on the entry identified by `id`.
272    fn update(
273        &self,
274        scope: &TenantScope,
275        id: &str,
276        content: String,
277    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>>;
278
279    /// Delete the entry identified by `id`; returns `true` if an entry was removed.
280    fn forget(
281        &self,
282        scope: &TenantScope,
283        id: &str,
284    ) -> Pin<Box<dyn Future<Output = Result<bool, Error>> + Send + '_>>;
285
286    /// Add a bidirectional link between two memory entries.
287    /// Default implementation is a no-op for backward compatibility.
288    fn add_link(
289        &self,
290        _scope: &TenantScope,
291        _id: &str,
292        _related_id: &str,
293    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
294        Box::pin(async { Ok(()) })
295    }
296
297    /// Remove entries whose strength has decayed below `min_strength`
298    /// and are older than `min_age`.
299    ///
300    /// When `agent_prefix` is `Some`, only entries whose `agent` field
301    /// starts with the given prefix are candidates for pruning. This
302    /// enables namespace-scoped pruning in multi-tenant setups where
303    /// `NamespacedMemory` must not delete entries from other namespaces.
304    ///
305    /// Returns the number of entries pruned.
306    /// Default implementation is a no-op for backward compatibility.
307    fn prune(
308        &self,
309        _scope: &TenantScope,
310        _min_strength: f64,
311        _min_age: chrono::Duration,
312        _agent_prefix: Option<&str>,
313    ) -> Pin<Box<dyn Future<Output = Result<usize, Error>> + Send + '_>> {
314        Box::pin(async { Ok(0) })
315    }
316}
317
318#[cfg(test)]
319mod tests {
320    use super::*;
321
322    fn make_entry(id: &str, content: &str) -> MemoryEntry {
323        MemoryEntry {
324            id: id.into(),
325            agent: "a".into(),
326            content: content.into(),
327            category: "fact".into(),
328            tags: vec![],
329            created_at: Utc::now(),
330            last_accessed: Utc::now(),
331            access_count: 0,
332            importance: 5,
333            memory_type: MemoryType::default(),
334            keywords: vec![],
335            summary: None,
336            strength: 1.0,
337            related_ids: vec![],
338            source_ids: vec![],
339            embedding: None,
340            confidentiality: Confidentiality::default(),
341            author_user_id: None,
342            author_tenant_id: None,
343        }
344    }
345
346    #[test]
347    fn memory_entry_serializes() {
348        let entry = MemoryEntry {
349            id: "m1".into(),
350            agent: "researcher".into(),
351            content: "Rust is fast".into(),
352            category: "fact".into(),
353            tags: vec!["rust".into()],
354            created_at: Utc::now(),
355            last_accessed: Utc::now(),
356            access_count: 0,
357            importance: 7,
358            memory_type: MemoryType::default(),
359            keywords: vec![],
360            summary: None,
361            strength: 1.0,
362            related_ids: vec![],
363            source_ids: vec![],
364            embedding: None,
365            confidentiality: Confidentiality::default(),
366            author_user_id: None,
367            author_tenant_id: None,
368        };
369        let json = serde_json::to_string(&entry).unwrap();
370        let parsed: MemoryEntry = serde_json::from_str(&json).unwrap();
371        assert_eq!(parsed.id, "m1");
372        assert_eq!(parsed.agent, "researcher");
373        assert_eq!(parsed.content, "Rust is fast");
374        assert_eq!(parsed.importance, 7);
375    }
376
377    #[test]
378    fn memory_entry_serializes_new_fields() {
379        let entry = MemoryEntry {
380            id: "m1".into(),
381            agent: "a".into(),
382            content: "test".into(),
383            category: "fact".into(),
384            tags: vec![],
385            created_at: Utc::now(),
386            last_accessed: Utc::now(),
387            access_count: 0,
388            importance: 7,
389            memory_type: MemoryType::Reflection,
390            keywords: vec!["rust".into(), "performance".into()],
391            summary: Some("Rust is fast for systems programming".into()),
392            strength: 0.85,
393            related_ids: vec!["m2".into(), "m3".into()],
394            source_ids: vec!["m0".into()],
395            embedding: None,
396            confidentiality: Confidentiality::default(),
397            author_user_id: None,
398            author_tenant_id: None,
399        };
400        let json = serde_json::to_string(&entry).unwrap();
401        let parsed: MemoryEntry = serde_json::from_str(&json).unwrap();
402        assert_eq!(parsed.memory_type, MemoryType::Reflection);
403        assert_eq!(parsed.keywords, vec!["rust", "performance"]);
404        assert_eq!(
405            parsed.summary.as_deref(),
406            Some("Rust is fast for systems programming")
407        );
408        assert!((parsed.strength - 0.85).abs() < f64::EPSILON);
409        assert_eq!(parsed.related_ids, vec!["m2", "m3"]);
410        assert_eq!(parsed.source_ids, vec!["m0"]);
411    }
412
413    #[test]
414    fn memory_entry_deserialize_without_new_fields() {
415        // Existing JSON without the new fields — backward compat
416        let json = r#"{"id":"m1","agent":"a","content":"test","category":"fact","tags":[],"created_at":"2024-01-01T00:00:00Z","last_accessed":"2024-01-01T00:00:00Z","access_count":0,"importance":9}"#;
417        let entry: MemoryEntry = serde_json::from_str(json).unwrap();
418        assert_eq!(entry.importance, 9);
419        assert_eq!(entry.memory_type, MemoryType::Episodic);
420        assert!(entry.keywords.is_empty());
421        assert!(entry.summary.is_none());
422        assert!((entry.strength - 1.0).abs() < f64::EPSILON);
423        assert!(entry.related_ids.is_empty());
424        assert!(entry.source_ids.is_empty());
425    }
426
427    #[test]
428    fn memory_type_default_is_episodic() {
429        assert_eq!(MemoryType::default(), MemoryType::Episodic);
430    }
431
432    #[test]
433    fn strength_default_is_one() {
434        assert!((default_strength() - 1.0).abs() < f64::EPSILON);
435    }
436
437    #[test]
438    fn memory_type_serialization_roundtrip() {
439        for mt in [
440            MemoryType::Episodic,
441            MemoryType::Semantic,
442            MemoryType::Reflection,
443        ] {
444            let json = serde_json::to_string(&mt).unwrap();
445            let parsed: MemoryType = serde_json::from_str(&json).unwrap();
446            assert_eq!(parsed, mt);
447        }
448    }
449
450    #[test]
451    fn memory_type_serializes_as_snake_case() {
452        assert_eq!(
453            serde_json::to_string(&MemoryType::Episodic).unwrap(),
454            "\"episodic\""
455        );
456        assert_eq!(
457            serde_json::to_string(&MemoryType::Semantic).unwrap(),
458            "\"semantic\""
459        );
460        assert_eq!(
461            serde_json::to_string(&MemoryType::Reflection).unwrap(),
462            "\"reflection\""
463        );
464    }
465
466    #[test]
467    fn memory_entry_default_importance() {
468        let entry = make_entry("m1", "test");
469        assert_eq!(entry.importance, 5);
470    }
471
472    #[test]
473    fn memory_entry_deserialize_without_importance() {
474        let json = r#"{"id":"m1","agent":"a","content":"test","category":"fact","tags":[],"created_at":"2024-01-01T00:00:00Z","last_accessed":"2024-01-01T00:00:00Z","access_count":0}"#;
475        let entry: MemoryEntry = serde_json::from_str(json).unwrap();
476        assert_eq!(entry.importance, 5); // default
477    }
478
479    #[test]
480    fn memory_entry_deserialize_with_importance() {
481        let json = r#"{"id":"m1","agent":"a","content":"test","category":"fact","tags":[],"created_at":"2024-01-01T00:00:00Z","last_accessed":"2024-01-01T00:00:00Z","access_count":0,"importance":9}"#;
482        let entry: MemoryEntry = serde_json::from_str(json).unwrap();
483        assert_eq!(entry.importance, 9);
484    }
485
486    #[test]
487    fn memory_query_default() {
488        let q = MemoryQuery::default();
489        assert!(q.text.is_none());
490        assert!(q.category.is_none());
491        assert!(q.tags.is_empty());
492        assert!(q.agent.is_none());
493        assert_eq!(q.limit, 0);
494        assert!(q.memory_type.is_none());
495        assert!(q.min_strength.is_none());
496        assert!(q.query_embedding.is_none());
497    }
498
499    #[test]
500    fn memory_trait_is_object_safe() {
501        // Verify Memory can be used as dyn trait (including new default methods)
502        fn _accepts_dyn(_m: &dyn Memory) {}
503    }
504
505    #[test]
506    fn memory_entry_embedding_serde_roundtrip() {
507        let entry = MemoryEntry {
508            id: "m1".into(),
509            agent: "a".into(),
510            content: "test".into(),
511            category: "fact".into(),
512            tags: vec![],
513            created_at: Utc::now(),
514            last_accessed: Utc::now(),
515            access_count: 0,
516            importance: 5,
517            memory_type: MemoryType::default(),
518            keywords: vec![],
519            summary: None,
520            strength: 1.0,
521            related_ids: vec![],
522            source_ids: vec![],
523            embedding: Some(vec![0.1, 0.2, 0.3]),
524            confidentiality: Confidentiality::default(),
525            author_user_id: None,
526            author_tenant_id: None,
527        };
528        let json = serde_json::to_string(&entry).unwrap();
529        assert!(json.contains("\"embedding\""));
530        let parsed: MemoryEntry = serde_json::from_str(&json).unwrap();
531        let emb = parsed.embedding.unwrap();
532        assert_eq!(emb.len(), 3);
533        assert!((emb[0] - 0.1).abs() < f32::EPSILON);
534    }
535
536    #[test]
537    fn memory_entry_backward_compat_no_embedding() {
538        // Old JSON without embedding field — should deserialize to None
539        let json = r#"{"id":"m1","agent":"a","content":"test","category":"fact","tags":[],"created_at":"2024-01-01T00:00:00Z","last_accessed":"2024-01-01T00:00:00Z","access_count":0,"importance":5}"#;
540        let entry: MemoryEntry = serde_json::from_str(json).unwrap();
541        assert!(entry.embedding.is_none());
542    }
543
544    #[test]
545    fn memory_entry_none_embedding_not_serialized() {
546        // When embedding is None, field should be omitted from JSON
547        let entry = make_entry("m1", "test");
548        let json = serde_json::to_string(&entry).unwrap();
549        assert!(!json.contains("embedding"));
550    }
551
552    #[test]
553    fn confidentiality_default_is_public() {
554        assert_eq!(Confidentiality::default(), Confidentiality::Public);
555    }
556
557    #[test]
558    fn confidentiality_ordering() {
559        assert!(Confidentiality::Public < Confidentiality::Internal);
560        assert!(Confidentiality::Internal < Confidentiality::Confidential);
561        assert!(Confidentiality::Confidential < Confidentiality::Restricted);
562    }
563
564    #[test]
565    fn confidentiality_serde_roundtrip() {
566        for c in [
567            Confidentiality::Public,
568            Confidentiality::Internal,
569            Confidentiality::Confidential,
570            Confidentiality::Restricted,
571        ] {
572            let json = serde_json::to_string(&c).unwrap();
573            let parsed: Confidentiality = serde_json::from_str(&json).unwrap();
574            assert_eq!(parsed, c);
575        }
576    }
577
578    #[test]
579    fn confidentiality_serializes_as_snake_case() {
580        assert_eq!(
581            serde_json::to_string(&Confidentiality::Public).unwrap(),
582            "\"public\""
583        );
584        assert_eq!(
585            serde_json::to_string(&Confidentiality::Confidential).unwrap(),
586            "\"confidential\""
587        );
588        assert_eq!(
589            serde_json::to_string(&Confidentiality::Restricted).unwrap(),
590            "\"restricted\""
591        );
592    }
593
594    #[test]
595    fn memory_entry_backward_compat_no_confidentiality() {
596        // Old JSON without confidentiality field — should deserialize as Public
597        let json = r#"{"id":"m1","agent":"a","content":"test","category":"fact","tags":[],"created_at":"2024-01-01T00:00:00Z","last_accessed":"2024-01-01T00:00:00Z","access_count":0,"importance":5}"#;
598        let entry: MemoryEntry = serde_json::from_str(json).unwrap();
599        assert_eq!(entry.confidentiality, Confidentiality::Public);
600    }
601
602    #[test]
603    fn memory_query_max_confidentiality_default_is_none() {
604        let q = MemoryQuery::default();
605        assert!(q.max_confidentiality.is_none());
606    }
607}
heartbit_core/memory/mod.rs

heartbit_core/memory/
mod.rs
