ceylon_next/memory/advanced/

working.rs

//! Working Memory - Recent Context Management
//!
//! Working memory maintains the most recent and relevant context for the agent.
//! It's designed for fast access and limited capacity, similar to human short-term memory.
//!
//! # Features
//!
//! - Fixed capacity with LRU (Least Recently Used) eviction
//! - Token-aware context management
//! - Automatic summarization of evicted memories
//! - Priority-based retention

use super::{EnhancedMemoryEntry, ImportanceLevel, MemoryConfig, MemoryType};
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::RwLock;

/// Working memory for managing recent context
pub struct WorkingMemory {
    /// Recent memories (ordered by recency)
    memories: Arc<RwLock<VecDeque<EnhancedMemoryEntry>>>,
    /// Configuration
    config: MemoryConfig,
    /// Current token count estimate
    token_count: Arc<RwLock<usize>>,
}

impl WorkingMemory {
    /// Create a new working memory instance
    pub fn new(config: MemoryConfig) -> Self {
        Self {
            memories: Arc::new(RwLock::new(VecDeque::new())),
            config,
            token_count: Arc::new(RwLock::new(0)),
        }
    }

    /// Add a memory to working memory
    pub async fn add(&self, mut entry: EnhancedMemoryEntry) -> Result<(), String> {
        entry.memory_type = MemoryType::Working;

        let mut memories = self.memories.write().await;
        let mut token_count = self.token_count.write().await;

        // Estimate tokens for this entry
        let entry_tokens = Self::estimate_tokens(&entry);
        *token_count += entry_tokens;

        // Add to front (most recent)
        memories.push_front(entry);

        // Evict if over capacity
        while memories.len() > self.config.working_memory_limit {
            if let Some(evicted) = memories.pop_back() {
                let evicted_tokens = Self::estimate_tokens(&evicted);
                *token_count = token_count.saturating_sub(evicted_tokens);
            }
        }

        Ok(())
    }

    /// Get all working memories
    pub async fn get_all(&self) -> Vec<EnhancedMemoryEntry> {
        let memories = self.memories.read().await;
        memories.iter().cloned().collect()
    }

    /// Get recent memories up to a token limit
    pub async fn get_recent_within_limit(&self, max_tokens: usize) -> Vec<EnhancedMemoryEntry> {
        let memories = self.memories.read().await;
        let mut result = Vec::new();
        let mut current_tokens = 0;

        for memory in memories.iter() {
            let memory_tokens = Self::estimate_tokens(memory);
            if current_tokens + memory_tokens > max_tokens {
                break;
            }
            result.push(memory.clone());
            current_tokens += memory_tokens;
        }

        result
    }

    /// Get memories by importance level
    pub async fn get_by_importance(&self, min_importance: ImportanceLevel) -> Vec<EnhancedMemoryEntry> {
        let memories = self.memories.read().await;
        memories
            .iter()
            .filter(|m| m.importance >= min_importance)
            .cloned()
            .collect()
    }

    /// Mark a memory as accessed (moves it to front)
    pub async fn mark_accessed(&self, memory_id: &str) -> Result<(), String> {
        let mut memories = self.memories.write().await;

        if let Some(pos) = memories.iter().position(|m| m.entry.id == memory_id) {
            if let Some(mut memory) = memories.remove(pos) {
                memory.mark_accessed();
                memories.push_front(memory);
            }
        }

        Ok(())
    }

    /// Clear all working memory
    pub async fn clear(&self) {
        let mut memories = self.memories.write().await;
        let mut token_count = self.token_count.write().await;
        memories.clear();
        *token_count = 0;
    }

    /// Get current token count
    pub async fn token_count(&self) -> usize {
        *self.token_count.read().await
    }

    /// Get current memory count
    pub async fn memory_count(&self) -> usize {
        self.memories.read().await.len()
    }

    /// Create context message from working memory
    pub async fn create_context(&self, max_tokens: Option<usize>) -> String {
        let memories = if let Some(limit) = max_tokens {
            self.get_recent_within_limit(limit).await
        } else {
            self.get_all().await
        };

        if memories.is_empty() {
            return String::new();
        }

        let mut context = String::from("RECENT CONTEXT:\n");

        for (i, memory) in memories.iter().enumerate() {
            // Use summary if available, otherwise use full conversation
            if let Some(summary) = &memory.summary {
                context.push_str(&format!("{}. {}\n", i + 1, summary));
            } else {
                context.push_str(&format!("{}. Conversation from {}:\n",
                    i + 1,
                    Self::format_timestamp(memory.entry.created_at)
                ));

                for msg in &memory.entry.messages {
                    if msg.role == "user" {
                        context.push_str(&format!("  User: {}\n", msg.content));
                    } else if msg.role == "assistant" {
                        context.push_str(&format!("  Assistant: {}\n", msg.content));
                    }
                }
            }

            // Add key points if available
            if !memory.key_points.is_empty() {
                context.push_str("  Key points:\n");
                for point in &memory.key_points {
                    context.push_str(&format!("    - {}\n", point));
                }
            }

            context.push_str("\n");
        }

        context
    }

    /// Estimate token count for a memory entry
    fn estimate_tokens(entry: &EnhancedMemoryEntry) -> usize {
        // Rough estimation: ~4 characters per token
        let mut total_chars = 0;

        // Count message content
        for msg in &entry.entry.messages {
            total_chars += msg.content.len();
        }

        // Count summary if present
        if let Some(summary) = &entry.summary {
            total_chars += summary.len();
        }

        // Count key points
        for point in &entry.key_points {
            total_chars += point.len();
        }

        (total_chars / 4).max(1)
    }

    /// Format timestamp for display
    fn format_timestamp(timestamp: u64) -> String {
        use std::time::{SystemTime, UNIX_EPOCH, Duration};

        let dt = UNIX_EPOCH + Duration::from_secs(timestamp);
        let elapsed = SystemTime::now().duration_since(dt).unwrap_or_default();

        let seconds = elapsed.as_secs();
        if seconds < 60 {
            format!("{} seconds ago", seconds)
        } else if seconds < 3600 {
            format!("{} minutes ago", seconds / 60)
        } else if seconds < 86400 {
            format!("{} hours ago", seconds / 3600)
        } else {
            format!("{} days ago", seconds / 86400)
        }
    }
}

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

    #[tokio::test]
    async fn test_working_memory_basic() {
        let config = MemoryConfig {
            working_memory_limit: 3,
            ..Default::default()
        };
        let wm = WorkingMemory::new(config);

        // Add memories
        for i in 0..5 {
            let entry = MemoryEntry::new(
                "agent-1".to_string(),
                format!("task-{}", i),
                vec![],
            );
            let enhanced = EnhancedMemoryEntry::new(entry, MemoryType::Working);
            wm.add(enhanced).await.unwrap();
        }

        // Should only keep 3 most recent
        assert_eq!(wm.memory_count().await, 3);
    }

    #[tokio::test]
    async fn test_working_memory_access() {
        let config = MemoryConfig::default();
        let wm = WorkingMemory::new(config);

        let entry1 = MemoryEntry::new("agent-1".to_string(), "task-1".to_string(), vec![]);
        let id1 = entry1.id.clone();
        let enhanced1 = EnhancedMemoryEntry::new(entry1, MemoryType::Working);
        wm.add(enhanced1).await.unwrap();

        let entry2 = MemoryEntry::new("agent-1".to_string(), "task-2".to_string(), vec![]);
        let enhanced2 = EnhancedMemoryEntry::new(entry2, MemoryType::Working);
        wm.add(enhanced2).await.unwrap();

        // Access first entry (should move it to front)
        wm.mark_accessed(&id1).await.unwrap();

        let memories = wm.get_all().await;
        assert_eq!(memories[0].entry.id, id1);
    }

    #[tokio::test]
    async fn test_working_memory_importance() {
        let config = MemoryConfig::default();
        let wm = WorkingMemory::new(config);

        let entry1 = MemoryEntry::new("agent-1".to_string(), "task-1".to_string(), vec![]);
        let mut enhanced1 = EnhancedMemoryEntry::new(entry1, MemoryType::Working);
        enhanced1.importance = ImportanceLevel::Critical;
        wm.add(enhanced1).await.unwrap();

        let entry2 = MemoryEntry::new("agent-1".to_string(), "task-2".to_string(), vec![]);
        let mut enhanced2 = EnhancedMemoryEntry::new(entry2, MemoryType::Working);
        enhanced2.importance = ImportanceLevel::Low;
        wm.add(enhanced2).await.unwrap();

        let important = wm.get_by_importance(ImportanceLevel::High).await;
        assert_eq!(important.len(), 1);
        assert_eq!(important[0].importance, ImportanceLevel::Critical);
    }
}