ceylon_next/memory/advanced/

manager.rs

//! Advanced Memory Manager - Main orchestrator

use super::{
    consolidation::{Consolidator, ConsolidationTask},
    episodic::{EpisodicMemory, TimeRange},
    semantic::SemanticMemory,
    summarization::{Summarizer, SummaryStrategy},
    working::WorkingMemory,
    EnhancedMemoryEntry, ImportanceLevel, MemoryConfig, MemoryType,
};
use crate::memory::{Memory, MemoryEntry};
use async_trait::async_trait;
use std::sync::Arc;

/// Advanced Memory Manager
///
/// Orchestrates all memory subsystems:
/// - Working Memory: Recent context
/// - Episodic Memory: Past conversations
/// - Semantic Memory: Facts and knowledge
/// - Summarization: Automatic compression
/// - Consolidation: Background processing
pub struct AdvancedMemoryManager {
    /// Backend storage
    backend: Arc<dyn Memory>,
    /// Working memory
    working: WorkingMemory,
    /// Episodic memory
    episodic: EpisodicMemory,
    /// Semantic memory
    semantic: SemanticMemory,
    /// Summarizer
    summarizer: Summarizer,
    /// Consolidator
    consolidator: Consolidator,
    /// Configuration
    config: MemoryConfig,
}

impl AdvancedMemoryManager {
    /// Create a new advanced memory manager
    pub async fn new(
        backend: Arc<dyn Memory>,
        config: MemoryConfig,
    ) -> Result<Self, String> {
        let working = WorkingMemory::new(config.clone());
        let episodic = EpisodicMemory::new(backend.clone(), config.clone()).await?;
        let semantic = SemanticMemory::new();
        let summarizer = Summarizer::new();
        let consolidator = Consolidator::new(config.clone());

        Ok(Self {
            backend,
            working,
            episodic,
            semantic,
            summarizer,
            consolidator,
            config,
        })
    }

    // ============================================
    // Core Memory Operations
    // ============================================

    /// Store a new conversation
    pub async fn store_conversation(&self, entry: MemoryEntry) -> Result<String, String> {
        // Create enhanced entry
        let mut enhanced = EnhancedMemoryEntry::new(entry, MemoryType::Working);

        // Summarize if needed
        if self.config.auto_summarize && enhanced.entry.messages.len() >= self.config.min_summary_length {
            let summary = self
                .summarizer
                .summarize(&enhanced, SummaryStrategy::Hybrid)
                .await?;

            enhanced.summary = Some(summary.summary);
            enhanced.key_points = summary.key_points;
            enhanced.importance = summary.importance;

            // Extract entities to string names
            enhanced.entities = summary.entities.iter().map(|e| e.name.clone()).collect();

            // Add entities to semantic memory
            for entity in summary.entities {
                self.semantic
                    .add_entity(entity.name, entity.entity_type, enhanced.entry.id.clone())
                    .await;
            }

            // Add facts to semantic memory
            for fact in summary.facts {
                self.semantic
                    .add_fact(fact.subject, fact.predicate, fact.object, enhanced.entry.id.clone())
                    .await;
            }
        }

        // Add to working memory
        self.working.add(enhanced.clone()).await?;

        // Store in episodic memory
        let id = self.episodic.store(enhanced).await?;

        Ok(id)
    }

    /// Retrieve a memory by ID
    pub async fn get(&self, id: &str) -> Result<Option<EnhancedMemoryEntry>, String> {
        self.episodic.get(id).await
    }

    /// Get recent working memory
    pub async fn get_working_memory(&self, max_tokens: Option<usize>) -> Vec<EnhancedMemoryEntry> {
        if let Some(limit) = max_tokens {
            self.working.get_recent_within_limit(limit).await
        } else {
            self.working.get_all().await
        }
    }

    /// Get episodic memories by time range
    pub async fn get_episodic_by_time(
        &self,
        agent_id: &str,
        range: TimeRange,
    ) -> Result<Vec<EnhancedMemoryEntry>, String> {
        self.episodic.get_by_time_range(agent_id, range).await
    }

    /// Get relevant memories by score
    pub async fn get_relevant_memories(
        &self,
        agent_id: &str,
        min_score: f32,
        limit: usize,
    ) -> Result<Vec<EnhancedMemoryEntry>, String> {
        self.episodic.get_by_relevance(agent_id, min_score, limit).await
    }

    /// Search memories
    pub async fn search(&self, agent_id: &str, query: &str) -> Result<Vec<EnhancedMemoryEntry>, String> {
        self.episodic.search(agent_id, query).await
    }

    // ============================================
    // Semantic Memory Operations
    // ============================================

    /// Get knowledge about an entity
    pub async fn get_knowledge(&self, entity_name: &str) -> Option<super::semantic::KnowledgeSummary> {
        self.semantic.get_knowledge_summary(entity_name).await
    }

    /// Get all facts about a subject
    pub async fn get_facts_about(&self, subject: &str) -> Vec<super::semantic::Fact> {
        self.semantic.get_facts_about(subject).await
    }

    /// Get semantic memory statistics
    pub async fn get_semantic_stats(&self) -> super::semantic::SemanticMemoryStats {
        self.semantic.get_statistics().await
    }

    // ============================================
    // Context Generation
    // ============================================

    /// Create a comprehensive context message for the agent
    pub async fn create_agent_context(
        &self,
        agent_id: &str,
        max_tokens: Option<usize>,
        include_semantic: bool,
    ) -> Result<String, String> {
        let mut context = String::new();

        // 1. Working Memory (recent context)
        let working_context = self.working.create_context(max_tokens).await;
        if !working_context.is_empty() {
            context.push_str(&working_context);
            context.push_str("\n");
        }

        // 2. Relevant Episodic Memories (based on importance)
        let important_memories = self
            .episodic
            .get_by_importance(agent_id, ImportanceLevel::High)
            .await?;

        if !important_memories.is_empty() {
            context.push_str("IMPORTANT PAST CONVERSATIONS:\n");
            for (i, memory) in important_memories.iter().take(3).enumerate() {
                if let Some(summary) = &memory.summary {
                    context.push_str(&format!("{}. {}\n", i + 1, summary));
                }
            }
            context.push_str("\n");
        }

        // 3. Semantic Knowledge (if requested)
        if include_semantic {
            let stats = self.semantic.get_statistics().await;
            if stats.total_entities > 0 || stats.total_facts > 0 {
                context.push_str("LEARNED KNOWLEDGE:\n");
                context.push_str(&format!("- {} entities tracked\n", stats.total_entities));
                context.push_str(&format!("- {} facts stored\n", stats.total_facts));

                // Include high-confidence facts
                let high_conf_facts = self.semantic.get_high_confidence_facts(0.8).await;
                if !high_conf_facts.is_empty() {
                    context.push_str("Key facts:\n");
                    for fact in high_conf_facts.iter().take(5) {
                        context.push_str(&format!(
                            "  - {} {} {}\n",
                            fact.subject, fact.predicate, fact.object
                        ));
                    }
                }
                context.push_str("\n");
            }
        }

        Ok(context)
    }

    // ============================================
    // Maintenance Operations
    // ============================================

    /// Run consolidation tasks
    pub async fn consolidate(&self, task: ConsolidationTask) -> Result<super::consolidation::ConsolidationResult, String> {
        // Get all memories from episodic
        let memories = self.episodic.get_by_time_range("", TimeRange::All).await?;

        // Convert to HashMap for consolidator
        let mut mem_map = std::collections::HashMap::new();
        for memory in memories {
            mem_map.insert(memory.entry.id.clone(), memory);
        }

        self.consolidator.set_memories(mem_map).await;

        // Run consolidation
        let result = self.consolidator.consolidate(task).await?;

        // Update episodic memory with consolidated memories
        // (In a real implementation, we'd persist these changes)

        Ok(result)
    }

    /// Start background consolidation
    pub async fn start_background_processing(&self) {
        if self.config.auto_consolidate {
            self.consolidator.start_background_consolidation().await;
        }
    }

    /// Stop background consolidation
    pub async fn stop_background_processing(&self) {
        self.consolidator.stop_background_consolidation().await;
    }

    /// Get statistics about memory usage
    pub async fn get_statistics(&self) -> MemoryStatistics {
        let semantic_stats = self.semantic.get_statistics().await;

        MemoryStatistics {
            working_memory_count: self.working.memory_count().await,
            working_memory_tokens: self.working.token_count().await,
            semantic_entities: semantic_stats.total_entities,
            semantic_facts: semantic_stats.total_facts,
            semantic_relationships: semantic_stats.total_relationships,
        }
    }

    /// Clear all memory for an agent
    pub async fn clear_agent_memory(&self, agent_id: &str) -> Result<(), String> {
        self.working.clear().await;
        self.episodic.clear(agent_id).await?;
        Ok(())
    }
}

/// Implement the base Memory trait for compatibility
#[async_trait]
impl Memory for AdvancedMemoryManager {
    async fn store(&self, entry: MemoryEntry) -> Result<String, String> {
        self.store_conversation(entry).await
    }

    async fn get(&self, id: &str) -> Result<Option<MemoryEntry>, String> {
        if let Some(enhanced) = self.get(id).await? {
            Ok(Some(enhanced.entry))
        } else {
            Ok(None)
        }
    }

    async fn get_agent_history(&self, agent_id: &str) -> Result<Vec<MemoryEntry>, String> {
        let enhanced_entries = self
            .episodic
            .get_by_time_range(agent_id, TimeRange::All)
            .await?;

        Ok(enhanced_entries.into_iter().map(|e| e.entry).collect())
    }

    async fn get_recent(&self, agent_id: &str, limit: usize) -> Result<Vec<MemoryEntry>, String> {
        let enhanced_entries = self
            .episodic
            .get_by_relevance(agent_id, 0.0, limit)
            .await?;

        Ok(enhanced_entries.into_iter().map(|e| e.entry).collect())
    }

    async fn search(&self, agent_id: &str, query: &str) -> Result<Vec<MemoryEntry>, String> {
        let enhanced_entries = self.search(agent_id, query).await?;
        Ok(enhanced_entries.into_iter().map(|e| e.entry).collect())
    }

    async fn clear_agent_memory(&self, agent_id: &str) -> Result<(), String> {
        self.clear_agent_memory(agent_id).await
    }
}

/// Memory usage statistics
#[derive(Debug, Clone)]
pub struct MemoryStatistics {
    pub working_memory_count: usize,
    pub working_memory_tokens: usize,
    pub semantic_entities: usize,
    pub semantic_facts: usize,
    pub semantic_relationships: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::memory::InMemoryStore;
    use crate::llm::types::Message;

    #[tokio::test]
    async fn test_advanced_memory_manager_basic() {
        let backend = Arc::new(InMemoryStore::new());
        let config = MemoryConfig {
            auto_consolidate: false, // Disable to prevent background task from blocking test
            ..Default::default()
        };
        let manager = AdvancedMemoryManager::new(backend, config).await.unwrap();

        // Store a conversation
        let messages = vec![
            Message {
                role: "user".to_string(),
                content: "Hello, tell me about Rust.".to_string(),
            },
            Message {
                role: "assistant".to_string(),
                content: "Rust is a systems programming language.".to_string(),
            },
        ];

        let entry = MemoryEntry::new("agent-1".to_string(), "task-1".to_string(), messages);
        let id = manager.store_conversation(entry).await.unwrap();

        // Retrieve it
        let retrieved = manager.get(&id).await.unwrap();
        assert!(retrieved.is_some());
    }

    #[tokio::test]
    async fn test_context_generation() {
        let backend = Arc::new(InMemoryStore::new());
        let config = MemoryConfig {
            auto_consolidate: false, // Disable to prevent background task from blocking test
            ..Default::default()
        };
        let manager = AdvancedMemoryManager::new(backend, config).await.unwrap();

        // Store some conversations
        for i in 0..3 {
            let messages = vec![Message {
                role: "user".to_string(),
                content: format!("Message {}", i),
            }];

            let entry = MemoryEntry::new(
                "agent-1".to_string(),
                format!("task-{}", i),
                messages,
            );
            manager.store_conversation(entry).await.unwrap();
        }

        // Generate context
        let context = manager
            .create_agent_context("agent-1", None, true)
            .await
            .unwrap();

        assert!(!context.is_empty());
    }
}