heartbit_core/memory/

mod.rs

//! Agent memory system — `Memory` trait, in-memory and PostgreSQL stores, BM25 and vector recall, Ebbinghaus decay, reflection, and consolidation.

#![allow(missing_docs)]
pub mod bm25;
pub mod consolidation;
pub mod embedding;
pub mod hybrid;
pub mod in_memory;
pub mod namespaced;
pub mod pruning;
pub mod reflection;
pub mod scoring;
pub mod shared_tools;
pub mod tools;

use std::future::Future;
use std::pin::Pin;

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

use crate::auth::TenantScope;
use crate::error::Error;

/// Classification of a memory entry's origin and purpose.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MemoryType {
    /// Direct experience or observation from an agent run.
    #[default]
    Episodic,
    /// Generalized knowledge derived from consolidation or reflection.
    Semantic,
    /// Higher-order insight generated by reflecting on episodic memories.
    Reflection,
}

/// Access classification for memory entries.
///
/// Ordered from least to most sensitive — `PartialOrd`/`Ord` derives use
/// variant declaration order, so `Public < Internal < Confidential < Restricted`.
#[derive(
    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default,
)]
#[serde(rename_all = "snake_case")]
pub enum Confidentiality {
    /// Shareable with anyone (public facts, general knowledge).
    #[default]
    Public,
    /// Internal context (work items, project details). Shareable with Verified+ senders.
    Internal,
    /// Personal/sensitive (expenses, health, private conversations). Owner only.
    Confidential,
    /// Secrets (API keys, passwords, tokens). Never included in LLM context.
    Restricted,
}

/// A single memory entry stored by an agent.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryEntry {
    pub id: String,
    pub agent: String,
    pub content: String,
    pub category: String,
    pub tags: Vec<String>,
    pub created_at: DateTime<Utc>,
    pub last_accessed: DateTime<Utc>,
    pub access_count: u32,
    /// Importance score (1-10). Default: 5. Set by agent at store time.
    #[serde(default = "default_importance")]
    pub importance: u8,
    /// Classification of memory origin (episodic, semantic, reflection).
    #[serde(default)]
    pub memory_type: MemoryType,
    /// LLM-generated keywords for improved retrieval.
    #[serde(default)]
    pub keywords: Vec<String>,
    /// One-sentence summary providing context for the memory content.
    #[serde(default)]
    pub summary: Option<String>,
    /// Ebbinghaus strength score. Starts at 1.0, decays over time,
    /// reinforced on access. Entries with low strength may be pruned.
    ///
    /// `Memory::store` preserves whatever value the caller supplies — no
    /// normalisation or clamping at insert time. `Memory::recall` reinforces
    /// the stored value by `+0.2` per access (capped at 1.0) unless the
    /// caller opts out via [`MemoryQuery::reinforce`] = `false`. Decay is
    /// applied lazily at read time via `effective_strength`.
    #[serde(default = "default_strength")]
    pub strength: f64,
    /// Bidirectional links to related memory entries.
    #[serde(default)]
    pub related_ids: Vec<String>,
    /// IDs of source entries that were consolidated into this one.
    #[serde(default)]
    pub source_ids: Vec<String>,
    /// Optional vector embedding for semantic search (hybrid retrieval).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub embedding: Option<Vec<f32>>,
    /// Access classification — controls which trust levels may read this entry.
    #[serde(default)]
    pub confidentiality: Confidentiality,
    /// User ID of the agent/user who authored this entry (multi-tenant authorship).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub author_user_id: Option<String>,
    /// Tenant ID of the agent/user who authored this entry (multi-tenant authorship).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub author_tenant_id: Option<String>,
}

pub(crate) fn default_importance() -> u8 {
    5
}

pub(crate) fn default_strength() -> f64 {
    1.0
}

pub(crate) fn default_category() -> String {
    "fact".into()
}

pub(crate) fn default_recall_limit() -> usize {
    10
}

/// Query parameters for recalling memories.
///
/// `limit` controls the maximum number of results returned. A value of `0`
/// means no limit (return all matching entries). This is the default.
#[derive(Debug, Clone)]
pub struct MemoryQuery {
    pub text: Option<String>,
    pub category: Option<String>,
    pub tags: Vec<String>,
    pub agent: Option<String>,
    /// Filter entries whose `agent` field starts with this prefix.
    /// Useful for cross-agent recall within a user namespace (e.g. `"tg:123"`
    /// matches `"tg:123:assistant"`, `"tg:123:researcher"`, etc.).
    /// Mutually exclusive with `agent` — if both are set, `agent` takes precedence.
    pub agent_prefix: Option<String>,
    /// Maximum number of results. `0` means unlimited.
    pub limit: usize,
    /// Filter by memory type.
    pub memory_type: Option<MemoryType>,
    /// Minimum strength threshold. Entries below this are excluded.
    pub min_strength: Option<f64>,
    /// Optional query embedding for hybrid (BM25 + vector) retrieval.
    /// When present and entries have stored embeddings, cosine similarity
    /// is computed and fused with BM25 via Reciprocal Rank Fusion.
    /// Populated automatically by `EmbeddingMemory::recall()`.
    pub query_embedding: Option<Vec<f32>>,
    /// When set, recall excludes entries with confidentiality above this level.
    /// `None` means no restriction (all levels returned).
    pub max_confidentiality: Option<Confidentiality>,
    /// Whether to reinforce the `strength` of returned entries on this read
    /// (Ebbinghaus reinforcement, +0.2 per access, capped at 1.0).
    ///
    /// Defaults to `true` to preserve historical recall semantics. Set to
    /// `false` for a pure read — useful when surfacing strength to a UI,
    /// driving deterministic decay tests, or letting `prune_weak_entries`
    /// observe a freshly-stored low-strength entry without first promoting
    /// it above the prune threshold. `last_accessed` and `access_count`
    /// are still updated regardless.
    pub reinforce: bool,

    /// Opt in to **exact-word** text matching instead of the default
    /// substring (`word.contains(token)`) semantics.
    ///
    /// When `true`, `InMemoryStore::recall` short-circuits to entries
    /// whose lowercased content / keyword tokens **exactly** equal at
    /// least one query token, looked up via the in-memory inverted
    /// index built at store time. Estimated gain at N=10k entries:
    /// 12.69 ms → 1–3 ms text-query recall (Phase 8 in
    /// `tasks/perf-audit-v2-2026-05-07.md`).
    ///
    /// Trade-off: queries whose tokens are *prefixes / substrings*
    /// of indexed words ("perf" matching "performance") will no
    /// longer match. Default is `false` — substring semantics
    /// preserved; opt-in when callers know their queries are full
    /// words. The Postgres path ignores this flag (it doesn't
    /// implement substring matching the same way).
    pub exact_words: bool,
}

impl Default for MemoryQuery {
    fn default() -> Self {
        Self {
            text: None,
            category: None,
            tags: Vec::new(),
            agent: None,
            agent_prefix: None,
            limit: 0,
            memory_type: None,
            min_strength: None,
            query_embedding: None,
            max_confidentiality: None,
            reinforce: true,
            exact_words: false,
        }
    }
}

/// Trait for persistent memory stores.
///
/// Every method requires a `&TenantScope` as the first parameter so the
/// compiler rejects code that accidentally drops tenant context. Single-tenant
/// deployments pass `&TenantScope::default()` (the empty-string sentinel).
///
/// Uses `Pin<Box<dyn Future>>` for dyn-compatibility, matching the `Tool` trait pattern.
///
/// # Example
///
/// Recalling memory entries with the in-memory backend:
///
/// ```rust,no_run
/// use heartbit_core::auth::TenantScope;
/// use heartbit_core::{InMemoryStore, Memory, MemoryQuery};
///
/// # async fn run() -> Result<(), heartbit_core::Error> {
/// let store = InMemoryStore::new();
/// let scope = TenantScope::default();
/// let hits = store
///     .recall(&scope, MemoryQuery {
///         agent: Some("assistant".into()),
///         text: Some("preferences".into()),
///         ..MemoryQuery::default()
///     })
///     .await?;
/// for entry in hits {
///     println!("{}: {}", entry.id, entry.content);
/// }
/// # Ok(()) }
/// ```
pub trait Memory: Send + Sync {
    /// Persist `entry` under `scope`.
    ///
    /// The caller-supplied [`MemoryEntry::strength`] is preserved verbatim;
    /// implementations must not normalise or clamp it at insert time.
    /// Reinforcement happens on read via [`Memory::recall`], not on write.
    fn store(
        &self,
        scope: &TenantScope,
        entry: MemoryEntry,
    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>>;

    /// Recall entries matching `query` from `scope`.
    ///
    /// By default, recall reinforces the `strength` of returned entries
    /// (Ebbinghaus reinforcement, +0.2 per access, capped at 1.0). To
    /// observe strength without modifying it, set
    /// [`MemoryQuery::reinforce`] to `false`.
    fn recall(
        &self,
        scope: &TenantScope,
        query: MemoryQuery,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryEntry>, Error>> + Send + '_>>;

    fn update(
        &self,
        scope: &TenantScope,
        id: &str,
        content: String,
    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>>;

    fn forget(
        &self,
        scope: &TenantScope,
        id: &str,
    ) -> Pin<Box<dyn Future<Output = Result<bool, Error>> + Send + '_>>;

    /// Add a bidirectional link between two memory entries.
    /// Default implementation is a no-op for backward compatibility.
    fn add_link(
        &self,
        _scope: &TenantScope,
        _id: &str,
        _related_id: &str,
    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
        Box::pin(async { Ok(()) })
    }

    /// Remove entries whose strength has decayed below `min_strength`
    /// and are older than `min_age`.
    ///
    /// When `agent_prefix` is `Some`, only entries whose `agent` field
    /// starts with the given prefix are candidates for pruning. This
    /// enables namespace-scoped pruning in multi-tenant setups where
    /// `NamespacedMemory` must not delete entries from other namespaces.
    ///
    /// Returns the number of entries pruned.
    /// Default implementation is a no-op for backward compatibility.
    fn prune(
        &self,
        _scope: &TenantScope,
        _min_strength: f64,
        _min_age: chrono::Duration,
        _agent_prefix: Option<&str>,
    ) -> Pin<Box<dyn Future<Output = Result<usize, Error>> + Send + '_>> {
        Box::pin(async { Ok(0) })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_entry(id: &str, content: &str) -> MemoryEntry {
        MemoryEntry {
            id: id.into(),
            agent: "a".into(),
            content: content.into(),
            category: "fact".into(),
            tags: vec![],
            created_at: Utc::now(),
            last_accessed: Utc::now(),
            access_count: 0,
            importance: 5,
            memory_type: MemoryType::default(),
            keywords: vec![],
            summary: None,
            strength: 1.0,
            related_ids: vec![],
            source_ids: vec![],
            embedding: None,
            confidentiality: Confidentiality::default(),
            author_user_id: None,
            author_tenant_id: None,
        }
    }

    #[test]
    fn memory_entry_serializes() {
        let entry = MemoryEntry {
            id: "m1".into(),
            agent: "researcher".into(),
            content: "Rust is fast".into(),
            category: "fact".into(),
            tags: vec!["rust".into()],
            created_at: Utc::now(),
            last_accessed: Utc::now(),
            access_count: 0,
            importance: 7,
            memory_type: MemoryType::default(),
            keywords: vec![],
            summary: None,
            strength: 1.0,
            related_ids: vec![],
            source_ids: vec![],
            embedding: None,
            confidentiality: Confidentiality::default(),
            author_user_id: None,
            author_tenant_id: None,
        };
        let json = serde_json::to_string(&entry).unwrap();
        let parsed: MemoryEntry = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.id, "m1");
        assert_eq!(parsed.agent, "researcher");
        assert_eq!(parsed.content, "Rust is fast");
        assert_eq!(parsed.importance, 7);
    }

    #[test]
    fn memory_entry_serializes_new_fields() {
        let entry = MemoryEntry {
            id: "m1".into(),
            agent: "a".into(),
            content: "test".into(),
            category: "fact".into(),
            tags: vec![],
            created_at: Utc::now(),
            last_accessed: Utc::now(),
            access_count: 0,
            importance: 7,
            memory_type: MemoryType::Reflection,
            keywords: vec!["rust".into(), "performance".into()],
            summary: Some("Rust is fast for systems programming".into()),
            strength: 0.85,
            related_ids: vec!["m2".into(), "m3".into()],
            source_ids: vec!["m0".into()],
            embedding: None,
            confidentiality: Confidentiality::default(),
            author_user_id: None,
            author_tenant_id: None,
        };
        let json = serde_json::to_string(&entry).unwrap();
        let parsed: MemoryEntry = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.memory_type, MemoryType::Reflection);
        assert_eq!(parsed.keywords, vec!["rust", "performance"]);
        assert_eq!(
            parsed.summary.as_deref(),
            Some("Rust is fast for systems programming")
        );
        assert!((parsed.strength - 0.85).abs() < f64::EPSILON);
        assert_eq!(parsed.related_ids, vec!["m2", "m3"]);
        assert_eq!(parsed.source_ids, vec!["m0"]);
    }

    #[test]
    fn memory_entry_deserialize_without_new_fields() {
        // Existing JSON without the new fields — backward compat
        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}"#;
        let entry: MemoryEntry = serde_json::from_str(json).unwrap();
        assert_eq!(entry.importance, 9);
        assert_eq!(entry.memory_type, MemoryType::Episodic);
        assert!(entry.keywords.is_empty());
        assert!(entry.summary.is_none());
        assert!((entry.strength - 1.0).abs() < f64::EPSILON);
        assert!(entry.related_ids.is_empty());
        assert!(entry.source_ids.is_empty());
    }

    #[test]
    fn memory_type_default_is_episodic() {
        assert_eq!(MemoryType::default(), MemoryType::Episodic);
    }

    #[test]
    fn strength_default_is_one() {
        assert!((default_strength() - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn memory_type_serialization_roundtrip() {
        for mt in [
            MemoryType::Episodic,
            MemoryType::Semantic,
            MemoryType::Reflection,
        ] {
            let json = serde_json::to_string(&mt).unwrap();
            let parsed: MemoryType = serde_json::from_str(&json).unwrap();
            assert_eq!(parsed, mt);
        }
    }

    #[test]
    fn memory_type_serializes_as_snake_case() {
        assert_eq!(
            serde_json::to_string(&MemoryType::Episodic).unwrap(),
            "\"episodic\""
        );
        assert_eq!(
            serde_json::to_string(&MemoryType::Semantic).unwrap(),
            "\"semantic\""
        );
        assert_eq!(
            serde_json::to_string(&MemoryType::Reflection).unwrap(),
            "\"reflection\""
        );
    }

    #[test]
    fn memory_entry_default_importance() {
        let entry = make_entry("m1", "test");
        assert_eq!(entry.importance, 5);
    }

    #[test]
    fn memory_entry_deserialize_without_importance() {
        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}"#;
        let entry: MemoryEntry = serde_json::from_str(json).unwrap();
        assert_eq!(entry.importance, 5); // default
    }

    #[test]
    fn memory_entry_deserialize_with_importance() {
        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}"#;
        let entry: MemoryEntry = serde_json::from_str(json).unwrap();
        assert_eq!(entry.importance, 9);
    }

    #[test]
    fn memory_query_default() {
        let q = MemoryQuery::default();
        assert!(q.text.is_none());
        assert!(q.category.is_none());
        assert!(q.tags.is_empty());
        assert!(q.agent.is_none());
        assert_eq!(q.limit, 0);
        assert!(q.memory_type.is_none());
        assert!(q.min_strength.is_none());
        assert!(q.query_embedding.is_none());
    }

    #[test]
    fn memory_trait_is_object_safe() {
        // Verify Memory can be used as dyn trait (including new default methods)
        fn _accepts_dyn(_m: &dyn Memory) {}
    }

    #[test]
    fn memory_entry_embedding_serde_roundtrip() {
        let entry = MemoryEntry {
            id: "m1".into(),
            agent: "a".into(),
            content: "test".into(),
            category: "fact".into(),
            tags: vec![],
            created_at: Utc::now(),
            last_accessed: Utc::now(),
            access_count: 0,
            importance: 5,
            memory_type: MemoryType::default(),
            keywords: vec![],
            summary: None,
            strength: 1.0,
            related_ids: vec![],
            source_ids: vec![],
            embedding: Some(vec![0.1, 0.2, 0.3]),
            confidentiality: Confidentiality::default(),
            author_user_id: None,
            author_tenant_id: None,
        };
        let json = serde_json::to_string(&entry).unwrap();
        assert!(json.contains("\"embedding\""));
        let parsed: MemoryEntry = serde_json::from_str(&json).unwrap();
        let emb = parsed.embedding.unwrap();
        assert_eq!(emb.len(), 3);
        assert!((emb[0] - 0.1).abs() < f32::EPSILON);
    }

    #[test]
    fn memory_entry_backward_compat_no_embedding() {
        // Old JSON without embedding field — should deserialize to None
        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}"#;
        let entry: MemoryEntry = serde_json::from_str(json).unwrap();
        assert!(entry.embedding.is_none());
    }

    #[test]
    fn memory_entry_none_embedding_not_serialized() {
        // When embedding is None, field should be omitted from JSON
        let entry = make_entry("m1", "test");
        let json = serde_json::to_string(&entry).unwrap();
        assert!(!json.contains("embedding"));
    }

    #[test]
    fn confidentiality_default_is_public() {
        assert_eq!(Confidentiality::default(), Confidentiality::Public);
    }

    #[test]
    fn confidentiality_ordering() {
        assert!(Confidentiality::Public < Confidentiality::Internal);
        assert!(Confidentiality::Internal < Confidentiality::Confidential);
        assert!(Confidentiality::Confidential < Confidentiality::Restricted);
    }

    #[test]
    fn confidentiality_serde_roundtrip() {
        for c in [
            Confidentiality::Public,
            Confidentiality::Internal,
            Confidentiality::Confidential,
            Confidentiality::Restricted,
        ] {
            let json = serde_json::to_string(&c).unwrap();
            let parsed: Confidentiality = serde_json::from_str(&json).unwrap();
            assert_eq!(parsed, c);
        }
    }

    #[test]
    fn confidentiality_serializes_as_snake_case() {
        assert_eq!(
            serde_json::to_string(&Confidentiality::Public).unwrap(),
            "\"public\""
        );
        assert_eq!(
            serde_json::to_string(&Confidentiality::Confidential).unwrap(),
            "\"confidential\""
        );
        assert_eq!(
            serde_json::to_string(&Confidentiality::Restricted).unwrap(),
            "\"restricted\""
        );
    }

    #[test]
    fn memory_entry_backward_compat_no_confidentiality() {
        // Old JSON without confidentiality field — should deserialize as Public
        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}"#;
        let entry: MemoryEntry = serde_json::from_str(json).unwrap();
        assert_eq!(entry.confidentiality, Confidentiality::Public);
    }

    #[test]
    fn memory_query_max_confidentiality_default_is_none() {
        let q = MemoryQuery::default();
        assert!(q.max_confidentiality.is_none());
    }
}