ceylon_next/memory/vector/

mod.rs

//! Vector memory and semantic search capabilities.
//!
//! This module provides vector embeddings and semantic search functionality
//! for agent memories. It enables similarity-based memory retrieval using
//! vector embeddings rather than simple keyword matching.
//!
//! # Features
//!
//! - **Embedding Generation**: Convert text to vector embeddings using various providers
//! - **Semantic Search**: Find similar memories based on meaning, not just keywords
//! - **Hybrid Search**: Combine semantic similarity with traditional keyword search
//! - **Multiple Backends**: Support for local and cloud vector databases
//!
//! # Architecture
//!
//! The vector memory system is built around three core traits:
//!
//! - [`EmbeddingProvider`] - Generates vector embeddings from text
//! - [`VectorStore`] - Stores and retrieves vectors with metadata
//! - [`VectorMemory`] - High-level interface combining embeddings and storage
//!
//! # Available Backends
//!
//! - [`LocalVectorStore`] - In-memory vector store with cosine similarity search
//!
//! # Example
//!
//! ```rust,no_run
//! use ceylon_next::memory::vector::{EmbeddingProvider, VectorMemory, LocalVectorStore};
//! use std::sync::Arc;
//!
//! #[tokio::main]
//! async fn main() {
//!     // Create an embedding provider (e.g., OpenAI, Ollama, etc.)
//!     // let embedder = OpenAIEmbeddings::new("api-key");
//!
//!     // Create a vector store
//!     let store = LocalVectorStore::new(384); // 384-dimensional embeddings
//!
//!     // Use for semantic search
//!     // let results = store.search(&query_vector, 5).await.unwrap();
//! }
//! ```

use async_trait::async_trait;
use serde::{Deserialize, Serialize};

// ============================================
// Core Types
// ============================================

/// A vector embedding with associated metadata.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VectorEntry {
    /// Unique identifier for this vector entry
    pub id: String,
    /// The vector embedding
    pub vector: Vec<f32>,
    /// Associated memory entry ID
    pub memory_id: String,
    /// Agent ID for filtering
    pub agent_id: String,
    /// Original text that was embedded
    pub text: String,
    /// Optional metadata
    pub metadata: Option<serde_json::Value>,
    /// Unix timestamp when created
    pub created_at: u64,
}

impl VectorEntry {
    /// Creates a new vector entry.
    pub fn new(
        memory_id: String,
        agent_id: String,
        text: String,
        vector: Vec<f32>,
        metadata: Option<serde_json::Value>,
    ) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            vector,
            memory_id,
            agent_id,
            text,
            metadata,
            created_at: Self::current_timestamp(),
        }
    }

    fn current_timestamp() -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs()
    }
}

/// A search result with similarity score.
#[derive(Debug, Clone)]
pub struct SearchResult {
    /// The vector entry that matched
    pub entry: VectorEntry,
    /// Similarity score (typically 0.0 to 1.0, higher is more similar)
    pub score: f32,
}

// ============================================
// Embedding Provider Trait
// ============================================

/// Trait for generating vector embeddings from text.
///
/// Implementations can use different embedding models:
/// - OpenAI's text-embedding models
/// - Local models via Ollama
/// - Hugging Face models
/// - Custom embedding models
///
/// # Example
///
/// ```rust,no_run
/// use ceylon_next::memory::vector::EmbeddingProvider;
/// use async_trait::async_trait;
///
/// struct MyEmbedder;
///
/// #[async_trait]
/// impl EmbeddingProvider for MyEmbedder {
///     async fn embed(&self, text: &str) -> Result<Vec<f32>, String> {
///         // Generate embedding...
///         Ok(vec![0.1, 0.2, 0.3])
///     }
///
///     fn dimension(&self) -> usize {
///         3
///     }
/// }
/// ```
#[async_trait]
pub trait EmbeddingProvider: Send + Sync {
    /// Generates an embedding vector for the given text.
    ///
    /// # Arguments
    ///
    /// * `text` - The text to embed
    ///
    /// # Returns
    ///
    /// A vector of floats representing the embedding, or an error message.
    async fn embed(&self, text: &str) -> Result<Vec<f32>, String>;

    /// Generates embeddings for multiple texts in batch.
    ///
    /// Default implementation calls `embed` for each text sequentially.
    /// Providers can override this for more efficient batch processing.
    async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, String> {
        let mut embeddings = Vec::with_capacity(texts.len());
        for text in texts {
            embeddings.push(self.embed(text).await?);
        }
        Ok(embeddings)
    }

    /// Returns the dimensionality of the embeddings produced by this provider.
    fn dimension(&self) -> usize;

    /// Returns the model name or identifier.
    fn model_name(&self) -> &str {
        "unknown"
    }
}

// ============================================
// Vector Store Trait
// ============================================

/// Trait for storing and retrieving vector embeddings.
///
/// Vector stores handle the low-level storage and similarity search
/// of vector embeddings. Different implementations can use different
/// backends like in-memory storage, vector databases, etc.
#[async_trait]
pub trait VectorStore: Send + Sync {
    /// Stores a vector entry.
    async fn store(&self, entry: VectorEntry) -> Result<String, String>;

    /// Stores multiple vector entries in batch.
    async fn store_batch(&self, entries: Vec<VectorEntry>) -> Result<Vec<String>, String> {
        let mut ids = Vec::with_capacity(entries.len());
        for entry in entries {
            ids.push(self.store(entry).await?);
        }
        Ok(ids)
    }

    /// Retrieves a vector entry by ID.
    async fn get(&self, id: &str) -> Result<Option<VectorEntry>, String>;

    /// Searches for similar vectors using cosine similarity.
    ///
    /// # Arguments
    ///
    /// * `query_vector` - The query vector to search for
    /// * `agent_id` - Optional agent ID to filter results
    /// * `limit` - Maximum number of results to return
    /// * `threshold` - Optional minimum similarity score threshold (0.0 to 1.0)
    ///
    /// # Returns
    ///
    /// A vector of search results ordered by similarity (highest first).
    async fn search(
        &self,
        query_vector: &[f32],
        agent_id: Option<&str>,
        limit: usize,
        threshold: Option<f32>,
    ) -> Result<Vec<SearchResult>, String>;

    /// Deletes all vectors for an agent.
    async fn clear_agent_vectors(&self, agent_id: &str) -> Result<(), String>;

    /// Returns the number of vectors stored.
    async fn count(&self) -> Result<usize, String>;

    /// Returns the dimension of vectors in this store.
    fn dimension(&self) -> usize;
}

// ============================================
// Vector Memory Trait
// ============================================

/// High-level interface for semantic memory operations.
///
/// This trait combines embedding generation and vector storage to provide
/// a complete semantic memory system. It handles the conversion of text
/// memories into vector embeddings and provides semantic search capabilities.
#[async_trait]
pub trait VectorMemory: Send + Sync {
    /// Stores a text memory by generating its embedding and storing it.
    ///
    /// # Arguments
    ///
    /// * `memory_id` - ID of the associated memory entry
    /// * `agent_id` - ID of the agent
    /// * `text` - The text content to embed and store
    /// * `metadata` - Optional metadata to associate with the vector
    async fn store_text(
        &self,
        memory_id: String,
        agent_id: String,
        text: String,
        metadata: Option<serde_json::Value>,
    ) -> Result<String, String>;

    /// Searches for similar memories using semantic search.
    ///
    /// # Arguments
    ///
    /// * `agent_id` - ID of the agent to search within
    /// * `query` - The search query text
    /// * `limit` - Maximum number of results
    /// * `threshold` - Optional minimum similarity threshold
    async fn semantic_search(
        &self,
        agent_id: &str,
        query: &str,
        limit: usize,
        threshold: Option<f32>,
    ) -> Result<Vec<SearchResult>, String>;

    /// Clears all vector memories for an agent.
    async fn clear(&self, agent_id: &str) -> Result<(), String>;
}

// ============================================
// Module Declarations
// ============================================

mod embedding;
mod local_store;
mod utils;

#[cfg(feature = "vector-openai")]
mod openai;

#[cfg(feature = "vector-ollama")]
mod ollama;

#[cfg(feature = "vector-huggingface")]
mod huggingface;

#[cfg(feature = "vector-huggingface-local")]
mod huggingface_local;

#[cfg(test)]
mod tests;

// ============================================
// Public Exports
// ============================================

pub use embedding::CachedEmbeddings;
pub use local_store::LocalVectorStore;
pub use utils::{cosine_similarity, normalize_vector};

#[cfg(feature = "vector-openai")]
pub use openai::{OpenAIEmbeddings, OpenAIModel};

#[cfg(feature = "vector-ollama")]
pub use ollama::{get_model_dimension, OllamaEmbeddings};

#[cfg(feature = "vector-huggingface")]
pub use huggingface::{HuggingFaceEmbeddings, HuggingFaceModel};

#[cfg(feature = "vector-huggingface-local")]
pub use huggingface_local::HuggingFaceLocalEmbeddings;