llm_brain/

db.rs

use std::path::Path;
use std::sync::Arc;

use async_trait::async_trait;
use serde::Deserialize;
use serde_json::Value as JsonValue;
use surrealdb::Surreal;
use surrealdb::engine::local::{Db, SurrealKv};
use surrealdb::sql::Thing;

use crate::error::{LLMBrainError, Result};
use crate::models::MemoryFragment;

/// Trait defining the interface for the memory storage layer.
#[async_trait]
pub trait Store: Send + Sync {
    /// Adds a new memory fragment to the store.
    /// Returns the ID (Thing) of the newly created fragment.
    async fn add_memory(&self, fragment: MemoryFragment) -> Result<Thing>;

    /// Retrieves a memory fragment by its unique ID.
    async fn get_memory(&self, id: Thing) -> Result<Option<MemoryFragment>>;

    /// Finds the top_k most similar memory fragments to the given embedding.
    /// Returns a list of (MemoryFragment, similarity_score) tuples.
    async fn find_similar(
        &self, embedding: &[f32], top_k: usize,
    ) -> Result<Vec<(MemoryFragment, f32)>>;

    /// Updates the last_accessed_at timestamp for a given memory fragment.
    async fn record_access(&self, id: Thing) -> Result<()>;
}

/// Concrete implementation of the Store trait using `SurrealDB`.
#[derive(Clone)]
pub struct SurrealStore {
    db: Arc<Surreal<Db>>,
}

impl SurrealStore {
    /// Establishes a connection to the embedded `SurrealDB` (`SurrealKV`)
    /// database at the specified path and selects the namespace and
    /// database.
    pub async fn connect(
        db_path: impl AsRef<Path>, namespace: &str, database: &str,
    ) -> Result<Self> {
        // Ensure the parent directory exists
        if let Some(parent) = db_path.as_ref().parent() {
            tokio::fs::create_dir_all(parent).await.map_err(|e| {
                LLMBrainError::InitializationError(format!(
                    "Failed to create database directory '{}': {}",
                    parent.display(),
                    e
                ))
            })?;
        }

        // Configure and open the SurrealKV engine
        let db = Surreal::new::<SurrealKv>(db_path.as_ref()).await?;

        // Select the namespace and database
        db.use_ns(namespace).use_db(database).await?;

        // Perform any necessary initialization
        Self::initialize_database(&db).await?;

        Ok(Self { db: Arc::new(db) })
    }

    /// Performs initial database setup like defining tables and vector indexes.
    async fn initialize_database(db: &Surreal<Db>) -> Result<()> {
        // Define the 'memory' table schema
        let _ = db
            .query(
                r#"
            DEFINE TABLE memory SCHEMAFULL;
            DEFINE FIELD content ON TABLE memory TYPE string;
            DEFINE FIELD embedding ON TABLE memory TYPE array<float>;
            DEFINE FIELD metadata ON TABLE memory TYPE object;
            DEFINE FIELD created_at ON TABLE memory TYPE datetime DEFAULT time::now();
            DEFINE FIELD last_accessed_at ON TABLE memory TYPE option<datetime>;
            "#,
            )
            .await?;

        // Define a vector index on the embedding field for efficient similarity search
        let _ = db
            .query(
                r#"
            DEFINE INDEX memory_embedding_index ON TABLE memory FIELDS embedding
                MTREE DIMENSION 1536 DISTANCE COSINE;
            "#,
            )
            .await?;

        // Define an index on created_at for potential time-based queries
        let _ = db
            .query(
                r#"
            DEFINE INDEX memory_created_at_index ON TABLE memory FIELDS created_at;
            "#,
            )
            .await?;

        Ok(())
    }
}

#[async_trait]
impl Store for SurrealStore {
    async fn add_memory(&self, fragment: MemoryFragment) -> Result<Thing> {
        #[derive(Deserialize, Debug)]
        struct CreatedId {
            id: Thing,
        }

        // Rely on database default for created_at and initial last_accessed_at.
        // Bind CLONES of content, embedding, and metadata.
        let mut response = self
            .db
            .query(
                "CREATE memory SET content = $content, embedding = $embedding, metadata = $metadata RETURN id"
            )
            .bind(("content", fragment.content.clone())) // Clone
            .bind(("embedding", fragment.embedding.clone())) // Clone
            .bind(("metadata", fragment.metadata.clone())) // Clone
            // Do not bind created_at or last_accessed_at
            .await?;

        let created: Option<CreatedId> = response.take(0)?;
        match created {
            Some(created) => Ok(created.id),
            None => {
                Err(LLMBrainError::DatabaseError(Box::new(
                    surrealdb::Error::Api(surrealdb::error::Api::Query(
                        "No record returned after creation".into(),
                    )),
                )))
            }
        }
    }

    async fn get_memory(&self, id: Thing) -> Result<Option<MemoryFragment>> {
        // Bind the owned Thing directly
        let mut response = self
            .db
            .query("SELECT * FROM $id")
            .bind(("id", id)) // Bind owned Thing
            .await?;

        let memory: Option<MemoryFragment> = response.take(0)?;
        Ok(memory)
    }

    async fn find_similar(
        &self, embedding: &[f32], top_k: usize,
    ) -> Result<Vec<(MemoryFragment, f32)>> {
        // Explicitly select fields and calculate similarity
        let query = "SELECT
                        content,
                        embedding,
                        metadata,
                        created_at,
                        last_accessed_at,
                        vector::similarity::cosine(embedding, $embedding) as similarity
                    FROM memory
                    ORDER BY similarity DESC
                    LIMIT $limit";

        // Define QueryResult matching the explicit SELECT list
        #[derive(Deserialize, Debug)]
        struct QueryResult {
            content: String,
            embedding: Vec<f32>,
            metadata: JsonValue,
            created_at: JsonValue,       // Expect JsonValue from DB datetime
            last_accessed_at: JsonValue, // Expect JsonValue from DB option<datetime>
            similarity: f32,
        }

        let embedding_vec = embedding.to_vec();
        let mut response = self
            .db
            .query(query)
            .bind(("embedding", embedding_vec))
            .bind(("limit", top_k))
            .await?;

        let results: Vec<QueryResult> = response.take(0)?;

        // Manually reconstruct MemoryFragment from QueryResult fields
        let scored_fragments = results
            .into_iter()
            .map(|res| {
                let fragment = MemoryFragment {
                    content: res.content,
                    embedding: res.embedding,
                    metadata: res.metadata,
                    created_at: res.created_at,
                    last_accessed_at: res.last_accessed_at,
                };
                (fragment, res.similarity.clamp(0.0, 1.0))
            })
            .collect();

        Ok(scored_fragments)
    }

    async fn record_access(&self, id: Thing) -> Result<()> {
        // Bind the owned Thing directly
        let _ = self
            .db
            .query("UPDATE $id SET last_accessed_at = time::now()")
            .bind(("id", id)) // Bind owned Thing
            .await?;
        Ok(())
    }
}