ceylon_next/memory/advanced/

consolidation.rs

//! Memory Consolidation - Background Processing and Maintenance
//!
//! This module handles:
//! - Background memory processing
//! - Duplicate detection and merging
//! - Memory aging and decay
//! - Memory cleanup and optimization

use super::{EnhancedMemoryEntry, ImportanceLevel, MemoryConfig};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio::time::{interval, Duration};

/// Tasks that can be performed during consolidation
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConsolidationTask {
    /// Detect and merge duplicate memories
    DeduplicateMemories,
    /// Apply decay to old memories
    ApplyDecay,
    /// Clean up low-value memories
    Cleanup,
    /// Reindex memories for faster retrieval
    Reindex,
    /// All consolidation tasks
    All,
}

/// Result of a consolidation operation
#[derive(Debug, Clone)]
pub struct ConsolidationResult {
    pub task: ConsolidationTask,
    pub memories_processed: usize,
    pub memories_merged: usize,
    pub memories_deleted: usize,
    pub duration_ms: u64,
}

/// Memory consolidator
pub struct Consolidator {
    /// Memories being managed
    memories: Arc<RwLock<HashMap<String, EnhancedMemoryEntry>>>,
    /// Configuration
    config: MemoryConfig,
    /// Running flag
    running: Arc<RwLock<bool>>,
}

impl Consolidator {
    /// Create a new consolidator
    pub fn new(config: MemoryConfig) -> Self {
        Self {
            memories: Arc::new(RwLock::new(HashMap::new())),
            config,
            running: Arc::new(RwLock::new(false)),
        }
    }

    /// Start background consolidation process
    pub async fn start_background_consolidation(&self) {
        let memories = self.memories.clone();
        let config = self.config.clone();
        let running = self.running.clone();

        // Set running flag
        *running.write().await = true;

        tokio::spawn(async move {
            let mut ticker = interval(Duration::from_secs(config.consolidation_interval));

            loop {
                ticker.tick().await;

                // Check if still running
                if !*running.read().await {
                    break;
                }

                // Perform consolidation
                let consolidator = Consolidator {
                    memories: memories.clone(),
                    config: config.clone(),
                    running: running.clone(),
                };

                if let Err(e) = consolidator.consolidate(ConsolidationTask::All).await {
                    eprintln!("Consolidation error: {}", e);
                }
            }
        });
    }

    /// Stop background consolidation
    pub async fn stop_background_consolidation(&self) {
        *self.running.write().await = false;
    }

    /// Perform consolidation tasks
    pub async fn consolidate(&self, task: ConsolidationTask) -> Result<ConsolidationResult, String> {
        let start = std::time::Instant::now();
        let mut result = ConsolidationResult {
            task: task.clone(),
            memories_processed: 0,
            memories_merged: 0,
            memories_deleted: 0,
            duration_ms: 0,
        };

        match task {
            ConsolidationTask::DeduplicateMemories => {
                result = self.deduplicate_memories().await?;
            }
            ConsolidationTask::ApplyDecay => {
                result = self.apply_decay().await?;
            }
            ConsolidationTask::Cleanup => {
                result = self.cleanup_memories().await?;
            }
            ConsolidationTask::Reindex => {
                result = self.reindex_memories().await?;
            }
            ConsolidationTask::All => {
                // Run all tasks
                let dedup = self.deduplicate_memories().await?;
                let decay = self.apply_decay().await?;
                let cleanup = self.cleanup_memories().await?;

                result.memories_processed = dedup.memories_processed + decay.memories_processed + cleanup.memories_processed;
                result.memories_merged = dedup.memories_merged;
                result.memories_deleted = cleanup.memories_deleted;
            }
        }

        result.duration_ms = start.elapsed().as_millis() as u64;
        Ok(result)
    }

    /// Detect and merge duplicate memories
    async fn deduplicate_memories(&self) -> Result<ConsolidationResult, String> {
        let mut memories = self.memories.write().await;
        let memory_list: Vec<EnhancedMemoryEntry> = memories.values().cloned().collect();

        let mut processed = 0;
        let mut merged = 0;
        let mut to_remove = Vec::new();

        // Compare each pair of memories
        for i in 0..memory_list.len() {
            for j in (i + 1)..memory_list.len() {
                processed += 1;

                let similarity = self.calculate_similarity(&memory_list[i], &memory_list[j]);

                if similarity >= self.config.duplicate_threshold {
                    // Merge the two memories
                    let (keep_id, remove_id) = if memory_list[i].relevance_score() >= memory_list[j].relevance_score() {
                        (&memory_list[i].entry.id, &memory_list[j].entry.id)
                    } else {
                        (&memory_list[j].entry.id, &memory_list[i].entry.id)
                    };

                    // Merge metadata - clone data first to avoid borrow conflicts
                    let merge_data = memories.get(remove_id).map(|remove_entry| {
                        (
                            remove_entry.access_count,
                            remove_entry.key_points.clone(),
                            remove_entry.entities.clone(),
                        )
                    });

                    if let Some((access_count, key_points, entities)) = merge_data {
                        if let Some(keep_entry) = memories.get_mut(keep_id) {
                            keep_entry.access_count += access_count;
                            keep_entry.related_memories.push(remove_id.clone());

                            // Merge key points
                            for point in &key_points {
                                if !keep_entry.key_points.contains(point) {
                                    keep_entry.key_points.push(point.clone());
                                }
                            }

                            // Merge entities
                            for entity in &entities {
                                if !keep_entry.entities.contains(entity) {
                                    keep_entry.entities.push(entity.clone());
                                }
                            }
                        }
                    }

                    to_remove.push(remove_id.clone());
                    merged += 1;
                }
            }
        }

        // Remove duplicates
        for id in to_remove {
            memories.remove(&id);
        }

        Ok(ConsolidationResult {
            task: ConsolidationTask::DeduplicateMemories,
            memories_processed: processed,
            memories_merged: merged,
            memories_deleted: 0,
            duration_ms: 0,
        })
    }

    /// Apply decay to memories
    async fn apply_decay(&self) -> Result<ConsolidationResult, String> {
        if !self.config.enable_decay {
            return Ok(ConsolidationResult {
                task: ConsolidationTask::ApplyDecay,
                memories_processed: 0,
                memories_merged: 0,
                memories_deleted: 0,
                duration_ms: 0,
            });
        }

        let mut memories = self.memories.write().await;
        let mut processed = 0;
        let current_time = Self::current_timestamp();

        for memory in memories.values_mut() {
            // Don't decay important memories
            if memory.importance >= self.config.min_importance_for_preservation {
                continue;
            }

            // Calculate age in days
            let age_seconds = current_time.saturating_sub(memory.entry.created_at);
            let age_days = age_seconds as f32 / 86400.0;

            // Apply exponential decay
            let decay_amount = 1.0 - (self.config.decay_rate * age_days);
            memory.decay_factor = (memory.decay_factor * decay_amount).max(0.0);

            processed += 1;
        }

        Ok(ConsolidationResult {
            task: ConsolidationTask::ApplyDecay,
            memories_processed: processed,
            memories_merged: 0,
            memories_deleted: 0,
            duration_ms: 0,
        })
    }

    /// Clean up low-value memories
    async fn cleanup_memories(&self) -> Result<ConsolidationResult, String> {
        let mut memories = self.memories.write().await;
        let mut processed = 0;
        let mut deleted = 0;

        // Collect IDs to remove
        let to_remove: Vec<String> = memories
            .iter()
            .filter(|(_, memory)| {
                processed += 1;
                // Remove memories with very low relevance score or fully decayed
                memory.relevance_score() < 0.1 || memory.decay_factor < 0.1
            })
            .map(|(id, _)| id.clone())
            .collect();

        deleted = to_remove.len();

        // Remove them
        for id in to_remove {
            memories.remove(&id);
        }

        Ok(ConsolidationResult {
            task: ConsolidationTask::Cleanup,
            memories_processed: processed,
            memories_merged: 0,
            memories_deleted: deleted,
            duration_ms: 0,
        })
    }

    /// Reindex memories for faster retrieval
    async fn reindex_memories(&self) -> Result<ConsolidationResult, String> {
        let memories = self.memories.read().await;

        Ok(ConsolidationResult {
            task: ConsolidationTask::Reindex,
            memories_processed: memories.len(),
            memories_merged: 0,
            memories_deleted: 0,
            duration_ms: 0,
        })
    }

    /// Calculate similarity between two memories
    fn calculate_similarity(&self, mem1: &EnhancedMemoryEntry, mem2: &EnhancedMemoryEntry) -> f32 {
        // Simple similarity based on:
        // 1. Time proximity
        // 2. Task ID match
        // 3. Content overlap (key points, entities)

        let mut similarity = 0.0;
        let mut factors = 0;

        // Time proximity (within 1 hour = high similarity)
        let time_diff = (mem1.entry.created_at as i64 - mem2.entry.created_at as i64).abs();
        if time_diff < 3600 {
            similarity += 0.3;
        } else if time_diff < 86400 {
            similarity += 0.1;
        }
        factors += 1;

        // Task ID match
        if mem1.entry.task_id == mem2.entry.task_id {
            similarity += 0.3;
            factors += 1;
        }

        // Agent ID match
        if mem1.entry.agent_id == mem2.entry.agent_id {
            similarity += 0.1;
            factors += 1;
        }

        // Key points overlap
        let common_points = mem1
            .key_points
            .iter()
            .filter(|p| mem2.key_points.contains(p))
            .count();
        if !mem1.key_points.is_empty() || !mem2.key_points.is_empty() {
            let max_points = mem1.key_points.len().max(mem2.key_points.len());
            if max_points > 0 {
                similarity += 0.2 * (common_points as f32 / max_points as f32);
                factors += 1;
            }
        }

        // Entities overlap
        let common_entities = mem1
            .entities
            .iter()
            .filter(|e| mem2.entities.contains(e))
            .count();
        if !mem1.entities.is_empty() || !mem2.entities.is_empty() {
            let max_entities = mem1.entities.len().max(mem2.entities.len());
            if max_entities > 0 {
                similarity += 0.1 * (common_entities as f32 / max_entities as f32);
                factors += 1;
            }
        }

        if factors > 0 {
            similarity / factors as f32
        } else {
            0.0
        }
    }

    /// Update memories reference
    pub async fn set_memories(&self, memories: HashMap<String, EnhancedMemoryEntry>) {
        *self.memories.write().await = memories;
    }

    /// Get current memories
    pub async fn get_memories(&self) -> HashMap<String, EnhancedMemoryEntry> {
        self.memories.read().await.clone()
    }

    fn current_timestamp() -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::memory::MemoryEntry;
    use super::super::MemoryType;

    #[tokio::test]
    async fn test_consolidator_decay() {
        let config = MemoryConfig {
            enable_decay: true,
            decay_rate: 0.1,
            min_importance_for_preservation: ImportanceLevel::High,
            ..Default::default()
        };

        let consolidator = Consolidator::new(config);

        // Add some memories
        let mut memories = HashMap::new();
        for i in 0..5 {
            let entry = MemoryEntry::new(
                "agent-1".to_string(),
                format!("task-{}", i),
                vec![],
            );
            let mut enhanced = EnhancedMemoryEntry::new(entry, MemoryType::Episodic);
            enhanced.importance = if i < 2 {
                ImportanceLevel::Critical
            } else {
                ImportanceLevel::Low
            };
            memories.insert(enhanced.entry.id.clone(), enhanced);
        }

        consolidator.set_memories(memories).await;

        let result = consolidator.apply_decay().await.unwrap();
        assert_eq!(result.memories_processed, 3); // Only low importance ones
    }

    #[tokio::test]
    async fn test_consolidator_cleanup() {
        let config = MemoryConfig::default();
        let consolidator = Consolidator::new(config);

        // Add memories with different relevance
        let mut memories = HashMap::new();
        for i in 0..5 {
            let entry = MemoryEntry::new(
                "agent-1".to_string(),
                format!("task-{}", i),
                vec![],
            );
            let mut enhanced = EnhancedMemoryEntry::new(entry, MemoryType::Episodic);
            enhanced.decay_factor = if i < 2 { 0.05 } else { 0.8 }; // Low decay = should be cleaned
            memories.insert(enhanced.entry.id.clone(), enhanced);
        }

        consolidator.set_memories(memories).await;

        let result = consolidator.cleanup_memories().await.unwrap();
        assert_eq!(result.memories_deleted, 2); // Two with low decay factor
    }
}