velesdb_core/agent/

procedural_memory.rs

//! Procedural Memory - Learned patterns storage (US-004)
//!
//! Stores action sequences and learned procedures with confidence scoring.
//! Supports pattern matching by similarity and reinforcement learning.
//! Includes extensible reinforcement strategies for adaptive confidence updates.

// Reason: Numeric casts in procedural memory are intentional:
// - u64->i64 casts for timestamps (SystemTime::elapsed returns u64, DB uses i64)
// - i64->u64 casts for display purposes (timestamps are always positive)
// - f64->f32 casts for confidence scores (f32 precision sufficient, values clamped to 0.0-1.0)
// - All timestamp values are bounded by reasonable time ranges
#![allow(clippy::cast_possible_wrap)]
#![allow(clippy::cast_sign_loss)]
#![allow(clippy::cast_precision_loss)]
#![allow(clippy::cast_possible_truncation)]

use crate::{Database, Point};
use parking_lot::RwLock;
use serde_json::json;
use std::collections::HashSet;
use std::sync::Arc;

use super::error::AgentMemoryError;
use super::memory_helpers;
use super::reinforcement::{
    power_law_decay, FixedRate, ReinforcementContext, ReinforcementStrategy,
};
use super::ttl::{MemoryKind, MemoryTtl};

struct ProcedureState {
    name: String,
    steps: Vec<String>,
    confidence: f32,
    usage_count: u64,
    created_at: i64,
    last_used_at: i64,
    success_count: u64,
    failure_count: u64,
}

impl ProcedureState {
    fn build_reinforcement_context(&self, now: i64) -> ReinforcementContext {
        let total_uses = self.success_count + self.failure_count;
        let success_rate = if total_uses > 0 {
            self.success_count as f32 / total_uses as f32
        } else {
            0.5
        };
        let mut custom = std::collections::HashMap::new();
        custom.insert("success_count".to_string(), self.success_count as f64);
        custom.insert("failure_count".to_string(), self.failure_count as f64);
        ReinforcementContext {
            usage_count: self.usage_count,
            created_at: self.created_at as u64,
            last_used: self.last_used_at as u64,
            current_time: now as u64,
            recent_success_rate: Some(success_rate),
            custom,
        }
    }
}

/// A procedure match result returned by [`ProceduralMemory::recall`].
#[derive(Debug, Clone)]
pub struct ProcedureMatch {
    /// Unique identifier for the procedure.
    pub id: u64,
    /// Human-readable name for the procedure.
    pub name: String,
    /// Ordered sequence of steps that constitute the procedure.
    pub steps: Vec<String>,
    /// Confidence score of this procedure (0.0 - 1.0).
    pub confidence: f32,
    /// Similarity score from the vector search.
    pub score: f32,
}

/// Procedural memory for storing learned action sequences with confidence scoring.
///
/// Stores procedures as embedding vectors with associated metadata.
/// Supports confidence-based recall, reinforcement learning, and TTL expiration.
///
/// ### ACT-R activation decay
///
/// When built with [`with_activation_decay`](Self::with_activation_decay), the
/// confidence returned by [`recall`](Self::recall) is modulated by the ACT-R
/// base-level power-law formula: `c × max(1, t_days)^(-d)` where `d` is the
/// decay exponent (Anderson 1996, default ≈ 0.5).  The stored value is
/// unchanged — decay is applied read-only at retrieval time.
pub struct ProceduralMemory {
    collection_name: String,
    db: Arc<Database>,
    dimension: usize,
    ttl: Arc<MemoryTtl>,
    reinforcement_strategy: Arc<dyn ReinforcementStrategy>,
    stored_ids: RwLock<HashSet<u64>>,
    /// ACT-R decay exponent `d` for passive confidence decay at recall time.
    /// `None` disables decay (default — backward-compatible behaviour).
    activation_decay_exponent: Option<f32>,
    /// Edge-id allocator for [`Self::relate`] (seeded past existing edges).
    next_edge_id: std::sync::atomic::AtomicU64,
}

impl ProceduralMemory {
    const COLLECTION_NAME: &'static str = "_procedural_memory";

    /// Returns the name of the underlying `VelesDB` collection.
    #[must_use]
    pub fn collection_name(&self) -> &str {
        &self.collection_name
    }

    /// Returns the embedding dimension for this collection.
    #[must_use]
    pub fn dimension(&self) -> usize {
        self.dimension
    }

    /// Creates or opens procedural memory.
    ///
    /// # Errors
    ///
    /// Returns an error when collection creation/opening fails or dimensions mismatch.
    pub fn new_from_db(db: Arc<Database>, dimension: usize) -> Result<Self, AgentMemoryError> {
        Self::new(db, dimension, Arc::new(MemoryTtl::new()))
    }

    pub(crate) fn new(
        db: Arc<Database>,
        dimension: usize,
        ttl: Arc<MemoryTtl>,
    ) -> Result<Self, AgentMemoryError> {
        let (collection_name, dimension, stored_ids) =
            memory_helpers::init_tracked_memory(&db, Self::COLLECTION_NAME, dimension)?;
        memory_helpers::rebuild_ttl_from_payloads(
            &db,
            &collection_name,
            &ttl,
            MemoryKind::Procedural,
        )?;

        let next_edge_id = memory_helpers::seed_edge_counter(&memory_helpers::get_collection(
            &db,
            &collection_name,
        )?);

        Ok(Self {
            collection_name,
            db,
            dimension,
            ttl,
            reinforcement_strategy: Arc::new(FixedRate::default()),
            stored_ids,
            activation_decay_exponent: None,
            next_edge_id,
        })
    }

    /// Overrides the default reinforcement strategy with a custom implementation.
    #[must_use]
    pub fn with_reinforcement_strategy(mut self, strategy: Arc<dyn ReinforcementStrategy>) -> Self {
        self.reinforcement_strategy = strategy;
        self
    }

    /// Enables ACT-R base-level activation decay at recall time.
    ///
    /// When set, the confidence returned by [`recall`](Self::recall) is
    /// multiplied by `max(1, t_days)^(-decay_exponent)` where `t_days` is the
    /// number of days since the procedure was last reinforced.
    ///
    /// The stored confidence is **not** modified — decay is applied read-only.
    /// ACT-R recommends `decay_exponent ≈ 0.5` (Anderson 1996).
    #[must_use]
    pub fn with_activation_decay(mut self, decay_exponent: f32) -> Self {
        self.activation_decay_exponent = Some(decay_exponent.clamp(0.0, 2.0));
        self
    }

    /// Learns a procedure and stores it in memory.
    ///
    /// # Errors
    ///
    /// Returns an error when embedding dimension is invalid, the collection is
    /// unavailable, or persistence fails.
    pub fn learn(
        &self,
        procedure_id: u64,
        name: &str,
        steps: &[String],
        embedding: Option<&[f32]>,
        confidence: f32,
    ) -> Result<(), AgentMemoryError> {
        self.learn_internal(procedure_id, name, steps, embedding, confidence, None)
    }

    /// Shared store path: persists the procedure, optionally with a durable
    /// `_veles_expires_at` payload field (epoch seconds) for TTL'd procedures.
    fn learn_internal(
        &self,
        procedure_id: u64,
        name: &str,
        steps: &[String],
        embedding: Option<&[f32]>,
        confidence: f32,
        expires_at: Option<u64>,
    ) -> Result<(), AgentMemoryError> {
        let vector = memory_helpers::resolve_embedding(self.dimension, embedding)?;
        let collection = memory_helpers::get_collection(&self.db, &self.collection_name)?;

        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_or(0, |d| d.as_secs() as i64);

        let mut payload = json!({
            "name": name,
            "steps": steps,
            "confidence": confidence,
            "usage_count": 0,
            "created_at": now,
            "last_used_at": now,
            "success_count": 0,
            "failure_count": 0
        });
        memory_helpers::attach_expiry(&mut payload, expires_at);
        let point = Point::new(procedure_id, vector, Some(payload));

        memory_helpers::upsert_points(&collection, vec![point])?;
        self.stored_ids.write().insert(procedure_id);
        Ok(())
    }

    /// Learns a procedure and assigns a TTL for auto-expiration.
    ///
    /// A `ttl_seconds` of `0` means "expire immediately": the procedure is
    /// eagerly removed (and any pre-existing point for `procedure_id` deleted),
    /// harmonising the behaviour with `SemanticMemory::store_with_ttl`. The
    /// embedding is still dimension-validated so callers get the same error
    /// contract as a real learn.
    ///
    /// The expiry is persisted as a reserved `_veles_expires_at` (epoch
    /// seconds) payload field, so the TTL survives a process restart: the
    /// in-memory map is rebuilt from payloads when the collection is reopened.
    ///
    /// # Errors
    ///
    /// Returns the same errors as [`Self::learn`].
    pub fn learn_with_ttl(
        &self,
        procedure_id: u64,
        name: &str,
        steps: &[String],
        embedding: Option<&[f32]>,
        confidence: f32,
        ttl_seconds: u64,
    ) -> Result<(), AgentMemoryError> {
        if ttl_seconds == 0 {
            if let Some(emb) = embedding {
                memory_helpers::validate_dimension(self.dimension, emb.len())?;
            }
            return self.delete(procedure_id);
        }
        let expires_at = MemoryTtl::now().saturating_add(ttl_seconds);
        self.learn_internal(
            procedure_id,
            name,
            steps,
            embedding,
            confidence,
            Some(expires_at),
        )?;
        self.ttl
            .set_expiry(MemoryKind::Procedural, procedure_id, expires_at);
        Ok(())
    }

    /// Durably sets (or refreshes) the TTL of an existing procedure.
    ///
    /// Unlike `AgentMemory::set_procedural_ttl` (in-memory map only, lost on
    /// restart), this persists the expiry to the reserved `_veles_expires_at`
    /// payload field, so it survives a restart. A `ttl_seconds` of 0 expires
    /// the procedure immediately.
    ///
    /// # Errors
    ///
    /// Returns `NotFound` when no procedure with `procedure_id` exists, or
    /// `CollectionError` when persistence fails.
    pub fn set_ttl_durable(
        &self,
        procedure_id: u64,
        ttl_seconds: u64,
    ) -> Result<(), AgentMemoryError> {
        memory_helpers::set_ttl_durable(
            &self.db,
            &self.collection_name,
            &self.ttl,
            MemoryKind::Procedural,
            procedure_id,
            ttl_seconds,
        )
    }

    /// Relates two live procedures with a typed, durable graph edge (e.g.
    /// `DEPENDS_ON`, `REFINES`); see `SemanticMemory::relate` for semantics.
    ///
    /// # Errors
    ///
    /// Returns `NotFound` when either endpoint is missing or expired, or
    /// `CollectionError` when the edge write fails.
    pub fn relate(
        &self,
        from_id: u64,
        to_id: u64,
        rel_type: &str,
        properties: Option<&serde_json::Map<String, serde_json::Value>>,
    ) -> Result<u64, AgentMemoryError> {
        memory_helpers::relate_memory_points(
            &memory_helpers::MemorySubsystem {
                db: &self.db,
                collection_name: &self.collection_name,
                ttl: &self.ttl,
                kind: MemoryKind::Procedural,
                next_edge_id: &self.next_edge_id,
            },
            from_id,
            to_id,
            rel_type,
            properties,
        )
    }

    /// Returns the outgoing relations of a procedure.
    ///
    /// # Errors
    ///
    /// Returns `CollectionError` when the collection cannot be resolved.
    pub fn relations(
        &self,
        id: u64,
    ) -> Result<Vec<crate::collection::graph::GraphEdge>, AgentMemoryError> {
        memory_helpers::relations_of(
            &self.db,
            &self.collection_name,
            id,
            &self.ttl,
            MemoryKind::Procedural,
        )
    }

    /// Removes a relation edge created by [`Self::relate`].
    ///
    /// # Errors
    ///
    /// Returns `CollectionError` when the collection cannot be resolved.
    pub fn unrelate(&self, edge_id: u64) -> Result<bool, AgentMemoryError> {
        memory_helpers::unrelate_edge(&self.db, &self.collection_name, edge_id)
    }

    /// Recalls matching procedures by vector similarity.
    ///
    /// When ACT-R activation decay is configured via
    /// [`with_activation_decay`](Self::with_activation_decay), the returned
    /// `confidence` reflects power-law decay since last use without modifying
    /// the stored value.
    ///
    /// # Errors
    ///
    /// Returns an error when embedding dimension is invalid, collection access fails,
    /// or vector search fails.
    pub fn recall(
        &self,
        query_embedding: &[f32],
        k: usize,
        min_confidence: f32,
    ) -> Result<Vec<ProcedureMatch>, AgentMemoryError> {
        let results = memory_helpers::search_filtered(
            &self.db,
            &self.collection_name,
            self.dimension,
            query_embedding,
            k,
            &self.ttl,
            MemoryKind::Procedural,
        )?;

        let now_secs = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_or(0, |d| d.as_secs());

        Ok(results
            .into_iter()
            .filter_map(|r| {
                let mut pm = extract_procedure_match(&r.point, r.score)?;
                if let Some(exponent) = self.activation_decay_exponent {
                    let last_used = r
                        .point
                        .payload
                        .as_ref()
                        .and_then(|p| p.get("last_used_at"))
                        .and_then(serde_json::Value::as_i64)
                        .unwrap_or(0);
                    let elapsed_secs = (now_secs as i64).saturating_sub(last_used).max(0) as u64;
                    pm.confidence = power_law_decay(pm.confidence, elapsed_secs, exponent);
                }
                if pm.confidence < min_confidence {
                    return None;
                }
                Some(pm)
            })
            .collect())
    }

    /// Reinforces a stored procedure using the configured strategy.
    ///
    /// # Errors
    ///
    /// Returns an error when procedure retrieval or update fails.
    pub fn reinforce(&self, procedure_id: u64, success: bool) -> Result<(), AgentMemoryError> {
        self.reinforce_with_strategy(procedure_id, success, &*self.reinforcement_strategy)
    }

    /// Reinforces a stored procedure using a custom strategy.
    ///
    /// # Errors
    ///
    /// Returns an error when procedure retrieval or update fails.
    pub fn reinforce_with_strategy(
        &self,
        procedure_id: u64,
        success: bool,
        strategy: &dyn ReinforcementStrategy,
    ) -> Result<(), AgentMemoryError> {
        let collection = memory_helpers::get_collection(&self.db, &self.collection_name)?;

        let points = collection.get(&[procedure_id]);
        let point = points
            .into_iter()
            .flatten()
            .next()
            .ok_or_else(|| AgentMemoryError::NotFound(format!("Procedure {procedure_id}")))?;

        let state = Self::extract_procedure_state(&point)?;
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_or(0, |d| d.as_secs() as i64);

        let context = state.build_reinforcement_context(now);
        let new_confidence = strategy.update_confidence(state.confidence, success, &context);
        let (new_success, new_failure) = if success {
            (state.success_count + 1, state.failure_count)
        } else {
            (state.success_count, state.failure_count + 1)
        };

        let mut payload = json!({
            "name": state.name,
            "steps": state.steps,
            "confidence": new_confidence,
            "usage_count": state.usage_count + 1,
            "created_at": state.created_at,
            "last_used_at": now,
            "success_count": new_success,
            "failure_count": new_failure
        });
        // Preserve the durable TTL field: reinforcing must not strip expiry.
        let prior_expiry = point
            .payload
            .as_ref()
            .and_then(|p| p.get(memory_helpers::EXPIRES_AT_KEY))
            .and_then(serde_json::Value::as_u64);
        memory_helpers::attach_expiry(&mut payload, prior_expiry);
        let updated_point = Point::new(procedure_id, point.vector.clone(), Some(payload));

        memory_helpers::upsert_points(&collection, vec![updated_point])?;

        Ok(())
    }

    fn extract_procedure_state(point: &Point) -> Result<ProcedureState, AgentMemoryError> {
        let payload = point
            .payload
            .as_ref()
            .ok_or_else(|| AgentMemoryError::CollectionError("Missing payload".to_string()))?;

        Ok(ProcedureState {
            name: payload
                .get("name")
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_string(),
            steps: payload
                .get("steps")
                .and_then(|v| v.as_array())
                .map(|arr| {
                    arr.iter()
                        .filter_map(|v| v.as_str().map(String::from))
                        .collect()
                })
                .unwrap_or_default(),
            confidence: payload
                .get("confidence")
                .and_then(serde_json::Value::as_f64)
                .unwrap_or(0.5) as f32,
            usage_count: payload
                .get("usage_count")
                .and_then(serde_json::Value::as_u64)
                .unwrap_or(0),
            created_at: payload
                .get("created_at")
                .and_then(serde_json::Value::as_i64)
                .unwrap_or(0),
            last_used_at: payload
                .get("last_used_at")
                .and_then(serde_json::Value::as_i64)
                .unwrap_or(0),
            success_count: payload
                .get("success_count")
                .and_then(serde_json::Value::as_u64)
                .unwrap_or(0),
            failure_count: payload
                .get("failure_count")
                .and_then(serde_json::Value::as_u64)
                .unwrap_or(0),
        })
    }

    /// Lists all tracked procedures.
    ///
    /// # Errors
    ///
    /// Returns an error when collection access fails.
    pub fn list_all(&self) -> Result<Vec<ProcedureMatch>, AgentMemoryError> {
        let collection = memory_helpers::get_collection(&self.db, &self.collection_name)?;

        let all_ids: Vec<u64> = self.stored_ids.read().iter().copied().collect();
        let points = collection.get(&all_ids);

        Ok(points
            .into_iter()
            .flatten()
            .filter(|p| !self.ttl.is_expired(MemoryKind::Procedural, p.id))
            .filter_map(|p| extract_procedure_match(&p, 0.0))
            .collect())
    }

    /// Deletes a procedure by id.
    ///
    /// # Errors
    ///
    /// Returns an error when collection access or deletion fails.
    pub fn delete(&self, id: u64) -> Result<(), AgentMemoryError> {
        memory_helpers::delete_tracked_point(
            &self.db,
            &self.collection_name,
            id,
            &self.stored_ids,
            &self.ttl,
            MemoryKind::Procedural,
        )
    }

    /// Serializes all procedures into snapshot bytes.
    ///
    /// # Errors
    ///
    /// Returns an error when collection access or JSON encoding fails.
    pub fn serialize(&self) -> Result<Vec<u8>, AgentMemoryError> {
        memory_helpers::serialize_tracked_points(&self.db, &self.collection_name, &self.stored_ids)
    }

    /// Replaces procedural memory state from serialized snapshot bytes.
    ///
    /// # Errors
    ///
    /// Returns an error when JSON decoding fails, collection access fails,
    /// or persistence operations fail.
    pub fn deserialize(&self, data: &[u8]) -> Result<(), AgentMemoryError> {
        memory_helpers::deserialize_tracked_points(
            &self.db,
            &self.collection_name,
            data,
            &self.stored_ids,
        )
    }
}

/// Extracts a `ProcedureMatch` from a point's payload with the given similarity score.
fn extract_procedure_match(point: &Point, score: f32) -> Option<ProcedureMatch> {
    let payload = point.payload.as_ref()?;
    let name = payload.get("name")?.as_str()?.to_string();
    let steps: Vec<String> = payload
        .get("steps")?
        .as_array()?
        .iter()
        .filter_map(|v| v.as_str().map(String::from))
        .collect();
    #[allow(clippy::cast_possible_truncation)]
    let confidence = payload.get("confidence")?.as_f64()? as f32;

    Some(ProcedureMatch {
        id: point.id,
        name,
        steps,
        confidence,
        score,
    })
}