ceylon_next/memory/vector/

embedding.rs

//! Embedding provider implementations and utilities.

use super::EmbeddingProvider;
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// A caching wrapper for embedding providers.
///
/// This wrapper adds a caching layer on top of any embedding provider to avoid
/// regenerating embeddings for the same text. This is especially useful when:
/// - The same queries or texts are embedded multiple times
/// - Embeddings are expensive to compute (API calls, local models)
/// - You want to reduce API costs and latency
///
/// # Cache Strategy
///
/// - Uses a simple in-memory HashMap with RwLock for thread-safety
/// - Cache key is the exact text string
/// - No cache eviction (unbounded cache - use with caution for large datasets)
/// - Thread-safe for concurrent access
///
/// # Example
///
/// ```rust,no_run
/// use ceylon_next::memory::vector::{EmbeddingProvider, CachedEmbeddings};
/// use std::sync::Arc;
///
/// #[tokio::main]
/// async fn main() {
///     // Wrap any embedding provider with caching
///     // let base_provider = OpenAIEmbeddings::new("api-key");
///     // let cached = CachedEmbeddings::new(Arc::new(base_provider));
///
///     // First call - computes embedding
///     // let embedding1 = cached.embed("hello world").await.unwrap();
///
///     // Second call - returns cached result
///     // let embedding2 = cached.embed("hello world").await.unwrap();
/// }
/// ```
pub struct CachedEmbeddings {
    /// The underlying embedding provider
    provider: Arc<dyn EmbeddingProvider>,
    /// Cache mapping text to embeddings
    cache: Arc<RwLock<HashMap<String, Vec<f32>>>>,
}

impl CachedEmbeddings {
    /// Creates a new cached embedding provider.
    ///
    /// # Arguments
    ///
    /// * `provider` - The underlying embedding provider to wrap
    pub fn new(provider: Arc<dyn EmbeddingProvider>) -> Self {
        Self {
            provider,
            cache: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Creates a new cached embedding provider with pre-allocated cache capacity.
    ///
    /// # Arguments
    ///
    /// * `provider` - The underlying embedding provider to wrap
    /// * `capacity` - Initial capacity for the cache
    pub fn with_capacity(provider: Arc<dyn EmbeddingProvider>, capacity: usize) -> Self {
        Self {
            provider,
            cache: Arc::new(RwLock::new(HashMap::with_capacity(capacity))),
        }
    }

    /// Returns the number of cached embeddings.
    pub async fn cache_size(&self) -> usize {
        self.cache.read().await.len()
    }

    /// Clears all cached embeddings.
    pub async fn clear_cache(&self) {
        self.cache.write().await.clear();
    }

    /// Checks if a text is in the cache.
    pub async fn is_cached(&self, text: &str) -> bool {
        self.cache.read().await.contains_key(text)
    }

    /// Preloads embeddings into the cache.
    ///
    /// This can be useful for warming up the cache with known queries.
    ///
    /// # Arguments
    ///
    /// * `texts` - Texts to preload into cache
    pub async fn preload(&self, texts: &[String]) -> Result<(), String> {
        let embeddings = self.provider.embed_batch(texts).await?;

        let mut cache = self.cache.write().await;
        for (text, embedding) in texts.iter().zip(embeddings.iter()) {
            cache.insert(text.clone(), embedding.clone());
        }

        Ok(())
    }
}

#[async_trait]
impl EmbeddingProvider for CachedEmbeddings {
    async fn embed(&self, text: &str) -> Result<Vec<f32>, String> {
        // Check cache first
        {
            let cache = self.cache.read().await;
            if let Some(embedding) = cache.get(text) {
                return Ok(embedding.clone());
            }
        }

        // Not in cache, compute embedding
        let embedding = self.provider.embed(text).await?;

        // Store in cache
        {
            let mut cache = self.cache.write().await;
            cache.insert(text.to_string(), embedding.clone());
        }

        Ok(embedding)
    }

    async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, String> {
        let mut results = Vec::with_capacity(texts.len());
        let mut uncached_indices = Vec::new();
        let mut uncached_texts = Vec::new();

        // Check which texts are cached
        {
            let cache = self.cache.read().await;
            for (i, text) in texts.iter().enumerate() {
                if let Some(embedding) = cache.get(text) {
                    results.push((i, embedding.clone()));
                } else {
                    uncached_indices.push(i);
                    uncached_texts.push(text.clone());
                }
            }
        }

        // Compute embeddings for uncached texts
        if !uncached_texts.is_empty() {
            let new_embeddings = self.provider.embed_batch(&uncached_texts).await?;

            // Store new embeddings in cache and results
            {
                let mut cache = self.cache.write().await;
                for (text, embedding) in uncached_texts.iter().zip(new_embeddings.iter()) {
                    cache.insert(text.clone(), embedding.clone());
                }
            }

            for (i, embedding) in uncached_indices.iter().zip(new_embeddings.iter()) {
                results.push((*i, embedding.clone()));
            }
        }

        // Sort by original index and extract embeddings
        results.sort_by_key(|(i, _)| *i);
        Ok(results.into_iter().map(|(_, emb)| emb).collect())
    }

    fn dimension(&self) -> usize {
        self.provider.dimension()
    }

    fn model_name(&self) -> &str {
        self.provider.model_name()
    }
}

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

    // Mock embedding provider for testing
    struct MockEmbedder {
        dimension: usize,
        call_count: Arc<RwLock<usize>>,
    }

    impl MockEmbedder {
        fn new(dimension: usize) -> Self {
            Self {
                dimension,
                call_count: Arc::new(RwLock::new(0)),
            }
        }

        async fn calls(&self) -> usize {
            *self.call_count.read().await
        }
    }

    #[async_trait]
    impl EmbeddingProvider for MockEmbedder {
        async fn embed(&self, text: &str) -> Result<Vec<f32>, String> {
            *self.call_count.write().await += 1;

            // Simple hash-based embedding for testing
            let hash = text.len() as f32;
            Ok(vec![hash; self.dimension])
        }

        fn dimension(&self) -> usize {
            self.dimension
        }

        fn model_name(&self) -> &str {
            "mock"
        }
    }

    #[tokio::test]
    async fn test_caching() {
        let mock = Arc::new(MockEmbedder::new(3));
        let cached = CachedEmbeddings::new(mock.clone());

        // First call should compute
        let emb1 = cached.embed("test").await.unwrap();
        assert_eq!(mock.calls().await, 1);
        assert_eq!(emb1.len(), 3);

        // Second call should use cache
        let emb2 = cached.embed("test").await.unwrap();
        assert_eq!(mock.calls().await, 1); // No new call
        assert_eq!(emb1, emb2);

        // Different text should compute again
        let _emb3 = cached.embed("other").await.unwrap();
        assert_eq!(mock.calls().await, 2);
    }

    #[tokio::test]
    async fn test_batch_caching() {
        let mock = Arc::new(MockEmbedder::new(2));
        let cached = CachedEmbeddings::new(mock.clone());

        let texts1 = vec!["a".to_string(), "b".to_string(), "c".to_string()];
        let embs1 = cached.embed_batch(&texts1).await.unwrap();
        assert_eq!(embs1.len(), 3);

        // All texts should be cached
        assert_eq!(cached.cache_size().await, 3);

        // Reusing some texts
        let texts2 = vec!["a".to_string(), "b".to_string(), "d".to_string()];
        let embs2 = cached.embed_batch(&texts2).await.unwrap();

        // Should only compute for "d"
        assert_eq!(cached.cache_size().await, 4);

        // Cached texts should return same embeddings
        assert_eq!(embs1[0], embs2[0]); // "a"
        assert_eq!(embs1[1], embs2[1]); // "b"
    }

    #[tokio::test]
    async fn test_cache_operations() {
        let mock = Arc::new(MockEmbedder::new(2));
        let cached = CachedEmbeddings::new(mock);

        assert_eq!(cached.cache_size().await, 0);
        assert!(!cached.is_cached("test").await);

        cached.embed("test").await.unwrap();

        assert_eq!(cached.cache_size().await, 1);
        assert!(cached.is_cached("test").await);

        cached.clear_cache().await;

        assert_eq!(cached.cache_size().await, 0);
        assert!(!cached.is_cached("test").await);
    }

    #[tokio::test]
    async fn test_preload() {
        let mock = Arc::new(MockEmbedder::new(2));
        let cached = CachedEmbeddings::new(mock.clone());

        let texts = vec!["a".to_string(), "b".to_string(), "c".to_string()];
        cached.preload(&texts).await.unwrap();

        assert_eq!(cached.cache_size().await, 3);
        assert!(cached.is_cached("a").await);
        assert!(cached.is_cached("b").await);
        assert!(cached.is_cached("c").await);

        let initial_calls = mock.calls().await;

        // These should all hit cache
        cached.embed("a").await.unwrap();
        cached.embed("b").await.unwrap();
        cached.embed("c").await.unwrap();

        assert_eq!(mock.calls().await, initial_calls);
    }
}