ceylon_next/memory/vector/

local_store.rs

//! In-memory vector storage with cosine similarity search.

use super::{SearchResult, VectorEntry, VectorStore};
use crate::memory::vector::utils::cosine_similarity;
use async_trait::async_trait;
use std::sync::Arc;
use tokio::sync::RwLock;

/// In-memory vector store using brute-force cosine similarity search.
///
/// This store keeps all vectors in memory and performs exact nearest neighbor
/// search using cosine similarity. While not as efficient as approximate methods
/// like HNSW for large datasets, it's simple, fast for small to medium datasets,
/// and provides exact results.
///
/// # Features
///
/// - Fast in-memory storage with RwLock for thread-safety
/// - Exact nearest neighbor search using cosine similarity
/// - Optional similarity threshold filtering
/// - Agent-specific filtering
/// - No external dependencies
///
/// # Example
///
/// ```rust
/// use ceylon_next::memory::vector::{LocalVectorStore, VectorEntry, VectorStore};
/// use std::sync::Arc;
///
/// #[tokio::main]
/// async fn main() {
///     let store = LocalVectorStore::new(384);
///
///     let entry = VectorEntry::new(
///         "memory-1".to_string(),
///         "agent-1".to_string(),
///         "Hello world".to_string(),
///         vec![0.1; 384],
///         None,
///     );
///
///     let id = store.store(entry).await.unwrap();
///     println!("Stored vector with ID: {}", id);
///
///     // Search for similar vectors
///     let query = vec![0.1; 384];
///     let results = store.search(&query, Some("agent-1"), 5, None).await.unwrap();
///     println!("Found {} similar vectors", results.len());
/// }
/// ```
pub struct LocalVectorStore {
    /// The dimensionality of vectors in this store
    dimension: usize,
    /// The stored vectors
    vectors: Arc<RwLock<Vec<VectorEntry>>>,
}

impl LocalVectorStore {
    /// Creates a new local vector store.
    ///
    /// # Arguments
    ///
    /// * `dimension` - The dimensionality of vectors to store
    ///
    /// # Example
    ///
    /// ```rust
    /// use ceylon_next::memory::vector::LocalVectorStore;
    ///
    /// let store = LocalVectorStore::new(384);
    /// ```
    pub fn new(dimension: usize) -> Self {
        Self {
            dimension,
            vectors: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// Creates a new local vector store with pre-allocated capacity.
    ///
    /// # Arguments
    ///
    /// * `dimension` - The dimensionality of vectors to store
    /// * `capacity` - Initial capacity for the vector storage
    pub fn with_capacity(dimension: usize, capacity: usize) -> Self {
        Self {
            dimension,
            vectors: Arc::new(RwLock::new(Vec::with_capacity(capacity))),
        }
    }
}

#[async_trait]
impl VectorStore for LocalVectorStore {
    async fn store(&self, entry: VectorEntry) -> Result<String, String> {
        // Validate vector dimension
        if entry.vector.len() != self.dimension {
            return Err(format!(
                "Vector dimension mismatch: expected {}, got {}",
                self.dimension,
                entry.vector.len()
            ));
        }

        let id = entry.id.clone();
        let mut vectors = self.vectors.write().await;
        vectors.push(entry);

        Ok(id)
    }

    async fn store_batch(&self, entries: Vec<VectorEntry>) -> Result<Vec<String>, String> {
        // Validate all vectors first
        for entry in &entries {
            if entry.vector.len() != self.dimension {
                return Err(format!(
                    "Vector dimension mismatch: expected {}, got {}",
                    self.dimension,
                    entry.vector.len()
                ));
            }
        }

        let ids: Vec<String> = entries.iter().map(|e| e.id.clone()).collect();
        let mut vectors = self.vectors.write().await;
        vectors.extend(entries);

        Ok(ids)
    }

    async fn get(&self, id: &str) -> Result<Option<VectorEntry>, String> {
        let vectors = self.vectors.read().await;
        Ok(vectors.iter().find(|v| v.id == id).cloned())
    }

    async fn search(
        &self,
        query_vector: &[f32],
        agent_id: Option<&str>,
        limit: usize,
        threshold: Option<f32>,
    ) -> Result<Vec<SearchResult>, String> {
        // Validate query vector dimension
        if query_vector.len() != self.dimension {
            return Err(format!(
                "Query vector dimension mismatch: expected {}, got {}",
                self.dimension,
                query_vector.len()
            ));
        }

        let vectors = self.vectors.read().await;

        // Filter by agent if specified
        let filtered: Vec<&VectorEntry> = if let Some(aid) = agent_id {
            vectors.iter().filter(|v| v.agent_id == aid).collect()
        } else {
            vectors.iter().collect()
        };

        // Compute similarities for all vectors
        let mut results: Vec<SearchResult> = Vec::new();

        for entry in filtered {
            match cosine_similarity(query_vector, &entry.vector) {
                Ok(score) => {
                    // Apply threshold if specified
                    if let Some(min_score) = threshold {
                        if score < min_score {
                            continue;
                        }
                    }

                    results.push(SearchResult {
                        entry: entry.clone(),
                        score,
                    });
                }
                Err(e) => {
                    log::warn!("Failed to compute similarity for vector {}: {}", entry.id, e);
                    continue;
                }
            }
        }

        // Sort by score descending (highest similarity first)
        results.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Limit results
        results.truncate(limit);

        Ok(results)
    }

    async fn clear_agent_vectors(&self, agent_id: &str) -> Result<(), String> {
        let mut vectors = self.vectors.write().await;
        vectors.retain(|v| v.agent_id != agent_id);
        Ok(())
    }

    async fn count(&self) -> Result<usize, String> {
        let vectors = self.vectors.read().await;
        Ok(vectors.len())
    }

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

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

    #[tokio::test]
    async fn test_store_and_get() {
        let store = LocalVectorStore::new(3);
        let entry = VectorEntry::new(
            "mem-1".to_string(),
            "agent-1".to_string(),
            "test".to_string(),
            vec![1.0, 2.0, 3.0],
            None,
        );

        let id = entry.id.clone();
        store.store(entry).await.unwrap();

        let retrieved = store.get(&id).await.unwrap();
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().text, "test");
    }

    #[tokio::test]
    async fn test_dimension_validation() {
        let store = LocalVectorStore::new(3);
        let entry = VectorEntry::new(
            "mem-1".to_string(),
            "agent-1".to_string(),
            "test".to_string(),
            vec![1.0, 2.0], // Wrong dimension
            None,
        );

        let result = store.store(entry).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_search() {
        let store = LocalVectorStore::new(3);

        // Store some vectors
        let entries = vec![
            VectorEntry::new(
                "mem-1".to_string(),
                "agent-1".to_string(),
                "cat".to_string(),
                vec![1.0, 0.0, 0.0],
                None,
            ),
            VectorEntry::new(
                "mem-2".to_string(),
                "agent-1".to_string(),
                "dog".to_string(),
                vec![0.9, 0.1, 0.0],
                None,
            ),
            VectorEntry::new(
                "mem-3".to_string(),
                "agent-1".to_string(),
                "car".to_string(),
                vec![0.0, 1.0, 0.0],
                None,
            ),
        ];

        for entry in entries {
            store.store(entry).await.unwrap();
        }

        // Search for similar to "cat"
        let query = vec![1.0, 0.0, 0.0];
        let results = store.search(&query, Some("agent-1"), 2, None).await.unwrap();

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].entry.text, "cat"); // Most similar
        assert_eq!(results[1].entry.text, "dog"); // Second most similar
        assert!(results[0].score > results[1].score);
    }

    #[tokio::test]
    async fn test_search_with_threshold() {
        let store = LocalVectorStore::new(2);

        store
            .store(VectorEntry::new(
                "mem-1".to_string(),
                "agent-1".to_string(),
                "similar".to_string(),
                vec![1.0, 0.0],
                None,
            ))
            .await
            .unwrap();

        store
            .store(VectorEntry::new(
                "mem-2".to_string(),
                "agent-1".to_string(),
                "different".to_string(),
                vec![0.0, 1.0],
                None,
            ))
            .await
            .unwrap();

        let query = vec![1.0, 0.0];
        let results = store
            .search(&query, Some("agent-1"), 10, Some(0.5))
            .await
            .unwrap();

        // Should only return the similar vector
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].entry.text, "similar");
    }

    #[tokio::test]
    async fn test_agent_filtering() {
        let store = LocalVectorStore::new(2);

        store
            .store(VectorEntry::new(
                "mem-1".to_string(),
                "agent-1".to_string(),
                "agent1".to_string(),
                vec![1.0, 0.0],
                None,
            ))
            .await
            .unwrap();

        store
            .store(VectorEntry::new(
                "mem-2".to_string(),
                "agent-2".to_string(),
                "agent2".to_string(),
                vec![1.0, 0.0],
                None,
            ))
            .await
            .unwrap();

        let query = vec![1.0, 0.0];
        let results = store.search(&query, Some("agent-1"), 10, None).await.unwrap();

        assert_eq!(results.len(), 1);
        assert_eq!(results[0].entry.agent_id, "agent-1");
    }

    #[tokio::test]
    async fn test_clear_agent_vectors() {
        let store = LocalVectorStore::new(2);

        store
            .store(VectorEntry::new(
                "mem-1".to_string(),
                "agent-1".to_string(),
                "test".to_string(),
                vec![1.0, 0.0],
                None,
            ))
            .await
            .unwrap();

        store
            .store(VectorEntry::new(
                "mem-2".to_string(),
                "agent-2".to_string(),
                "test".to_string(),
                vec![1.0, 0.0],
                None,
            ))
            .await
            .unwrap();

        assert_eq!(store.count().await.unwrap(), 2);

        store.clear_agent_vectors("agent-1").await.unwrap();

        assert_eq!(store.count().await.unwrap(), 1);

        let query = vec![1.0, 0.0];
        let results = store.search(&query, None, 10, None).await.unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].entry.agent_id, "agent-2");
    }

    #[tokio::test]
    async fn test_store_batch() {
        let store = LocalVectorStore::new(2);

        let entries = vec![
            VectorEntry::new(
                "mem-1".to_string(),
                "agent-1".to_string(),
                "test1".to_string(),
                vec![1.0, 0.0],
                None,
            ),
            VectorEntry::new(
                "mem-2".to_string(),
                "agent-1".to_string(),
                "test2".to_string(),
                vec![0.0, 1.0],
                None,
            ),
        ];

        let ids = store.store_batch(entries).await.unwrap();
        assert_eq!(ids.len(), 2);
        assert_eq!(store.count().await.unwrap(), 2);
    }
}