ceylon_next/memory/

file_store.rs

//! File-based persistent memory implementation.
//!
//! This module provides file-based storage with support for both JSON
//! (human-readable) and MessagePack (compact binary) formats.

use super::{Memory, MemoryEntry};
use async_trait::async_trait;
use std::collections::HashMap;
use std::fs;
use std::io::{Read, Write};
use std::path::PathBuf;
use std::sync::{Arc, RwLock};

/// Serialization format for file-based storage
#[derive(Debug, Clone, Copy)]
pub enum StorageFormat {
    /// JSON format (human-readable)
    Json,
    /// MessagePack format (compact binary)
    MessagePack,
}

/// A file-based implementation of the [`Memory`] trait.
///
/// This provides persistent storage of conversation history using files.
/// Supports both JSON (human-readable) and MessagePack (compact binary) formats.
///
/// # Examples
///
/// ```rust,no_run
/// use ceylon_next::memory::{FileStore, StorageFormat};
/// use std::sync::Arc;
///
/// #[tokio::main]
/// async fn main() {
///     // JSON format
///     let memory = Arc::new(
///         FileStore::new("memory.json", StorageFormat::Json).await.unwrap()
///     );
///
///     // MessagePack format
///     let memory_mp = Arc::new(
///         FileStore::new("memory.msgpack", StorageFormat::MessagePack).await.unwrap()
///     );
/// }
/// ```
pub struct FileStore {
    file_path: PathBuf,
    format: StorageFormat,
    entries: Arc<RwLock<HashMap<String, MemoryEntry>>>,
    agent_index: Arc<RwLock<HashMap<String, Vec<String>>>>,
}

impl FileStore {
    /// Creates a new file-based store with the given file path and format.
    ///
    /// # Arguments
    ///
    /// * `file_path` - Path to the storage file
    /// * `format` - Serialization format (JSON or MessagePack)
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::memory::{FileStore, StorageFormat};
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let store = FileStore::new("memory.json", StorageFormat::Json).await.unwrap();
    /// }
    /// ```
    pub async fn new<P: Into<PathBuf>>(file_path: P, format: StorageFormat) -> Result<Self, String> {
        let file_path = file_path.into();
        let mut entries: HashMap<String, MemoryEntry> = HashMap::new();
        let mut agent_index: HashMap<String, Vec<String>> = HashMap::new();

        // Load existing data if file exists
        if file_path.exists() {
            let mut file = fs::File::open(&file_path)
                .map_err(|e| format!("Failed to open file: {}", e))?;

            let mut buffer = Vec::new();
            file.read_to_end(&mut buffer)
                .map_err(|e| format!("Failed to read file: {}", e))?;

            if !buffer.is_empty() {
                entries = match format {
                    StorageFormat::Json => serde_json::from_slice(&buffer)
                        .map_err(|e| format!("Failed to deserialize JSON: {}", e))?,
                    StorageFormat::MessagePack => rmp_serde::from_slice(&buffer)
                        .map_err(|e| format!("Failed to deserialize MessagePack: {}", e))?,
                };

                // Rebuild agent index
                for (entry_id, entry) in &entries {
                    agent_index
                        .entry(entry.agent_id.clone())
                        .or_insert_with(Vec::new)
                        .push(entry_id.clone());
                }
            }
        }

        Ok(Self {
            file_path,
            format,
            entries: Arc::new(RwLock::new(entries)),
            agent_index: Arc::new(RwLock::new(agent_index)),
        })
    }

    /// Returns the path to the storage file.
    pub fn file_path(&self) -> &PathBuf {
        &self.file_path
    }

    /// Returns the storage format being used.
    pub fn format(&self) -> StorageFormat {
        self.format
    }

    /// Persists the current in-memory data to disk.
    fn save_to_disk(&self) -> Result<(), String> {
        let entries = self.entries.read().unwrap();

        let serialized = match self.format {
            StorageFormat::Json => serde_json::to_vec_pretty(&*entries)
                .map_err(|e| format!("Failed to serialize to JSON: {}", e))?,
            StorageFormat::MessagePack => rmp_serde::to_vec(&*entries)
                .map_err(|e| format!("Failed to serialize to MessagePack: {}", e))?,
        };

        let mut file = fs::File::create(&self.file_path)
            .map_err(|e| format!("Failed to create file: {}", e))?;

        file.write_all(&serialized)
            .map_err(|e| format!("Failed to write to file: {}", e))?;

        Ok(())
    }
}

#[async_trait]
impl Memory for FileStore {
    async fn store(&self, entry: MemoryEntry) -> Result<String, String> {
        let entry_id = entry.id.clone();
        let agent_id = entry.agent_id.clone();

        // Store the entry
        {
            let mut entries = self.entries.write().unwrap();
            entries.insert(entry_id.clone(), entry);
        }

        // Update agent index
        {
            let mut index = self.agent_index.write().unwrap();
            index
                .entry(agent_id)
                .or_insert_with(Vec::new)
                .push(entry_id.clone());
        }

        // Persist to disk
        self.save_to_disk()?;

        Ok(entry_id)
    }

    async fn get(&self, id: &str) -> Result<Option<MemoryEntry>, String> {
        let entries = self.entries.read().unwrap();
        Ok(entries.get(id).cloned())
    }

    async fn get_agent_history(&self, agent_id: &str) -> Result<Vec<MemoryEntry>, String> {
        let index = self.agent_index.read().unwrap();
        let entries = self.entries.read().unwrap();

        let entry_ids = index.get(agent_id);

        if entry_ids.is_none() {
            return Ok(Vec::new());
        }

        let mut result = Vec::new();
        for entry_id in entry_ids.unwrap() {
            if let Some(entry) = entries.get(entry_id) {
                result.push(entry.clone());
            }
        }

        // Sort by timestamp (newest first)
        result.sort_by(|a, b| b.created_at.cmp(&a.created_at));

        Ok(result)
    }

    async fn get_recent(&self, agent_id: &str, limit: usize) -> Result<Vec<MemoryEntry>, String> {
        let mut history = self.get_agent_history(agent_id).await?;
        history.truncate(limit);
        Ok(history)
    }

    async fn search(&self, agent_id: &str, query: &str) -> Result<Vec<MemoryEntry>, String> {
        let history = self.get_agent_history(agent_id).await?;
        let query_lower = query.to_lowercase();

        let results: Vec<MemoryEntry> = history
            .into_iter()
            .filter(|entry| {
                entry
                    .messages
                    .iter()
                    .any(|msg| msg.content.to_lowercase().contains(&query_lower))
            })
            .collect();

        Ok(results)
    }

    async fn clear_agent_memory(&self, agent_id: &str) -> Result<(), String> {
        let index = self.agent_index.read().unwrap();
        let entry_ids = index.get(agent_id);

        if let Some(ids) = entry_ids {
            let mut entries = self.entries.write().unwrap();
            for id in ids {
                entries.remove(id);
            }
        }

        let mut index = self.agent_index.write().unwrap();
        index.remove(agent_id);

        // Persist to disk
        drop(index); // Release the lock before calling save_to_disk
        self.save_to_disk()?;

        Ok(())
    }
}