ceylon_next/memory/

sqlite.rs

//! SQLite-backed persistent memory implementation.
//!
//! This module provides persistent storage using SQLite database,
//! suitable for production use with efficient indexed queries.

use super::{Memory, MemoryEntry};
use crate::llm::types::Message;
use async_trait::async_trait;
use rusqlite::{params, Connection, OptionalExtension};
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

/// A SQLite-backed implementation of the [`Memory`] trait.
///
/// This provides persistent storage of conversation history using SQLite.
/// All conversations are stored in a local SQLite database file.
///
/// # Examples
///
/// ```rust,no_run
/// use ceylon_next::memory::SqliteStore;
/// use std::sync::Arc;
///
/// #[tokio::main]
/// async fn main() {
///     let memory = Arc::new(
///         SqliteStore::new("agent_memory.db").await.unwrap()
///     );
/// }
/// ```
pub struct SqliteStore {
    db_path: PathBuf,
    connection: Arc<Mutex<Connection>>,
}

impl SqliteStore {
    /// Creates a new SQLite store with the given database path.
    ///
    /// # Arguments
    ///
    /// * `db_path` - Path to the SQLite database file (will be created if it doesn't exist)
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use ceylon_next::memory::SqliteStore;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let store = SqliteStore::new("memory.db").await.unwrap();
    /// }
    /// ```
    pub async fn new<P: Into<PathBuf>>(db_path: P) -> Result<Self, String> {
        let db_path = db_path.into();
        let connection = Connection::open(&db_path)
            .map_err(|e| format!("Failed to open database: {}", e))?;

        // Create the memory table if it doesn't exist
        connection
            .execute(
                "CREATE TABLE IF NOT EXISTS memory_entries (
                    id TEXT PRIMARY KEY,
                    agent_id TEXT NOT NULL,
                    task_id TEXT NOT NULL,
                    messages TEXT NOT NULL,
                    created_at INTEGER NOT NULL
                )",
                [],
            )
            .map_err(|e| format!("Failed to create table: {}", e))?;

        // Create index for faster agent lookups
        connection
            .execute(
                "CREATE INDEX IF NOT EXISTS idx_agent_id ON memory_entries(agent_id, created_at DESC)",
                [],
            )
            .map_err(|e| format!("Failed to create index: {}", e))?;

        Ok(Self {
            db_path,
            connection: Arc::new(Mutex::new(connection)),
        })
    }

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

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

        tokio::task::spawn_blocking(move || {
            let messages_json = serde_json::to_string(&entry.messages)
                .map_err(|e| format!("Failed to serialize messages: {}", e))?;

            let conn = connection.lock().unwrap();
            conn.execute(
                "INSERT INTO memory_entries (id, agent_id, task_id, messages, created_at)
                 VALUES (?1, ?2, ?3, ?4, ?5)",
                params![
                    entry.id,
                    entry.agent_id,
                    entry.task_id,
                    messages_json,
                    entry.created_at as i64
                ],
            )
            .map_err(|e| format!("Failed to store entry: {}", e))?;

            Ok(entry_id)
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }

    async fn get(&self, id: &str) -> Result<Option<MemoryEntry>, String> {
        let id = id.to_string();
        let connection = Arc::clone(&self.connection);

        tokio::task::spawn_blocking(move || {
            let conn = connection.lock().unwrap();
            let mut stmt = conn
                .prepare("SELECT id, agent_id, task_id, messages, created_at FROM memory_entries WHERE id = ?1")
                .map_err(|e| format!("Failed to prepare statement: {}", e))?;

            let result = stmt
                .query_row(params![id], |row| {
                    let messages_json: String = row.get(3)?;
                    let messages: Vec<Message> = serde_json::from_str(&messages_json)
                        .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;

                    Ok(MemoryEntry {
                        id: row.get(0)?,
                        agent_id: row.get(1)?,
                        task_id: row.get(2)?,
                        messages,
                        created_at: row.get::<_, i64>(4)? as u64,
                    })
                })
                .optional()
                .map_err(|e| format!("Failed to query entry: {}", e))?;

            Ok(result)
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }

    async fn get_agent_history(&self, agent_id: &str) -> Result<Vec<MemoryEntry>, String> {
        let agent_id = agent_id.to_string();
        let connection = Arc::clone(&self.connection);

        tokio::task::spawn_blocking(move || {
            let conn = connection.lock().unwrap();
            let mut stmt = conn
                .prepare(
                    "SELECT id, agent_id, task_id, messages, created_at
                     FROM memory_entries
                     WHERE agent_id = ?1
                     ORDER BY created_at DESC",
                )
                .map_err(|e| format!("Failed to prepare statement: {}", e))?;

            let entries = stmt
                .query_map(params![agent_id], |row| {
                    let messages_json: String = row.get(3)?;
                    let messages: Vec<Message> = serde_json::from_str(&messages_json)
                        .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;

                    Ok(MemoryEntry {
                        id: row.get(0)?,
                        agent_id: row.get(1)?,
                        task_id: row.get(2)?,
                        messages,
                        created_at: row.get::<_, i64>(4)? as u64,
                    })
                })
                .map_err(|e| format!("Failed to query entries: {}", e))?
                .collect::<Result<Vec<_>, _>>()
                .map_err(|e| format!("Failed to collect entries: {}", e))?;

            Ok(entries)
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }

    async fn get_recent(&self, agent_id: &str, limit: usize) -> Result<Vec<MemoryEntry>, String> {
        let agent_id = agent_id.to_string();
        let connection = Arc::clone(&self.connection);

        tokio::task::spawn_blocking(move || {
            let conn = connection.lock().unwrap();
            let mut stmt = conn
                .prepare(
                    "SELECT id, agent_id, task_id, messages, created_at
                     FROM memory_entries
                     WHERE agent_id = ?1
                     ORDER BY created_at DESC
                     LIMIT ?2",
                )
                .map_err(|e| format!("Failed to prepare statement: {}", e))?;

            let entries = stmt
                .query_map(params![agent_id, limit as i64], |row| {
                    let messages_json: String = row.get(3)?;
                    let messages: Vec<Message> = serde_json::from_str(&messages_json)
                        .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;

                    Ok(MemoryEntry {
                        id: row.get(0)?,
                        agent_id: row.get(1)?,
                        task_id: row.get(2)?,
                        messages,
                        created_at: row.get::<_, i64>(4)? as u64,
                    })
                })
                .map_err(|e| format!("Failed to query entries: {}", e))?
                .collect::<Result<Vec<_>, _>>()
                .map_err(|e| format!("Failed to collect entries: {}", e))?;

            Ok(entries)
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }

    async fn search(&self, agent_id: &str, query: &str) -> Result<Vec<MemoryEntry>, String> {
        let agent_id = agent_id.to_string();
        let query = query.to_string();
        let connection = Arc::clone(&self.connection);

        tokio::task::spawn_blocking(move || {
            let conn = connection.lock().unwrap();
            let mut stmt = conn
                .prepare(
                    "SELECT id, agent_id, task_id, messages, created_at
                     FROM memory_entries
                     WHERE agent_id = ?1 AND messages LIKE ?2
                     ORDER BY created_at DESC",
                )
                .map_err(|e| format!("Failed to prepare statement: {}", e))?;

            let search_pattern = format!("%{}%", query);
            let entries = stmt
                .query_map(params![agent_id, search_pattern], |row| {
                    let messages_json: String = row.get(3)?;
                    let messages: Vec<Message> = serde_json::from_str(&messages_json)
                        .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;

                    Ok(MemoryEntry {
                        id: row.get(0)?,
                        agent_id: row.get(1)?,
                        task_id: row.get(2)?,
                        messages,
                        created_at: row.get::<_, i64>(4)? as u64,
                    })
                })
                .map_err(|e| format!("Failed to query entries: {}", e))?
                .collect::<Result<Vec<_>, _>>()
                .map_err(|e| format!("Failed to collect entries: {}", e))?;

            Ok(entries)
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }

    async fn clear_agent_memory(&self, agent_id: &str) -> Result<(), String> {
        let agent_id = agent_id.to_string();
        let connection = Arc::clone(&self.connection);

        tokio::task::spawn_blocking(move || {
            let conn = connection.lock().unwrap();
            conn.execute(
                "DELETE FROM memory_entries WHERE agent_id = ?1",
                params![agent_id],
            )
            .map_err(|e| format!("Failed to clear agent memory: {}", e))?;

            Ok(())
        })
        .await
        .map_err(|e| format!("Task join error: {}", e))?
    }
}