Skip to main content

ipfrs_semantic/
hybrid.rs

1//! Hybrid search combining vector similarity with metadata filtering
2//!
3//! This module provides a unified search interface that combines
4//! semantic vector search with attribute-based filtering.
5
6use crate::hnsw::{DistanceMetric, SearchResult, VectorIndex};
7use crate::metadata::{Metadata, MetadataFilter, MetadataStore, TemporalOptions};
8use crate::stats::{IndexHealth, IndexStats, MemoryUsage, PerfTimer, StatsSnapshot};
9use ipfrs_core::{Cid, Error, Result};
10use lru::LruCache;
11use serde::{Deserialize, Serialize};
12use std::collections::HashSet;
13use std::num::NonZeroUsize;
14use std::sync::{Arc, RwLock};
15
16/// Hybrid search configuration
17#[derive(Debug, Clone)]
18pub struct HybridConfig {
19    /// Vector dimension
20    pub dimension: usize,
21    /// Distance metric
22    pub metric: DistanceMetric,
23    /// HNSW max connections
24    pub max_connections: usize,
25    /// HNSW ef_construction
26    pub ef_construction: usize,
27    /// Default ef_search
28    pub ef_search: usize,
29    /// Query cache size
30    pub cache_size: usize,
31    /// Enable statistics collection
32    pub collect_stats: bool,
33    /// Filtering strategy
34    pub filter_strategy: FilterStrategy,
35}
36
37impl Default for HybridConfig {
38    fn default() -> Self {
39        Self {
40            dimension: 768,
41            metric: DistanceMetric::Cosine,
42            max_connections: 16,
43            ef_construction: 200,
44            ef_search: 50,
45            cache_size: 1000,
46            collect_stats: true,
47            filter_strategy: FilterStrategy::Auto,
48        }
49    }
50}
51
52/// Strategy for applying filters
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
54pub enum FilterStrategy {
55    /// Automatically choose based on selectivity
56    Auto,
57    /// Filter before vector search (pre-filtering)
58    PreFilter,
59    /// Filter after vector search (post-filtering)
60    PostFilter,
61}
62
63/// Hybrid search query
64#[derive(Debug, Clone)]
65pub struct HybridQuery {
66    /// Query vector
67    pub vector: Vec<f32>,
68    /// Number of results to return
69    pub k: usize,
70    /// Metadata filter (optional)
71    pub filter: Option<MetadataFilter>,
72    /// Temporal options (optional)
73    pub temporal: Option<TemporalOptions>,
74    /// Minimum similarity score
75    pub min_score: Option<f32>,
76    /// Override ef_search parameter
77    pub ef_search: Option<usize>,
78    /// Include metadata in results
79    pub include_metadata: bool,
80}
81
82impl HybridQuery {
83    /// Create a simple k-NN query
84    pub fn knn(vector: Vec<f32>, k: usize) -> Self {
85        Self {
86            vector,
87            k,
88            filter: None,
89            temporal: None,
90            min_score: None,
91            ef_search: None,
92            include_metadata: false,
93        }
94    }
95
96    /// Add a metadata filter
97    pub fn with_filter(mut self, filter: MetadataFilter) -> Self {
98        self.filter = Some(filter);
99        self
100    }
101
102    /// Add temporal options
103    pub fn with_temporal(mut self, temporal: TemporalOptions) -> Self {
104        self.temporal = Some(temporal);
105        self
106    }
107
108    /// Set minimum score threshold
109    pub fn with_min_score(mut self, min_score: f32) -> Self {
110        self.min_score = Some(min_score);
111        self
112    }
113
114    /// Include metadata in results
115    pub fn with_metadata(mut self) -> Self {
116        self.include_metadata = true;
117        self
118    }
119
120    /// Override ef_search parameter
121    pub fn with_ef_search(mut self, ef_search: usize) -> Self {
122        self.ef_search = Some(ef_search);
123        self
124    }
125}
126
127/// Hybrid search result with optional metadata
128#[derive(Debug, Clone)]
129pub struct HybridResult {
130    /// Content identifier
131    pub cid: Cid,
132    /// Similarity score
133    pub score: f32,
134    /// Metadata (if requested)
135    pub metadata: Option<Metadata>,
136}
137
138impl From<SearchResult> for HybridResult {
139    fn from(result: SearchResult) -> Self {
140        Self {
141            cid: result.cid,
142            score: result.score,
143            metadata: None,
144        }
145    }
146}
147
148/// Hybrid search response
149#[derive(Debug, Clone)]
150pub struct HybridResponse {
151    /// Search results
152    pub results: Vec<HybridResult>,
153    /// Total candidates evaluated
154    pub total_evaluated: usize,
155    /// Search latency in microseconds
156    pub latency_us: u64,
157    /// Filter strategy used
158    pub strategy_used: FilterStrategy,
159}
160
161/// Hybrid search index combining HNSW with metadata
162pub struct HybridIndex {
163    /// Vector index
164    vector_index: Arc<RwLock<VectorIndex>>,
165    /// Metadata store
166    metadata_store: Arc<MetadataStore>,
167    /// Configuration
168    config: HybridConfig,
169    /// Statistics
170    stats: Arc<IndexStats>,
171    /// Query cache
172    cache: Arc<RwLock<LruCache<u64, Vec<HybridResult>>>>,
173}
174
175impl HybridIndex {
176    /// Create a new hybrid index
177    pub fn new(config: HybridConfig) -> Result<Self> {
178        let vector_index = VectorIndex::new(
179            config.dimension,
180            config.metric,
181            config.max_connections,
182            config.ef_construction,
183        )?;
184
185        let cache_size = NonZeroUsize::new(config.cache_size)
186            .unwrap_or(NonZeroUsize::new(1000).expect("1000 > 0"));
187
188        Ok(Self {
189            vector_index: Arc::new(RwLock::new(vector_index)),
190            metadata_store: Arc::new(MetadataStore::new()),
191            config,
192            stats: Arc::new(IndexStats::new()),
193            cache: Arc::new(RwLock::new(LruCache::new(cache_size))),
194        })
195    }
196
197    /// Create with default configuration
198    pub fn with_defaults() -> Result<Self> {
199        Self::new(HybridConfig::default())
200    }
201
202    /// Insert a vector with metadata
203    pub fn insert(&self, cid: &Cid, vector: &[f32], metadata: Option<Metadata>) -> Result<()> {
204        let timer = PerfTimer::start();
205
206        // Insert into vector index
207        self.vector_index
208            .write()
209            .unwrap_or_else(|e| e.into_inner())
210            .insert(cid, vector)?;
211
212        // Insert metadata if provided
213        if let Some(meta) = metadata {
214            self.metadata_store.insert(*cid, meta)?;
215        } else {
216            // Create minimal metadata with timestamp
217            self.metadata_store.insert(*cid, Metadata::new())?;
218        }
219
220        if self.config.collect_stats {
221            self.stats.record_insert(timer.stop());
222        }
223
224        // Invalidate cache
225        self.cache
226            .write()
227            .unwrap_or_else(|e| e.into_inner())
228            .clear();
229
230        Ok(())
231    }
232
233    /// Insert multiple vectors with metadata in batch
234    pub fn insert_batch(&self, items: &[(Cid, Vec<f32>, Option<Metadata>)]) -> Result<()> {
235        for (cid, vector, metadata) in items {
236            self.insert(cid, vector, metadata.clone())?;
237        }
238        Ok(())
239    }
240
241    /// Delete a vector and its metadata
242    pub fn delete(&self, cid: &Cid) -> Result<()> {
243        self.vector_index
244            .write()
245            .unwrap_or_else(|e| e.into_inner())
246            .delete(cid)?;
247        self.metadata_store.remove(cid)?;
248
249        if self.config.collect_stats {
250            self.stats.record_delete();
251        }
252
253        // Invalidate cache
254        self.cache
255            .write()
256            .unwrap_or_else(|e| e.into_inner())
257            .clear();
258
259        Ok(())
260    }
261
262    /// Perform hybrid search
263    pub async fn search(&self, query: HybridQuery) -> Result<HybridResponse> {
264        let timer = PerfTimer::start();
265
266        // Determine filter strategy
267        let strategy = self.determine_strategy(&query);
268        let mut total_evaluated = 0;
269
270        let results = match strategy {
271            FilterStrategy::PreFilter => {
272                self.search_pre_filter(&query, &mut total_evaluated).await?
273            }
274            FilterStrategy::PostFilter | FilterStrategy::Auto => {
275                self.search_post_filter(&query, &mut total_evaluated)
276                    .await?
277            }
278        };
279
280        let latency = timer.stop();
281
282        if self.config.collect_stats {
283            self.stats.record_search(latency, query.k, results.len());
284        }
285
286        Ok(HybridResponse {
287            results,
288            total_evaluated,
289            latency_us: latency.as_micros() as u64,
290            strategy_used: strategy,
291        })
292    }
293
294    /// Pre-filter strategy: filter first, then search on subset
295    async fn search_pre_filter(
296        &self,
297        query: &HybridQuery,
298        total_evaluated: &mut usize,
299    ) -> Result<Vec<HybridResult>> {
300        // Get candidate CIDs from filter
301        let candidates: HashSet<Cid> = if let Some(ref filter) = query.filter {
302            self.metadata_store.filter(filter).into_iter().collect()
303        } else {
304            // No filter, use all CIDs
305            self.metadata_store.cids().into_iter().collect()
306        };
307
308        // Apply temporal filter if present
309        let candidates = if let Some(ref temporal) = query.temporal {
310            let time_filtered = self
311                .metadata_store
312                .get_by_time_range(temporal.start, temporal.end);
313            candidates
314                .intersection(&time_filtered.into_iter().collect())
315                .copied()
316                .collect()
317        } else {
318            candidates
319        };
320
321        *total_evaluated = candidates.len();
322
323        if candidates.is_empty() {
324            return Ok(Vec::new());
325        }
326
327        // Search vector index
328        let ef_search = query.ef_search.unwrap_or(self.config.ef_search);
329        let fetch_k = (query.k * 3).max(100); // Fetch more to account for filtering
330
331        let search_results = self
332            .vector_index
333            .read()
334            .unwrap_or_else(|e| e.into_inner())
335            .search(&query.vector, fetch_k, ef_search)?;
336
337        // Filter results to candidates
338        let mut results: Vec<HybridResult> = search_results
339            .into_iter()
340            .filter(|r| candidates.contains(&r.cid))
341            .map(|r| {
342                let mut hr = HybridResult::from(r);
343                // Apply recency boost
344                if let Some(ref temporal) = query.temporal {
345                    if let Some(meta) = self.metadata_store.get(&hr.cid) {
346                        let boost = temporal.recency_multiplier(meta.created_at);
347                        hr.score *= boost;
348                    }
349                }
350                hr
351            })
352            .collect();
353
354        // Apply min score filter
355        if let Some(min_score) = query.min_score {
356            results.retain(|r| r.score >= min_score);
357        }
358
359        // Sort by score and truncate
360        results.sort_by(|a, b| {
361            b.score
362                .partial_cmp(&a.score)
363                .unwrap_or(std::cmp::Ordering::Equal)
364        });
365        results.truncate(query.k);
366
367        // Add metadata if requested
368        if query.include_metadata {
369            for result in &mut results {
370                result.metadata = self.metadata_store.get(&result.cid);
371            }
372        }
373
374        Ok(results)
375    }
376
377    /// Post-filter strategy: search first, then filter results
378    async fn search_post_filter(
379        &self,
380        query: &HybridQuery,
381        total_evaluated: &mut usize,
382    ) -> Result<Vec<HybridResult>> {
383        let ef_search = query.ef_search.unwrap_or(self.config.ef_search);
384
385        // Fetch more results to account for filtering
386        let fetch_k = if query.filter.is_some() || query.temporal.is_some() {
387            (query.k * 5).max(100)
388        } else {
389            query.k
390        };
391
392        let search_results = self
393            .vector_index
394            .read()
395            .unwrap_or_else(|e| e.into_inner())
396            .search(&query.vector, fetch_k, ef_search)?;
397
398        *total_evaluated = search_results.len();
399
400        let mut results: Vec<HybridResult> = search_results
401            .into_iter()
402            .filter_map(|r| {
403                // Apply metadata filter
404                if let Some(ref filter) = query.filter {
405                    if let Some(meta) = self.metadata_store.get(&r.cid) {
406                        if !filter.matches(&meta) {
407                            return None;
408                        }
409                    } else {
410                        return None; // No metadata, filter out
411                    }
412                }
413
414                // Apply temporal filter
415                if let Some(ref temporal) = query.temporal {
416                    if let Some(meta) = self.metadata_store.get(&r.cid) {
417                        if let (Some(start), Some(end)) = (temporal.start, temporal.end) {
418                            if meta.created_at < start || meta.created_at > end {
419                                return None;
420                            }
421                        }
422                    }
423                }
424
425                let mut hr = HybridResult::from(r);
426
427                // Apply recency boost
428                if let Some(ref temporal) = query.temporal {
429                    if let Some(meta) = self.metadata_store.get(&hr.cid) {
430                        let boost = temporal.recency_multiplier(meta.created_at);
431                        hr.score *= boost;
432                    }
433                }
434
435                Some(hr)
436            })
437            .collect();
438
439        // Apply min score filter
440        if let Some(min_score) = query.min_score {
441            results.retain(|r| r.score >= min_score);
442        }
443
444        // Re-sort if recency boost was applied
445        if query.temporal.is_some() {
446            results.sort_by(|a, b| {
447                b.score
448                    .partial_cmp(&a.score)
449                    .unwrap_or(std::cmp::Ordering::Equal)
450            });
451        }
452
453        results.truncate(query.k);
454
455        // Add metadata if requested
456        if query.include_metadata {
457            for result in &mut results {
458                result.metadata = self.metadata_store.get(&result.cid);
459            }
460        }
461
462        Ok(results)
463    }
464
465    /// Determine the best filter strategy
466    fn determine_strategy(&self, query: &HybridQuery) -> FilterStrategy {
467        if self.config.filter_strategy != FilterStrategy::Auto {
468            return self.config.filter_strategy;
469        }
470
471        // Estimate selectivity
472        let total_count = self.metadata_store.len();
473        if total_count == 0 {
474            return FilterStrategy::PostFilter;
475        }
476
477        // If no filter, use post-filter (simpler path)
478        if query.filter.is_none() && query.temporal.is_none() {
479            return FilterStrategy::PostFilter;
480        }
481
482        // Estimate filter selectivity
483        let filtered_count = if let Some(ref filter) = query.filter {
484            self.metadata_store.filter(filter).len()
485        } else {
486            total_count
487        };
488
489        let selectivity = filtered_count as f64 / total_count as f64;
490
491        // Pre-filter if highly selective (< 10% of data)
492        // Post-filter if less selective (more data passes)
493        if selectivity < 0.1 {
494            FilterStrategy::PreFilter
495        } else {
496            FilterStrategy::PostFilter
497        }
498    }
499
500    /// Get the number of indexed vectors
501    pub fn len(&self) -> usize {
502        self.vector_index
503            .read()
504            .unwrap_or_else(|e| e.into_inner())
505            .len()
506    }
507
508    /// Check if the index is empty
509    pub fn is_empty(&self) -> bool {
510        self.len() == 0
511    }
512
513    /// Check if a CID exists
514    pub fn contains(&self, cid: &Cid) -> bool {
515        self.vector_index
516            .read()
517            .unwrap_or_else(|e| e.into_inner())
518            .contains(cid)
519    }
520
521    /// Get metadata for a CID
522    pub fn get_metadata(&self, cid: &Cid) -> Option<Metadata> {
523        self.metadata_store.get(cid)
524    }
525
526    /// Update metadata for a CID (without changing the vector)
527    pub fn update_metadata(&self, cid: &Cid, metadata: Metadata) -> Result<()> {
528        if !self.contains(cid) {
529            return Err(Error::NotFound(format!("CID not in index: {}", cid)));
530        }
531        self.metadata_store.insert(*cid, metadata)?;
532        Ok(())
533    }
534
535    /// Get statistics snapshot
536    pub fn stats(&self) -> StatsSnapshot {
537        self.stats.snapshot()
538    }
539
540    /// Get index health metrics
541    pub fn health(&self) -> IndexHealth {
542        let stats = self.stats.snapshot();
543        IndexHealth::analyze(self.len(), self.config.dimension, Some(&stats))
544    }
545
546    /// Get memory usage estimate
547    pub fn memory_usage(&self) -> MemoryUsage {
548        MemoryUsage::estimate(
549            self.len(),
550            self.config.dimension,
551            self.metadata_store.len(),
552            self.config.cache_size,
553        )
554    }
555
556    /// Get facet counts for a field
557    pub fn facet_counts(&self, field: &str) -> std::collections::HashMap<String, usize> {
558        self.metadata_store.get_facet_counts(field)
559    }
560
561    /// Clear the search cache
562    pub fn clear_cache(&self) {
563        self.cache
564            .write()
565            .unwrap_or_else(|e| e.into_inner())
566            .clear();
567    }
568
569    /// Reset statistics
570    pub fn reset_stats(&self) {
571        self.stats.reset();
572    }
573
574    /// Save the index to a path
575    pub async fn save(&self, path: impl AsRef<std::path::Path>) -> Result<()> {
576        self.vector_index
577            .read()
578            .unwrap_or_else(|e| e.into_inner())
579            .save(path)
580    }
581
582    /// Clear all data
583    pub fn clear(&self) -> Result<()> {
584        // Create new empty vector index
585        let new_index = VectorIndex::new(
586            self.config.dimension,
587            self.config.metric,
588            self.config.max_connections,
589            self.config.ef_construction,
590        )?;
591
592        *self.vector_index.write().unwrap_or_else(|e| e.into_inner()) = new_index;
593        self.metadata_store.clear();
594        self.cache
595            .write()
596            .unwrap_or_else(|e| e.into_inner())
597            .clear();
598        self.stats.reset();
599
600        Ok(())
601    }
602
603    /// Prune entries older than the given TTL (time-to-live in seconds)
604    ///
605    /// Removes vectors and metadata for entries that were created more than
606    /// `ttl_seconds` ago.
607    ///
608    /// # Arguments
609    /// * `ttl_seconds` - Maximum age in seconds for entries to keep
610    ///
611    /// # Returns
612    /// Number of entries pruned
613    pub fn prune_by_ttl(&self, ttl_seconds: u64) -> Result<usize> {
614        let now = std::time::SystemTime::now()
615            .duration_since(std::time::UNIX_EPOCH)
616            .unwrap_or_default()
617            .as_secs();
618
619        let cutoff = now.saturating_sub(ttl_seconds);
620
621        self.prune_older_than(cutoff)
622    }
623
624    /// Prune entries created before a specific timestamp
625    ///
626    /// # Arguments
627    /// * `timestamp` - Unix timestamp; entries created before this are removed
628    ///
629    /// # Returns
630    /// Number of entries pruned
631    pub fn prune_older_than(&self, timestamp: u64) -> Result<usize> {
632        // Find CIDs to remove
633        let cids_to_remove: Vec<Cid> = self
634            .metadata_store
635            .cids()
636            .into_iter()
637            .filter(|cid| {
638                self.metadata_store
639                    .get(cid)
640                    .map(|m| m.created_at < timestamp)
641                    .unwrap_or(false)
642            })
643            .collect();
644
645        let count = cids_to_remove.len();
646
647        // Remove from both indexes
648        for cid in &cids_to_remove {
649            let _ = self
650                .vector_index
651                .write()
652                .unwrap_or_else(|e| e.into_inner())
653                .delete(cid);
654            let _ = self.metadata_store.remove(cid);
655        }
656
657        // Clear cache since data has changed
658        self.cache
659            .write()
660            .unwrap_or_else(|e| e.into_inner())
661            .clear();
662
663        Ok(count)
664    }
665
666    /// Prune entries keeping only the N most recently created
667    ///
668    /// # Arguments
669    /// * `max_entries` - Maximum number of entries to keep
670    ///
671    /// # Returns
672    /// Number of entries pruned
673    pub fn prune_to_max_entries(&self, max_entries: usize) -> Result<usize> {
674        let current_count = self.len();
675        if current_count <= max_entries {
676            return Ok(0);
677        }
678
679        // Get all CIDs with their creation timestamps
680        let mut entries: Vec<(Cid, u64)> = self
681            .metadata_store
682            .cids()
683            .into_iter()
684            .filter_map(|cid| self.metadata_store.get(&cid).map(|m| (cid, m.created_at)))
685            .collect();
686
687        // Sort by creation time (oldest first)
688        entries.sort_by_key(|(_, ts)| *ts);
689
690        // Calculate how many to remove
691        let to_remove = current_count - max_entries;
692
693        // Remove the oldest entries
694        for (cid, _) in entries.iter().take(to_remove) {
695            let _ = self
696                .vector_index
697                .write()
698                .unwrap_or_else(|e| e.into_inner())
699                .delete(cid);
700            let _ = self.metadata_store.remove(cid);
701        }
702
703        // Clear cache
704        self.cache
705            .write()
706            .unwrap_or_else(|e| e.into_inner())
707            .clear();
708
709        Ok(to_remove)
710    }
711
712    /// Prune entries by LRU (Least Recently Updated)
713    ///
714    /// Removes entries that haven't been updated recently, keeping
715    /// only the most recently updated entries.
716    ///
717    /// # Arguments
718    /// * `max_entries` - Maximum number of entries to keep
719    ///
720    /// # Returns
721    /// Number of entries pruned
722    pub fn prune_lru(&self, max_entries: usize) -> Result<usize> {
723        let current_count = self.len();
724        if current_count <= max_entries {
725            return Ok(0);
726        }
727
728        // Get all CIDs with their update timestamps
729        let mut entries: Vec<(Cid, u64)> = self
730            .metadata_store
731            .cids()
732            .into_iter()
733            .filter_map(|cid| self.metadata_store.get(&cid).map(|m| (cid, m.updated_at)))
734            .collect();
735
736        // Sort by update time (least recent first)
737        entries.sort_by_key(|(_, ts)| *ts);
738
739        // Calculate how many to remove
740        let to_remove = current_count - max_entries;
741
742        // Remove the least recently updated entries
743        for (cid, _) in entries.iter().take(to_remove) {
744            let _ = self
745                .vector_index
746                .write()
747                .unwrap_or_else(|e| e.into_inner())
748                .delete(cid);
749            let _ = self.metadata_store.remove(cid);
750        }
751
752        // Clear cache
753        self.cache
754            .write()
755            .unwrap_or_else(|e| e.into_inner())
756            .clear();
757
758        Ok(to_remove)
759    }
760
761    /// Get pruning statistics
762    pub fn pruning_stats(&self) -> PruningStats {
763        let now = std::time::SystemTime::now()
764            .duration_since(std::time::UNIX_EPOCH)
765            .unwrap_or_default()
766            .as_secs();
767
768        let entries: Vec<(u64, u64)> = self
769            .metadata_store
770            .cids()
771            .into_iter()
772            .filter_map(|cid| {
773                self.metadata_store
774                    .get(&cid)
775                    .map(|m| (m.created_at, m.updated_at))
776            })
777            .collect();
778
779        if entries.is_empty() {
780            return PruningStats::default();
781        }
782
783        let oldest_created = entries.iter().map(|(c, _)| *c).min().unwrap_or(now);
784        let newest_created = entries.iter().map(|(c, _)| *c).max().unwrap_or(now);
785        let oldest_updated = entries.iter().map(|(_, u)| *u).min().unwrap_or(now);
786
787        let age_1day = entries.iter().filter(|(c, _)| now - *c < 86400).count();
788        let age_7days = entries.iter().filter(|(c, _)| now - *c < 86400 * 7).count();
789        let age_30days = entries
790            .iter()
791            .filter(|(c, _)| now - *c < 86400 * 30)
792            .count();
793
794        PruningStats {
795            total_entries: entries.len(),
796            oldest_entry_age: now.saturating_sub(oldest_created),
797            newest_entry_age: now.saturating_sub(newest_created),
798            oldest_update_age: now.saturating_sub(oldest_updated),
799            entries_last_day: age_1day,
800            entries_last_week: age_7days,
801            entries_last_month: age_30days,
802        }
803    }
804}
805
806/// Pruning statistics
807#[derive(Debug, Clone, Default, Serialize, Deserialize)]
808pub struct PruningStats {
809    /// Total number of entries
810    pub total_entries: usize,
811    /// Age of the oldest entry in seconds
812    pub oldest_entry_age: u64,
813    /// Age of the newest entry in seconds
814    pub newest_entry_age: u64,
815    /// Age of the least recently updated entry in seconds
816    pub oldest_update_age: u64,
817    /// Number of entries created in the last day
818    pub entries_last_day: usize,
819    /// Number of entries created in the last week
820    pub entries_last_week: usize,
821    /// Number of entries created in the last month
822    pub entries_last_month: usize,
823}
824
825impl PruningStats {
826    /// Get a summary string
827    pub fn summary(&self) -> String {
828        format!(
829            "Total: {}, Last day: {}, Last week: {}, Last month: {}, Oldest: {}s ago",
830            self.total_entries,
831            self.entries_last_day,
832            self.entries_last_week,
833            self.entries_last_month,
834            self.oldest_entry_age
835        )
836    }
837
838    /// Estimate entries that would be pruned for a given TTL
839    pub fn would_prune_for_ttl(&self, ttl_seconds: u64) -> usize {
840        // Approximate based on time buckets
841        if ttl_seconds < 86400 {
842            self.total_entries - self.entries_last_day
843        } else if ttl_seconds < 86400 * 7 {
844            self.total_entries - self.entries_last_week
845        } else if ttl_seconds < 86400 * 30 {
846            self.total_entries - self.entries_last_month
847        } else {
848            0
849        }
850    }
851}
852
853#[cfg(test)]
854mod tests {
855    use super::*;
856    use crate::metadata::MetadataValue;
857
858    fn test_cid(n: u8) -> Cid {
859        // Use different valid CID strings
860        let cids = [
861            "bafybeigdyrzt5sfp7udm7hu76uh7y26nf3efuylqabf3oclgtqy55fbzdi",
862            "bafybeiczsscdsbs7ffqz55asqdf3smv6klcw3gofszvwlyarci47bgf354",
863            "bafybeibvfkifsqbapirjrj7zbfwddz5qz5awvbftjgktpcqcxjkzstszlm",
864        ];
865        cids[n as usize % cids.len()]
866            .parse()
867            .expect("test: parse test cid")
868    }
869
870    #[tokio::test]
871    async fn test_hybrid_index_basic() {
872        let config = HybridConfig {
873            dimension: 4,
874            ..Default::default()
875        };
876
877        let index = HybridIndex::new(config).expect("test: create hybrid index basic");
878
879        let cid1 = test_cid(0);
880        let vec1 = vec![1.0, 0.0, 0.0, 0.0];
881        let meta1 = Metadata::new().with_string("type", "image");
882
883        index
884            .insert(&cid1, &vec1, Some(meta1))
885            .expect("test: insert cid1 basic");
886
887        assert_eq!(index.len(), 1);
888        assert!(index.contains(&cid1));
889    }
890
891    #[tokio::test]
892    async fn test_hybrid_search() {
893        let config = HybridConfig {
894            dimension: 4,
895            ..Default::default()
896        };
897
898        let index = HybridIndex::new(config).expect("test: create hybrid index for search");
899
900        // Insert some vectors with metadata (more vectors for better HNSW graph connectivity)
901        let cid1 = test_cid(0);
902        let vec1 = vec![1.0, 0.0, 0.0, 0.0];
903        let meta1 = Metadata::new()
904            .with_string("type", "image")
905            .with_integer("size", 1024);
906
907        let cid2 = test_cid(1);
908        let vec2 = vec![0.9, 0.1, 0.0, 0.0];
909        let meta2 = Metadata::new()
910            .with_string("type", "document")
911            .with_integer("size", 2048);
912
913        let cid3 = test_cid(2);
914        let vec3 = vec![0.0, 1.0, 0.0, 0.0];
915        let meta3 = Metadata::new()
916            .with_string("type", "audio")
917            .with_integer("size", 512);
918
919        index
920            .insert(&cid1, &vec1, Some(meta1))
921            .expect("test: insert cid1 for search");
922        index
923            .insert(&cid2, &vec2, Some(meta2))
924            .expect("test: insert cid2 for search");
925        index
926            .insert(&cid3, &vec3, Some(meta3))
927            .expect("test: insert cid3 for search");
928
929        // Simple k-NN search with explicit ef_search to ensure results are found
930        let mut query = HybridQuery::knn(vec![1.0, 0.0, 0.0, 0.0], 2);
931        query.ef_search = Some(50); // Ensure we search enough candidates
932        let response = index.search(query).await.expect("test: hybrid search");
933
934        assert!(
935            !response.results.is_empty(),
936            "Expected at least 1 result, got {}",
937            response.results.len()
938        );
939        // With 3 vectors and k=2, we should get 2 results
940        assert!(
941            !response.results.is_empty() && response.results.len() <= 2,
942            "Expected 1-2 results, got {}",
943            response.results.len()
944        );
945        // First result should be exact match (cid1)
946        assert_eq!(response.results[0].cid, cid1);
947    }
948
949    #[tokio::test]
950    async fn test_filtered_search() {
951        let config = HybridConfig {
952            dimension: 4,
953            ..Default::default()
954        };
955
956        let index =
957            HybridIndex::new(config).expect("test: create hybrid index for filtered search");
958
959        let cid1 = test_cid(0);
960        let vec1 = vec![1.0, 0.0, 0.0, 0.0];
961        let meta1 = Metadata::new().with_string("category", "tech");
962
963        let cid2 = test_cid(1);
964        let vec2 = vec![0.9, 0.1, 0.0, 0.0];
965        let meta2 = Metadata::new().with_string("category", "science");
966
967        index
968            .insert(&cid1, &vec1, Some(meta1))
969            .expect("test: insert cid1 filtered");
970        index
971            .insert(&cid2, &vec2, Some(meta2))
972            .expect("test: insert cid2 filtered");
973
974        // Search with filter
975        let filter = MetadataFilter::eq("category", MetadataValue::String("tech".to_string()));
976        let query = HybridQuery::knn(vec![0.9, 0.1, 0.0, 0.0], 10).with_filter(filter);
977        let response = index.search(query).await.expect("test: filtered search");
978
979        // Should only return tech category
980        assert_eq!(response.results.len(), 1);
981        assert_eq!(response.results[0].cid, cid1);
982    }
983
984    #[tokio::test]
985    async fn test_search_with_metadata() {
986        let config = HybridConfig {
987            dimension: 4,
988            ..Default::default()
989        };
990
991        let index = HybridIndex::new(config).expect("test: create hybrid index with metadata");
992
993        let cid1 = test_cid(0);
994        let vec1 = vec![1.0, 0.0, 0.0, 0.0];
995        let meta1 = Metadata::new().with_string("title", "Test Document");
996
997        index
998            .insert(&cid1, &vec1, Some(meta1))
999            .expect("test: insert cid1 with metadata");
1000
1001        let query = HybridQuery::knn(vec![1.0, 0.0, 0.0, 0.0], 1).with_metadata();
1002        let response = index
1003            .search(query)
1004            .await
1005            .expect("test: search with metadata");
1006
1007        assert_eq!(response.results.len(), 1);
1008        assert!(response.results[0].metadata.is_some());
1009
1010        let meta = response.results[0]
1011            .metadata
1012            .as_ref()
1013            .expect("test: result should have metadata");
1014        assert_eq!(
1015            meta.get("title"),
1016            Some(&MetadataValue::String("Test Document".to_string()))
1017        );
1018    }
1019
1020    #[test]
1021    fn test_health_and_stats() {
1022        let config = HybridConfig {
1023            dimension: 4,
1024            ..Default::default()
1025        };
1026
1027        let index = HybridIndex::new(config).expect("test: create hybrid index for health stats");
1028
1029        let health = index.health();
1030        assert_eq!(health.size, 0);
1031
1032        let stats = index.stats();
1033        assert_eq!(stats.search_count, 0);
1034    }
1035
1036    #[test]
1037    fn test_pruning_to_max_entries() {
1038        let config = HybridConfig {
1039            dimension: 4,
1040            ..Default::default()
1041        };
1042
1043        let index = HybridIndex::new(config).expect("test: create hybrid index for pruning");
1044
1045        // Insert 3 entries
1046        for i in 0..3 {
1047            let cid = test_cid(i);
1048            let vec = vec![i as f32, 0.0, 0.0, 0.0];
1049            let meta = Metadata::new().with_integer("order", i as i64);
1050            index
1051                .insert(&cid, &vec, Some(meta))
1052                .expect("test: insert vector for pruning");
1053        }
1054
1055        assert_eq!(index.len(), 3);
1056
1057        // Prune to max 2 entries
1058        let pruned = index
1059            .prune_to_max_entries(2)
1060            .expect("test: prune to max entries");
1061        assert_eq!(pruned, 1);
1062        assert_eq!(index.len(), 2);
1063    }
1064
1065    #[test]
1066    fn test_pruning_stats() {
1067        let config = HybridConfig {
1068            dimension: 4,
1069            ..Default::default()
1070        };
1071
1072        let index = HybridIndex::new(config).expect("test: create hybrid index for pruning stats");
1073
1074        // Insert some entries
1075        for i in 0..3 {
1076            let cid = test_cid(i);
1077            let vec = vec![i as f32, 0.0, 0.0, 0.0];
1078            index
1079                .insert(&cid, &vec, None)
1080                .expect("test: insert vector for pruning stats");
1081        }
1082
1083        let stats = index.pruning_stats();
1084        assert_eq!(stats.total_entries, 3);
1085        // All entries should be recent (created just now)
1086        assert_eq!(stats.entries_last_day, 3);
1087        assert_eq!(stats.entries_last_week, 3);
1088    }
1089}