Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'

vendor/ruvector/examples/refrag-pipeline/src/store.rs

@@ -0,0 +1,581 @@
+//! REFRAG Store - Unified storage layer with hybrid search
+//!
+//! This module integrates the Compress, Sense, and Expand layers
+//! into a cohesive REFRAG-enabled vector store.
+
+use crate::compress::{BatchCompressor, CompressionStrategy, TensorCompressor};
+use crate::expand::{ExpandLayer, ProjectorRegistry};
+use crate::sense::{PolicyDecision, PolicyNetwork, RefragAction};
+use crate::types::{RefragConfig, RefragEntry, RefragSearchResult, RefragStats};
+
+use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
+use ruvector_core::{SearchQuery, SearchResult, VectorEntry};
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, RwLock};
+use std::time::Instant;
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum StoreError {
+    #[error("Entry not found: {0}")]
+    NotFound(String),
+
+    #[error("Dimension mismatch: expected {expected}, got {actual}")]
+    DimensionMismatch { expected: usize, actual: usize },
+
+    #[error("Compression error: {0}")]
+    CompressionError(String),
+
+    #[error("Policy error: {0}")]
+    PolicyError(String),
+
+    #[error("Projection error: {0}")]
+    ProjectionError(String),
+
+    #[error("Core error: {0}")]
+    CoreError(String),
+}
+
+pub type Result<T> = std::result::Result<T, StoreError>;
+
+/// REFRAG-enabled vector store
+///
+/// Wraps ruvector-core with REFRAG capabilities:
+/// - Stores both search vectors and representation tensors
+/// - Uses policy network to decide COMPRESS vs EXPAND
+/// - Projects tensors to target LLM dimensions
+pub struct RefragStore {
+    /// Configuration
+    config: RefragConfig,
+    /// Stored entries (in-memory for this example)
+    entries: RwLock<HashMap<String, RefragEntry>>,
+    /// Tensor compressor
+    compressor: TensorCompressor,
+    /// Policy network
+    policy: PolicyNetwork,
+    /// Expand layer
+    expand: ExpandLayer,
+    /// Statistics
+    stats: RefragStoreStats,
+}
+
+/// Thread-safe statistics
+struct RefragStoreStats {
+    total_searches: AtomicU64,
+    expand_count: AtomicU64,
+    compress_count: AtomicU64,
+    total_policy_time_us: AtomicU64,
+    total_projection_time_us: AtomicU64,
+}
+
+impl RefragStoreStats {
+    fn new() -> Self {
+        Self {
+            total_searches: AtomicU64::new(0),
+            expand_count: AtomicU64::new(0),
+            compress_count: AtomicU64::new(0),
+            total_policy_time_us: AtomicU64::new(0),
+            total_projection_time_us: AtomicU64::new(0),
+        }
+    }
+
+    fn to_stats(&self) -> RefragStats {
+        let total = self.total_searches.load(Ordering::Relaxed);
+        RefragStats {
+            total_searches: total,
+            expand_count: self.expand_count.load(Ordering::Relaxed),
+            compress_count: self.compress_count.load(Ordering::Relaxed),
+            avg_policy_time_us: if total > 0 {
+                self.total_policy_time_us.load(Ordering::Relaxed) as f64 / total as f64
+            } else {
+                0.0
+            },
+            avg_projection_time_us: if total > 0 {
+                self.total_projection_time_us.load(Ordering::Relaxed) as f64 / total as f64
+            } else {
+                0.0
+            },
+            bytes_saved: 0, // Would need per-entry tracking
+        }
+    }
+}
+
+impl RefragStore {
+    /// Create a new REFRAG store with default configuration
+    pub fn new(search_dim: usize, tensor_dim: usize) -> Result<Self> {
+        let config = RefragConfig {
+            search_dimensions: search_dim,
+            tensor_dimensions: tensor_dim,
+            ..Default::default()
+        };
+
+        Self::with_config(config)
+    }
+
+    /// Create with custom configuration
+    pub fn with_config(config: RefragConfig) -> Result<Self> {
+        let compressor = TensorCompressor::new(config.tensor_dimensions)
+            .with_strategy(CompressionStrategy::None);
+
+        let policy = PolicyNetwork::threshold(config.compress_threshold);
+
+        let expand = ExpandLayer::new(
+            ProjectorRegistry::with_defaults(config.tensor_dimensions),
+            "llama3-8b",
+        );
+
+        Ok(Self {
+            config,
+            entries: RwLock::new(HashMap::new()),
+            compressor,
+            policy,
+            expand,
+            stats: RefragStoreStats::new(),
+        })
+    }
+
+    /// Set custom policy network
+    pub fn with_policy(mut self, policy: PolicyNetwork) -> Self {
+        self.policy = policy;
+        self
+    }
+
+    /// Set custom expand layer
+    pub fn with_expand(mut self, expand: ExpandLayer) -> Self {
+        self.expand = expand;
+        self
+    }
+
+    /// Insert a REFRAG entry
+    pub fn insert(&self, entry: RefragEntry) -> Result<String> {
+        if entry.search_vector.len() != self.config.search_dimensions {
+            return Err(StoreError::DimensionMismatch {
+                expected: self.config.search_dimensions,
+                actual: entry.search_vector.len(),
+            });
+        }
+
+        let id = entry.id.clone();
+        self.entries.write().unwrap().insert(id.clone(), entry);
+        Ok(id)
+    }
+
+    /// Insert with automatic tensor compression
+    pub fn insert_with_tensor(
+        &self,
+        id: impl Into<String>,
+        search_vector: Vec<f32>,
+        representation_vector: Vec<f32>,
+        text: impl Into<String>,
+        model_id: impl Into<String>,
+    ) -> Result<String> {
+        // Compress the representation tensor
+        let tensor = self
+            .compressor
+            .compress(&representation_vector)
+            .map_err(|e| StoreError::CompressionError(e.to_string()))?;
+
+        let entry = RefragEntry::new(id, search_vector, text).with_tensor(tensor, model_id);
+
+        self.insert(entry)
+    }
+
+    /// Batch insert
+    pub fn insert_batch(&self, entries: Vec<RefragEntry>) -> Result<Vec<String>> {
+        let mut ids = Vec::with_capacity(entries.len());
+        for entry in entries {
+            ids.push(self.insert(entry)?);
+        }
+        Ok(ids)
+    }
+
+    /// Get entry by ID
+    pub fn get(&self, id: &str) -> Result<RefragEntry> {
+        self.entries
+            .read()
+            .unwrap()
+            .get(id)
+            .cloned()
+            .ok_or_else(|| StoreError::NotFound(id.to_string()))
+    }
+
+    /// Delete entry
+    pub fn delete(&self, id: &str) -> Result<bool> {
+        Ok(self.entries.write().unwrap().remove(id).is_some())
+    }
+
+    /// Standard vector search (returns text only)
+    pub fn search(&self, query: &[f32], k: usize) -> Result<Vec<RefragSearchResult>> {
+        self.search_with_options(query, k, None, false)
+    }
+
+    /// Hybrid search with REFRAG policy decisions
+    ///
+    /// Returns mixed COMPRESS/EXPAND results based on policy network decisions.
+    pub fn search_hybrid(
+        &self,
+        query: &[f32],
+        k: usize,
+        threshold: Option<f32>,
+    ) -> Result<Vec<RefragSearchResult>> {
+        self.search_with_options(query, k, threshold, true)
+    }
+
+    /// Full-featured search
+    fn search_with_options(
+        &self,
+        query: &[f32],
+        k: usize,
+        threshold: Option<f32>,
+        use_policy: bool,
+    ) -> Result<Vec<RefragSearchResult>> {
+        if query.len() != self.config.search_dimensions {
+            return Err(StoreError::DimensionMismatch {
+                expected: self.config.search_dimensions,
+                actual: query.len(),
+            });
+        }
+
+        let entries = self.entries.read().unwrap();
+
+        // Compute similarities (brute force for this example)
+        let mut scored: Vec<(&RefragEntry, f32)> = entries
+            .values()
+            .map(|entry| {
+                let similarity = cosine_similarity(query, &entry.search_vector);
+                (entry, similarity)
+            })
+            .collect();
+
+        // Sort by score descending
+        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
+
+        // Apply threshold filter
+        let threshold_val = threshold.unwrap_or(0.0);
+        let filtered: Vec<_> = scored
+            .into_iter()
+            .filter(|(_, score)| *score >= threshold_val)
+            .take(k)
+            .collect();
+
+        // Process results with policy
+        let mut results = Vec::with_capacity(filtered.len());
+
+        for (entry, score) in filtered {
+            self.stats.total_searches.fetch_add(1, Ordering::Relaxed);
+
+            let result = if use_policy && entry.has_tensor() {
+                self.process_with_policy(entry, query, score)?
+            } else {
+                // Default to EXPAND (text)
+                self.stats.expand_count.fetch_add(1, Ordering::Relaxed);
+                RefragSearchResult::expand(entry.id.clone(), score, entry.text_content.clone(), 1.0)
+            };
+
+            results.push(result);
+        }
+
+        Ok(results)
+    }
+
+    /// Process a single result through the REFRAG policy
+    fn process_with_policy(
+        &self,
+        entry: &RefragEntry,
+        query: &[f32],
+        score: f32,
+    ) -> Result<RefragSearchResult> {
+        let tensor_bytes = entry.representation_tensor.as_ref().unwrap();
+
+        // Decompress tensor for policy evaluation
+        let tensor = self
+            .compressor
+            .decompress(tensor_bytes)
+            .map_err(|e| StoreError::CompressionError(e.to_string()))?;
+
+        // Run policy
+        let start = Instant::now();
+        let decision = self
+            .policy
+            .decide(&tensor, query)
+            .map_err(|e| StoreError::PolicyError(e.to_string()))?;
+        let policy_time = start.elapsed().as_micros() as u64;
+        self.stats
+            .total_policy_time_us
+            .fetch_add(policy_time, Ordering::Relaxed);
+
+        match decision.action {
+            RefragAction::Compress => {
+                self.stats.compress_count.fetch_add(1, Ordering::Relaxed);
+
+                // Optionally project to target LLM dimensions
+                let (final_tensor, projection_time) = if self.config.auto_project {
+                    let model_id = entry.alignment_model_id.as_deref();
+                    let start = Instant::now();
+                    let projected = self
+                        .expand
+                        .expand_auto(&tensor, model_id)
+                        .map_err(|e| StoreError::ProjectionError(e.to_string()))?;
+                    let time = start.elapsed().as_micros() as u64;
+                    (projected, time)
+                } else {
+                    (tensor, 0)
+                };
+
+                self.stats
+                    .total_projection_time_us
+                    .fetch_add(projection_time, Ordering::Relaxed);
+
+                // Encode tensor as base64
+                let tensor_bytes: Vec<u8> =
+                    final_tensor.iter().flat_map(|f| f.to_le_bytes()).collect();
+                let tensor_b64 = BASE64.encode(&tensor_bytes);
+
+                Ok(RefragSearchResult::compress(
+                    entry.id.clone(),
+                    score,
+                    tensor_b64,
+                    final_tensor.len(),
+                    entry.alignment_model_id.clone(),
+                    decision.confidence,
+                ))
+            }
+            RefragAction::Expand => {
+                self.stats.expand_count.fetch_add(1, Ordering::Relaxed);
+
+                Ok(RefragSearchResult::expand(
+                    entry.id.clone(),
+                    score,
+                    entry.text_content.clone(),
+                    decision.confidence,
+                ))
+            }
+        }
+    }
+
+    /// Get store statistics
+    pub fn stats(&self) -> RefragStats {
+        self.stats.to_stats()
+    }
+
+    /// Get entry count
+    pub fn len(&self) -> usize {
+        self.entries.read().unwrap().len()
+    }
+
+    /// Check if empty
+    pub fn is_empty(&self) -> bool {
+        self.entries.read().unwrap().is_empty()
+    }
+
+    /// Get configuration
+    pub fn config(&self) -> &RefragConfig {
+        &self.config
+    }
+}
+
+/// Cosine similarity helper
+fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
+    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
+    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
+    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
+
+    if norm_a > f32::EPSILON && norm_b > f32::EPSILON {
+        dot / (norm_a * norm_b)
+    } else {
+        0.0
+    }
+}
+
+/// Builder for RefragStore
+pub struct RefragStoreBuilder {
+    config: RefragConfig,
+    policy: Option<PolicyNetwork>,
+    expand: Option<ExpandLayer>,
+    compression: CompressionStrategy,
+}
+
+impl RefragStoreBuilder {
+    pub fn new() -> Self {
+        Self {
+            config: RefragConfig::default(),
+            policy: None,
+            expand: None,
+            compression: CompressionStrategy::None,
+        }
+    }
+
+    pub fn search_dimensions(mut self, dim: usize) -> Self {
+        self.config.search_dimensions = dim;
+        self
+    }
+
+    pub fn tensor_dimensions(mut self, dim: usize) -> Self {
+        self.config.tensor_dimensions = dim;
+        self
+    }
+
+    pub fn target_dimensions(mut self, dim: usize) -> Self {
+        self.config.target_dimensions = dim;
+        self
+    }
+
+    pub fn compress_threshold(mut self, threshold: f32) -> Self {
+        self.config.compress_threshold = threshold;
+        self
+    }
+
+    pub fn auto_project(mut self, enabled: bool) -> Self {
+        self.config.auto_project = enabled;
+        self
+    }
+
+    pub fn policy(mut self, policy: PolicyNetwork) -> Self {
+        self.policy = Some(policy);
+        self
+    }
+
+    pub fn expand_layer(mut self, expand: ExpandLayer) -> Self {
+        self.expand = Some(expand);
+        self
+    }
+
+    pub fn compression(mut self, strategy: CompressionStrategy) -> Self {
+        self.compression = strategy;
+        self
+    }
+
+    pub fn build(self) -> Result<RefragStore> {
+        let mut store = RefragStore::with_config(self.config)?;
+
+        if let Some(policy) = self.policy {
+            store = store.with_policy(policy);
+        }
+
+        if let Some(expand) = self.expand {
+            store = store.with_expand(expand);
+        }
+
+        Ok(store)
+    }
+}
+
+impl Default for RefragStoreBuilder {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::types::RefragResponseType;
+
+    fn create_test_entry(id: &str, dim: usize) -> RefragEntry {
+        let search_vec: Vec<f32> = (0..dim).map(|i| (i as f32) / (dim as f32)).collect();
+        let tensor_vec: Vec<f32> = (0..768).map(|i| (i as f32) / 768.0).collect();
+        let tensor_bytes: Vec<u8> = tensor_vec.iter().flat_map(|f| f.to_le_bytes()).collect();
+
+        RefragEntry::new(id, search_vec, format!("Text content for {}", id))
+            .with_tensor(tensor_bytes, "llama3-8b")
+    }
+
+    #[test]
+    fn test_store_creation() {
+        let store = RefragStore::new(384, 768).unwrap();
+        assert_eq!(store.config().search_dimensions, 384);
+        assert_eq!(store.config().tensor_dimensions, 768);
+        assert!(store.is_empty());
+    }
+
+    #[test]
+    fn test_insert_and_get() {
+        let store = RefragStore::new(4, 768).unwrap();
+        let entry = create_test_entry("doc_1", 4);
+
+        let id = store.insert(entry.clone()).unwrap();
+        assert_eq!(id, "doc_1");
+        assert_eq!(store.len(), 1);
+
+        let retrieved = store.get("doc_1").unwrap();
+        assert_eq!(retrieved.id, "doc_1");
+        assert!(retrieved.has_tensor());
+    }
+
+    #[test]
+    fn test_standard_search() {
+        let store = RefragStore::new(4, 768).unwrap();
+
+        // Insert test entries
+        for i in 0..5 {
+            store
+                .insert(create_test_entry(&format!("doc_{}", i), 4))
+                .unwrap();
+        }
+
+        let query: Vec<f32> = (0..4).map(|i| (i as f32) / 4.0).collect();
+        let results = store.search(&query, 3).unwrap();
+
+        assert_eq!(results.len(), 3);
+        // All should be EXPAND since we used standard search
+        for result in &results {
+            assert_eq!(result.response_type, RefragResponseType::Expand);
+            assert!(result.content.is_some());
+        }
+    }
+
+    #[test]
+    fn test_hybrid_search() {
+        // Use lower threshold to get COMPRESS results
+        let store = RefragStoreBuilder::new()
+            .search_dimensions(4)
+            .tensor_dimensions(768)
+            .compress_threshold(0.5)
+            .build()
+            .unwrap();
+
+        for i in 0..5 {
+            store
+                .insert(create_test_entry(&format!("doc_{}", i), 4))
+                .unwrap();
+        }
+
+        let query: Vec<f32> = (0..4).map(|i| (i as f32) / 4.0).collect();
+        let results = store.search_hybrid(&query, 3, None).unwrap();
+
+        assert_eq!(results.len(), 3);
+
+        // Check that we got some policy decisions
+        let stats = store.stats();
+        assert!(stats.total_searches > 0);
+    }
+
+    #[test]
+    fn test_statistics() {
+        let store = RefragStore::new(4, 768).unwrap();
+
+        for i in 0..3 {
+            store
+                .insert(create_test_entry(&format!("doc_{}", i), 4))
+                .unwrap();
+        }
+
+        let query: Vec<f32> = (0..4).map(|i| (i as f32) / 4.0).collect();
+        let _ = store.search_hybrid(&query, 3, None).unwrap();
+
+        let stats = store.stats();
+        assert_eq!(stats.total_searches, 3);
+        assert_eq!(stats.expand_count + stats.compress_count, 3);
+    }
+
+    #[test]
+    fn test_dimension_mismatch() {
+        let store = RefragStore::new(4, 768).unwrap();
+
+        let bad_entry = RefragEntry::new("bad", vec![1.0, 2.0, 3.0], "text"); // Only 3 dims
+        let result = store.insert(bad_entry);
+
+        assert!(matches!(result, Err(StoreError::DimensionMismatch { .. })));
+    }
+}