pe_core/

cognitive_memory.rs

//! Cognitive memory types — working notes, categories, and memory configuration.
//!
//! These types form the agent's scratchpad — structured notes that the agent
//! writes to itself during execution. Unlike free-text logs, these are
//! queryable by category, sortable by relevance, and decayable over time.

use serde::{Deserialize, Serialize};

/// A note the agent writes to itself during execution.
///
/// Like a human jotting thoughts while working. Notes have a category
/// for structured queries and a relevance score that can decay over time.
///
/// # Example
///
/// ```
/// use pe_core::cognitive_memory::{WorkingNote, NoteCategory};
///
/// let note = WorkingNote::new(
///     "User prefers concise responses",
///     NoteCategory::Observation,
/// );
/// assert_eq!(note.relevance, 1.0);
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct WorkingNote {
    /// The note content.
    pub content: String,

    /// Classification for structured queries.
    pub category: NoteCategory,

    /// ISO 8601 timestamp when the note was created.
    pub created_at: String,

    /// Relevance score — 1.0 = just created, decays toward 0.0 over time.
    pub relevance: f64,
}

impl WorkingNote {
    /// Create a new working note with full relevance.
    pub fn new(content: impl Into<String>, category: NoteCategory) -> Self {
        Self {
            content: content.into(),
            category,
            created_at: String::new(), // Caller sets timestamp
            relevance: 1.0,
        }
    }

    /// Create a note with a specific timestamp.
    ///
    /// If not called, `created_at` will be empty — meaning "unknown creation time".
    /// Callers are responsible for setting timestamps when recency matters.
    #[must_use]
    pub fn with_timestamp(mut self, timestamp: impl Into<String>) -> Self {
        self.created_at = timestamp.into();
        self
    }

    /// Decay the relevance by a factor (0.0 to 1.0).
    pub fn decay(&mut self, factor: f64) {
        self.relevance = (self.relevance * factor).max(0.0);
    }
}

/// Classification for working notes.
///
/// Enables structured queries — a lobe can ask "give me all Concerns"
/// instead of parsing free text.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum NoteCategory {
    /// "I noticed X" — things observed during execution.
    Observation,
    /// "I should try X next" — planned actions.
    Plan,
    /// "X might be a problem" — potential issues flagged.
    Concern,
    /// "I found that X" — new information discovered.
    Discovery,
    /// "Looking back, X worked/didn't work" — lessons learned.
    Reflection,
    /// "Come back to X later" — deferred items.
    Bookmark,
    /// User-defined category.
    Custom(String),
}

/// Configuration for the agent's memory tiers.
///
/// Controls how much memory the agent allocates to each tier
/// and how aggressively to compress when limits are reached.
///
/// # Example
///
/// ```
/// use pe_core::cognitive_memory::MemoryConfig;
///
/// let config = MemoryConfig::default();
/// assert_eq!(config.short_term_limit, 8192);
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct MemoryConfig {
    /// Max tokens for short-term memory (current conversation).
    pub short_term_limit: u32,

    /// Max tokens for working memory (session scratchpad).
    pub working_limit: u32,

    /// Max tokens for long-term memory (cross-session, compressed).
    pub long_term_limit: u32,

    /// How aggressively to compress (0.0 = keep everything, 1.0 = compress hard).
    pub compression_ratio: f64,
}

impl Default for MemoryConfig {
    fn default() -> Self {
        Self {
            short_term_limit: 8192,
            working_limit: 2048,
            long_term_limit: 16384,
            compression_ratio: 0.5,
        }
    }
}

/// Configuration for the meditate (memory consolidation) operation.
///
/// Meditate is the agent's "sleep cycle" — it prunes stale notes, decays relevance,
/// consolidates related notes, and optionally uses an LLM for deeper synthesis.
/// All fields have sensible defaults for typical agent workloads.
///
/// # Example
///
/// ```
/// use pe_core::cognitive_memory::MeditateConfig;
///
/// let config = MeditateConfig::default();
/// assert_eq!(config.prune_threshold, 0.1);
/// assert!(config.use_llm);
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct MeditateConfig {
    /// Relevance threshold below which notes are pruned.
    /// Notes with `relevance < prune_threshold` are removed during the prune phase.
    pub prune_threshold: f64,

    /// Decay factor applied to all notes during prune (0.0-1.0).
    /// Lower values are more aggressive. Applied before the threshold check.
    pub decay_factor: f64,

    /// Max working notes before meditate auto-triggers (for MeditateLobe).
    pub max_notes_before_trigger: usize,

    /// Max failure records before meditate auto-triggers.
    pub max_failures_before_trigger: usize,

    /// Whether to use an LLM for consolidation and indexing phases.
    /// When false, uses algorithmic fallback (group-by-category, truncate).
    pub use_llm: bool,

    /// Max notes to keep after consolidation. Excess pruned by relevance.
    pub max_notes_after_consolidation: usize,

    /// Max failure records to keep after pruning. Oldest resolved go first.
    pub max_failures_after_prune: usize,
}

impl Default for MeditateConfig {
    fn default() -> Self {
        Self {
            prune_threshold: 0.1,
            decay_factor: 0.9,
            max_notes_before_trigger: 50,
            max_failures_before_trigger: 20,
            use_llm: true,
            max_notes_after_consolidation: 30,
            max_failures_after_prune: 10,
        }
    }
}

/// Summary of what a meditate operation did.
///
/// Returned after each meditate cycle for logging, debugging, and agent
/// self-awareness. An agent can inspect this to understand how its memory
/// was reshaped.
///
/// # Example
///
/// ```
/// use pe_core::cognitive_memory::MeditateResult;
///
/// let result = MeditateResult {
///     notes_before: 45,
///     notes_after: 28,
///     notes_pruned: 12,  // includes threshold + cap removals
///     notes_merged: 5,
///     failures_pruned: 3,
///     constraints_removed: 1,
///     used_llm: true,
///     insights_summary: Some("Consolidated 5 observation notes into 2".into()),
/// };
/// assert!(result.notes_after <= result.notes_before);
/// ```
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct MeditateResult {
    /// Number of working notes before meditate ran.
    pub notes_before: usize,

    /// Number of working notes after meditate completed.
    pub notes_after: usize,

    /// Number of notes removed during prune phase (threshold + cap combined).
    pub notes_pruned: usize,

    /// Number of notes merged during consolidation.
    pub notes_merged: usize,

    /// Number of failure records pruned (resolved or stale).
    pub failures_pruned: usize,

    /// Number of constraints removed (no longer applicable).
    pub constraints_removed: usize,

    /// Whether the LLM was used for consolidation.
    pub used_llm: bool,

    /// Optional summary of insights generated during consolidation.
    pub insights_summary: Option<String>,
}

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

    #[test]
    fn test_working_note_creation() {
        let note = WorkingNote::new("test note", NoteCategory::Observation);
        assert_eq!(note.content, "test note");
        assert_eq!(note.category, NoteCategory::Observation);
        assert_eq!(note.relevance, 1.0);
    }

    #[test]
    fn test_working_note_with_timestamp() {
        let note =
            WorkingNote::new("test", NoteCategory::Plan).with_timestamp("2026-03-23T12:00:00Z");
        assert_eq!(note.created_at, "2026-03-23T12:00:00Z");
    }

    #[test]
    fn test_working_note_decay() {
        let mut note = WorkingNote::new("test", NoteCategory::Discovery);
        assert_eq!(note.relevance, 1.0);
        note.decay(0.9);
        assert!((note.relevance - 0.9).abs() < f64::EPSILON);
        note.decay(0.5);
        assert!((note.relevance - 0.45).abs() < f64::EPSILON);
    }

    #[test]
    fn test_decay_floor_at_zero() {
        let mut note = WorkingNote::new("test", NoteCategory::Concern);
        note.relevance = 0.01;
        note.decay(0.0);
        assert_eq!(note.relevance, 0.0);
    }

    #[test]
    fn test_note_category_custom() {
        let cat = NoteCategory::Custom("legal_flag".into());
        let json = serde_json::to_string(&cat).unwrap();
        let back: NoteCategory = serde_json::from_str(&json).unwrap();
        assert_eq!(back, NoteCategory::Custom("legal_flag".into()));
    }

    #[test]
    fn test_memory_config_defaults() {
        let config = MemoryConfig::default();
        assert_eq!(config.short_term_limit, 8192);
        assert_eq!(config.working_limit, 2048);
        assert_eq!(config.long_term_limit, 16384);
        assert!((config.compression_ratio - 0.5).abs() < f64::EPSILON);
    }

    #[test]
    fn test_working_note_serialization() {
        let note = WorkingNote::new("important", NoteCategory::Reflection)
            .with_timestamp("2026-03-23T10:00:00Z");
        let json = serde_json::to_string(&note).unwrap();
        let back: WorkingNote = serde_json::from_str(&json).unwrap();
        assert_eq!(back, note);
    }

    #[test]
    fn test_meditate_config_defaults() {
        let config = MeditateConfig::default();
        assert!((config.prune_threshold - 0.1).abs() < f64::EPSILON);
        assert!((config.decay_factor - 0.9).abs() < f64::EPSILON);
        assert_eq!(config.max_notes_before_trigger, 50);
        assert_eq!(config.max_failures_before_trigger, 20);
        assert!(config.use_llm);
        assert_eq!(config.max_notes_after_consolidation, 30);
        assert_eq!(config.max_failures_after_prune, 10);
    }

    #[test]
    fn test_meditate_config_serialization() {
        let config = MeditateConfig {
            prune_threshold: 0.2,
            decay_factor: 0.8,
            max_notes_before_trigger: 100,
            max_failures_before_trigger: 10,
            use_llm: false,
            max_notes_after_consolidation: 50,
            max_failures_after_prune: 5,
        };
        let json = serde_json::to_string(&config).unwrap();
        let back: MeditateConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(back, config);
    }

    #[test]
    fn test_meditate_result_default() {
        let result = MeditateResult::default();
        assert_eq!(result.notes_before, 0);
        assert_eq!(result.notes_after, 0);
        assert_eq!(result.notes_pruned, 0);
        assert_eq!(result.notes_merged, 0);
        assert_eq!(result.failures_pruned, 0);
        assert_eq!(result.constraints_removed, 0);
        assert!(!result.used_llm);
        assert_eq!(result.insights_summary, None);
    }

    #[test]
    fn test_meditate_result_serialization() {
        let result = MeditateResult {
            notes_before: 45,
            notes_after: 28,
            notes_pruned: 12,
            notes_merged: 5,
            failures_pruned: 3,
            constraints_removed: 1,
            used_llm: true,
            insights_summary: Some("Consolidated observation notes".into()),
        };
        let json = serde_json::to_string(&result).unwrap();
        let back: MeditateResult = serde_json::from_str(&json).unwrap();
        assert_eq!(back, result);
        assert_eq!(
            back.insights_summary.as_deref(),
            Some("Consolidated observation notes")
        );
    }
}