velesdb_core/agent/

memory.rs

//! `AgentMemory` - Unified memory interface for AI agents (EPIC-010)
//!
//! Provides three memory subsystems for AI agents:
//! - **`SemanticMemory`**: Long-term knowledge facts with vector similarity search
//! - **`EpisodicMemory`**: Event timeline with temporal and similarity queries
//! - **`ProceduralMemory`**: Learned patterns with confidence scoring
//!
//! # Enhanced Features
//!
//! - **TTL/Eviction**: Automatic expiration and memory consolidation
//! - **Snapshots**: Versioned state persistence and rollback
//! - **Temporal Index**: Efficient O(log N) time-based queries
//! - **Adaptive Reinforcement**: Extensible confidence update strategies

// Reason: Numeric casts in agent memory are intentional:
// - u64 <-> i64 casts for timestamps (SystemTime uses u64, DB schema uses i64)
// - Values are always positive (elapsed time) and bounded by reasonable ranges
// - Casts verified by temporal index tests and snapshot functionality
#![allow(clippy::cast_possible_wrap)]
#![allow(clippy::cast_sign_loss)]

use crate::Database;
use std::sync::Arc;

pub use super::episodic_memory::EpisodicMemory;
pub use super::error::AgentMemoryError;
pub use super::procedural_memory::{ProceduralMemory, ProcedureMatch};
pub use super::semantic_memory::SemanticMemory;
pub use super::snapshot::{MemoryState, SnapshotManager};
pub use super::temporal_index::TemporalIndex;
pub use super::ttl::{EvictionConfig, ExpireResult, MemoryKind, MemoryTtl};

/// Default embedding dimension for memory collections.
pub const DEFAULT_DIMENSION: usize = 384;

/// Unified memory interface for AI agents.
///
/// Provides access to three memory subsystems:
/// - `semantic`: Long-term knowledge (vector storage; graph linkage planned)
/// - `episodic`: Event timeline with temporal context
/// - `procedural`: Learned patterns and action sequences
///
/// # Enhanced Features
///
/// - TTL management for automatic expiration
/// - Snapshot/restore for state persistence
/// - Temporal indexing for efficient time-based queries
/// - Configurable eviction policies
pub struct AgentMemory {
    db: Arc<Database>,
    semantic: SemanticMemory,
    episodic: EpisodicMemory,
    procedural: ProceduralMemory,
    ttl: Arc<MemoryTtl>,
    eviction_config: EvictionConfig,
    snapshot_manager: Option<SnapshotManager>,
}

impl AgentMemory {
    /// Creates a new `AgentMemory` instance from a `Database`.
    ///
    /// Initializes or connects to the three memory subsystem collections:
    /// - `_semantic_memory`: For knowledge facts
    /// - `_episodic_memory`: For event timeline
    /// - `_procedural_memory`: For learned patterns
    ///
    /// Uses the default embedding dimension (384).
    ///
    /// # Errors
    ///
    /// Returns an error when one of the underlying memory subsystems cannot be initialized.
    pub fn new(db: Arc<Database>) -> Result<Self, AgentMemoryError> {
        Self::with_dimension(db, DEFAULT_DIMENSION)
    }

    /// Creates a new `AgentMemory` with a custom embedding dimension.
    ///
    /// # Errors
    ///
    /// Returns an error when one of the underlying memory subsystems cannot be initialized.
    pub fn with_dimension(db: Arc<Database>, dimension: usize) -> Result<Self, AgentMemoryError> {
        let ttl = Arc::new(MemoryTtl::new());

        let semantic = SemanticMemory::new(Arc::clone(&db), dimension, Arc::clone(&ttl))?;
        let episodic = EpisodicMemory::new(
            Arc::clone(&db),
            dimension,
            Arc::clone(&ttl),
            Arc::new(TemporalIndex::new()),
        )?;
        let procedural = ProceduralMemory::new(Arc::clone(&db), dimension, Arc::clone(&ttl))?;

        Ok(Self {
            db,
            semantic,
            episodic,
            procedural,
            ttl,
            eviction_config: EvictionConfig::default(),
            snapshot_manager: None,
        })
    }

    /// Configures eviction policies for automatic memory cleanup.
    #[must_use]
    pub fn with_eviction_config(mut self, config: EvictionConfig) -> Self {
        self.eviction_config = config;
        self
    }

    /// Enables snapshot management with a storage directory.
    ///
    /// # Arguments
    ///
    /// * `snapshot_dir` - Directory path for storing snapshots
    /// * `max_snapshots` - Maximum number of snapshots to retain
    #[must_use]
    pub fn with_snapshots(mut self, snapshot_dir: &str, max_snapshots: usize) -> Self {
        self.snapshot_manager = Some(SnapshotManager::new(snapshot_dir, max_snapshots));
        self
    }

    /// Returns a reference to the semantic memory subsystem.
    #[must_use]
    pub fn semantic(&self) -> &SemanticMemory {
        &self.semantic
    }

    /// Returns a reference to the episodic memory subsystem.
    #[must_use]
    pub fn episodic(&self) -> &EpisodicMemory {
        &self.episodic
    }

    /// Returns a reference to the procedural memory subsystem.
    #[must_use]
    pub fn procedural(&self) -> &ProceduralMemory {
        &self.procedural
    }

    /// Sets TTL for a semantic memory entry (in-memory only; lost on restart).
    ///
    /// Use [`Self::set_semantic_ttl_durable`] to persist the expiry.
    pub fn set_semantic_ttl(&self, id: u64, ttl_seconds: u64) {
        self.ttl.set_ttl(MemoryKind::Semantic, id, ttl_seconds);
    }

    /// Sets TTL for an episodic memory entry (in-memory only; lost on restart).
    ///
    /// Use [`Self::set_episodic_ttl_durable`] to persist the expiry.
    pub fn set_episodic_ttl(&self, id: u64, ttl_seconds: u64) {
        self.ttl.set_ttl(MemoryKind::Episodic, id, ttl_seconds);
    }

    /// Sets TTL for a procedural memory entry (in-memory only; lost on restart).
    ///
    /// Use [`Self::set_procedural_ttl_durable`] to persist the expiry.
    pub fn set_procedural_ttl(&self, id: u64, ttl_seconds: u64) {
        self.ttl.set_ttl(MemoryKind::Procedural, id, ttl_seconds);
    }

    /// Durably sets the TTL of an existing semantic fact: the expiry is
    /// persisted to the reserved `_veles_expires_at` payload field and
    /// survives a restart.
    ///
    /// # Errors
    ///
    /// Returns `NotFound` when no fact with `id` exists, or
    /// `CollectionError` when persistence fails.
    pub fn set_semantic_ttl_durable(
        &self,
        id: u64,
        ttl_seconds: u64,
    ) -> Result<(), AgentMemoryError> {
        self.semantic.set_ttl_durable(id, ttl_seconds)
    }

    /// Durably sets the TTL of an existing episodic event: the expiry is
    /// persisted to the reserved `_veles_expires_at` payload field and
    /// survives a restart.
    ///
    /// # Errors
    ///
    /// Returns `NotFound` when no event with `id` exists, or
    /// `CollectionError` when persistence fails.
    pub fn set_episodic_ttl_durable(
        &self,
        id: u64,
        ttl_seconds: u64,
    ) -> Result<(), AgentMemoryError> {
        self.episodic.set_ttl_durable(id, ttl_seconds)
    }

    /// Durably sets the TTL of an existing procedure: the expiry is
    /// persisted to the reserved `_veles_expires_at` payload field and
    /// survives a restart.
    ///
    /// # Errors
    ///
    /// Returns `NotFound` when no procedure with `id` exists, or
    /// `CollectionError` when persistence fails.
    pub fn set_procedural_ttl_durable(
        &self,
        id: u64,
        ttl_seconds: u64,
    ) -> Result<(), AgentMemoryError> {
        self.procedural.set_ttl_durable(id, ttl_seconds)
    }

    /// Performs automatic expiration of entries that have exceeded their TTL.
    ///
    /// This method should be called periodically to clean up expired entries.
    /// It also handles consolidation of old episodic memories to semantic memory
    /// based on the configured eviction policy.
    ///
    /// # Returns
    ///
    /// Statistics about the expiration operation.
    ///
    /// # Errors
    ///
    /// Returns an error when consolidation operations fail.
    pub fn auto_expire(&self) -> Result<ExpireResult, AgentMemoryError> {
        // Read expired keys WITHOUT dropping their TTL entries yet. The entry is
        // removed (via the subsystem `delete`, which calls `ttl.remove`) only
        // after the point is actually deleted, so a failed delete leaves the TTL
        // entry intact and the id is retried on the next `auto_expire`. This
        // preserves the expiry invariant: a tracked-expired id is never forgotten
        // while its point still exists.
        //
        // Each key carries its owning `MemoryKind`, so an expired semantic id is
        // only ever deleted from semantic memory — it can never clobber a live
        // row that happens to share the numeric id in another subsystem.
        let expired_keys = self.ttl.get_expired();
        let mut result = ExpireResult::default();

        for &(kind, id) in &expired_keys {
            self.expire_one(kind, id, &mut result)?;
        }

        if self.eviction_config.consolidation_age_threshold > 0 {
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map_or(0, |d| d.as_secs() as i64);
            let cutoff = now - self.eviction_config.consolidation_age_threshold as i64;
            let outcome = self.consolidate_old_episodes(cutoff)?;
            result.episodic_consolidated = outcome.consolidated;
            result.consolidation_truncated = outcome.truncated;
        }

        result.procedural_evicted =
            self.evict_low_confidence_procedures(self.eviction_config.min_confidence_threshold)?;

        Ok(result)
    }

    /// Deletes a single expired entry from the subsystem that owns it and
    /// increments the matching counter only on a real deletion.
    fn expire_one(
        &self,
        kind: MemoryKind,
        id: u64,
        result: &mut ExpireResult,
    ) -> Result<(), AgentMemoryError> {
        match kind {
            MemoryKind::Semantic => {
                self.semantic.delete(id)?;
                result.semantic_expired += 1;
            }
            MemoryKind::Episodic => {
                self.episodic.delete(id)?;
                result.episodic_expired += 1;
            }
            MemoryKind::Procedural => {
                self.procedural.delete(id)?;
                result.procedural_expired += 1;
            }
        }
        Ok(())
    }

    /// Evicts procedures with confidence below the threshold.
    ///
    /// # Arguments
    ///
    /// * `min_confidence` - Minimum confidence threshold (0.0 - 1.0)
    ///
    /// # Returns
    ///
    /// Number of procedures evicted.
    ///
    /// # Errors
    ///
    /// Returns an error when listing or deleting procedures fails.
    pub fn evict_low_confidence_procedures(
        &self,
        min_confidence: f32,
    ) -> Result<usize, AgentMemoryError> {
        let all_procedures = self.procedural.list_all()?;
        let mut evicted = 0;

        for proc in all_procedures {
            if proc.confidence < min_confidence {
                self.procedural.delete(proc.id)?;
                evicted += 1;
            }
        }

        Ok(evicted)
    }

    /// Returns the snapshot manager, or an error if not configured.
    ///
    /// RF-DEDUP: Eliminates the repeated `ok_or_else(|| SnapshotError(...))` pattern
    /// across `snapshot`, `load_latest_snapshot`, `load_snapshot_version`, and
    /// `list_snapshot_versions`.
    fn require_snapshot_manager(&self) -> Result<&SnapshotManager, AgentMemoryError> {
        self.snapshot_manager.as_ref().ok_or_else(|| {
            AgentMemoryError::SnapshotError("Snapshot manager not configured".to_string())
        })
    }

    /// Creates a snapshot of the current memory state.
    ///
    /// # Returns
    ///
    /// The version number of the created snapshot.
    ///
    /// # Errors
    ///
    /// Returns an error when snapshot manager is not configured or snapshot persistence fails.
    pub fn snapshot(&self) -> Result<u64, AgentMemoryError> {
        let manager = self.require_snapshot_manager()?;

        let state = MemoryState {
            semantic: self.semantic.serialize()?,
            episodic: self.episodic.serialize()?,
            procedural: self.procedural.serialize()?,
            ttl: self.ttl.serialize(),
        };

        Ok(manager.create_versioned_snapshot(&state)?)
    }

    /// Loads the most recent snapshot.
    ///
    /// # Returns
    ///
    /// The version number of the loaded snapshot.
    ///
    /// # Errors
    ///
    /// Returns an error when snapshot manager is not configured, loading fails,
    /// or state restoration fails.
    pub fn load_latest_snapshot(&self) -> Result<u64, AgentMemoryError> {
        let manager = self.require_snapshot_manager()?;

        let (version, state) = manager.load_latest()?;
        self.restore_state(&state)?;
        Ok(version)
    }

    /// Loads a specific snapshot version.
    ///
    /// # Errors
    ///
    /// Returns an error when snapshot manager is not configured, loading fails,
    /// or state restoration fails.
    pub fn load_snapshot_version(&self, version: u64) -> Result<(), AgentMemoryError> {
        let manager = self.require_snapshot_manager()?;

        let state = manager.load_version(version)?;
        self.restore_state(&state)?;
        Ok(())
    }

    /// Lists all available snapshot versions.
    ///
    /// # Errors
    ///
    /// Returns an error when snapshot manager is not configured or listing fails.
    pub fn list_snapshot_versions(&self) -> Result<Vec<u64>, AgentMemoryError> {
        let manager = self.require_snapshot_manager()?;
        Ok(manager.list_versions()?)
    }

    /// Executes a `VelesQL` query against the semantic memory collection.
    ///
    /// Delegates to `Collection::execute_query_str` on the `_semantic_memory`
    /// collection. Use standard `VelesQL` syntax including `WHERE vector NEAR $v`,
    /// payload filters, `ORDER BY`, and `WITH` options. TTL-expired entries are
    /// filtered from the results, matching the native query APIs.
    ///
    /// # Errors
    ///
    /// Returns an error if the collection is missing or the query fails.
    pub fn query_semantic(
        &self,
        sql: &str,
        params: &std::collections::HashMap<String, serde_json::Value>,
    ) -> Result<Vec<crate::SearchResult>, AgentMemoryError> {
        super::memory_helpers::execute_velesql(
            &self.db,
            self.semantic.collection_name(),
            sql,
            params,
            &self.ttl,
            MemoryKind::Semantic,
        )
    }

    /// Executes a `VelesQL` query against the episodic memory collection.
    ///
    /// Delegates to `Collection::execute_query_str` on the `_episodic_memory`
    /// collection. Supports payload field filters like `WHERE timestamp > N`,
    /// `ORDER BY timestamp DESC`, and similarity search via `NEAR`. TTL-expired
    /// entries are filtered from the results, matching the native query APIs.
    ///
    /// # Errors
    ///
    /// Returns an error if the collection is missing or the query fails.
    pub fn query_episodic(
        &self,
        sql: &str,
        params: &std::collections::HashMap<String, serde_json::Value>,
    ) -> Result<Vec<crate::SearchResult>, AgentMemoryError> {
        super::memory_helpers::execute_velesql(
            &self.db,
            self.episodic.collection_name(),
            sql,
            params,
            &self.ttl,
            MemoryKind::Episodic,
        )
    }

    /// Executes a `VelesQL` query against the procedural memory collection.
    ///
    /// Delegates to `Collection::execute_query_str` on the `_procedural_memory`
    /// collection. Supports payload field filters like `WHERE confidence > 0.7`,
    /// `ORDER BY confidence DESC`, and scan queries. TTL-expired entries are
    /// filtered from the results, matching the native query APIs.
    ///
    /// # Errors
    ///
    /// Returns an error if the collection is missing or the query fails.
    pub fn query_procedural(
        &self,
        sql: &str,
        params: &std::collections::HashMap<String, serde_json::Value>,
    ) -> Result<Vec<crate::SearchResult>, AgentMemoryError> {
        super::memory_helpers::execute_velesql(
            &self.db,
            self.procedural.collection_name(),
            sql,
            params,
            &self.ttl,
            MemoryKind::Procedural,
        )
    }

    fn restore_state(&self, state: &MemoryState) -> Result<(), AgentMemoryError> {
        self.semantic.deserialize(&state.semantic)?;
        self.episodic.deserialize(&state.episodic)?;
        self.procedural.deserialize(&state.procedural)?;

        // An empty TTL section is legitimately empty state; anything else that
        // fails to deserialize is a corrupt/incompatible snapshot and must not
        // silently drop every TTL (which would make expiring entries immortal).
        if state.ttl.is_empty() {
            self.ttl.clear();
        } else if let Some(ttl) = MemoryTtl::deserialize(&state.ttl) {
            self.ttl.replace_from(&ttl);
        } else {
            return Err(AgentMemoryError::SnapshotError(
                "TTL state failed to deserialize (corrupt or incompatible snapshot)".to_string(),
            ));
        }

        Ok(())
    }

    /// Migrates episodic events older than `cutoff_timestamp` into semantic
    /// memory, capped at `eviction_config.max_entries_per_cycle` per call.
    ///
    /// # Preconditions
    ///
    /// The episodic id is *not* reused verbatim as the semantic id: semantic
    /// stores are upserts, so [`SemanticMemory::store_unique`] relocates the
    /// fact to a fresh semantic id on collision. This guarantees consolidation
    /// never clobbers an existing semantic fact even when the two subsystems
    /// share a numeric id.
    fn consolidate_old_episodes(
        &self,
        cutoff_timestamp: i64,
    ) -> Result<ConsolidationOutcome, AgentMemoryError> {
        let cap = self.eviction_config.max_entries_per_cycle;
        // Fetch one past the cap so we can tell whether more old events remain
        // than this cycle will process (the "truncated" signal).
        let old_events = self
            .episodic
            .older_than(cutoff_timestamp, cap.saturating_add(1))?;
        let truncated = old_events.len() > cap;
        let mut consolidated = 0;

        for (id, _description, _timestamp) in old_events.into_iter().take(cap) {
            if let Some((description, _ts, embedding)) = self.episodic.get_with_embedding(id)? {
                self.semantic.store_unique(id, &description, &embedding)?;
                self.episodic.delete(id)?;
                consolidated += 1;
            }
        }

        Ok(ConsolidationOutcome {
            consolidated,
            truncated,
        })
    }
}

/// Outcome of one consolidation pass.
struct ConsolidationOutcome {
    /// Number of episodes migrated to semantic memory this cycle.
    consolidated: usize,
    /// `true` when more old episodes remained than the per-cycle cap allowed.
    truncated: bool,
}