1use async_trait::async_trait;
125use ipfrs_core::{Cid, Error, Result};
126use serde::{Deserialize, Serialize};
127use std::collections::HashMap;
128
129use crate::hnsw::{DistanceMetric, VectorIndex};
130use crate::metadata::{Metadata, MetadataFilter};
131
132#[derive(Debug, Clone, Serialize, Deserialize)]
134pub struct BackendConfig {
135 pub dimension: usize,
137 pub metric: DistanceMetric,
139 pub params: HashMap<String, String>,
141}
142
143impl Default for BackendConfig {
144 fn default() -> Self {
145 Self {
146 dimension: 768,
147 metric: DistanceMetric::Cosine,
148 params: HashMap::new(),
149 }
150 }
151}
152
153#[derive(Debug, Clone)]
155pub struct BackendSearchResult {
156 pub cid: Cid,
158 pub score: f32,
160 pub metadata: Option<Metadata>,
162}
163
164#[async_trait]
166pub trait VectorBackend: Send + Sync {
167 fn insert(&mut self, cid: Cid, vector: &[f32], metadata: Option<Metadata>) -> Result<()>;
169
170 fn insert_batch(&mut self, items: &[(Cid, Vec<f32>, Option<Metadata>)]) -> Result<()> {
172 for (cid, vector, metadata) in items {
173 self.insert(*cid, vector, metadata.clone())?;
174 }
175 Ok(())
176 }
177
178 fn search(
180 &mut self,
181 query: &[f32],
182 k: usize,
183 filter: Option<&MetadataFilter>,
184 ) -> Result<Vec<BackendSearchResult>>;
185
186 fn search_batch(
188 &mut self,
189 queries: &[Vec<f32>],
190 k: usize,
191 filter: Option<&MetadataFilter>,
192 ) -> Result<Vec<Vec<BackendSearchResult>>> {
193 let mut results = Vec::new();
194 for query in queries {
195 results.push(self.search(query, k, filter)?);
196 }
197 Ok(results)
198 }
199
200 fn delete(&mut self, cid: &Cid) -> Result<()>;
202
203 fn update(&mut self, cid: &Cid, vector: &[f32], metadata: Option<Metadata>) -> Result<()> {
205 self.delete(cid)?;
206 self.insert(*cid, vector, metadata)
207 }
208
209 fn get(&self, cid: &Cid) -> Result<Option<(Vec<f32>, Option<Metadata>)>>;
211
212 fn count(&self) -> Result<usize>;
214
215 fn clear(&mut self) -> Result<()>;
217
218 fn backend_name(&self) -> &str;
220
221 fn stats(&self) -> BackendStats;
223}
224
225#[derive(Debug, Clone, Default)]
227pub struct BackendStats {
228 pub vector_count: usize,
230 pub searches: usize,
232 pub insertions: usize,
234 pub custom_metrics: HashMap<String, f64>,
236}
237
238pub struct IpfrsBackend {
240 index: VectorIndex,
242 vector_store: HashMap<Cid, Vec<f32>>,
244 metadata_store: HashMap<Cid, Metadata>,
246 config: BackendConfig,
248 stats: BackendStats,
250}
251
252impl IpfrsBackend {
253 pub fn new(config: BackendConfig) -> Result<Self> {
255 let index = VectorIndex::new(
256 config.dimension,
257 config.metric,
258 16, 200, )?;
261
262 Ok(Self {
263 index,
264 vector_store: HashMap::new(),
265 metadata_store: HashMap::new(),
266 config,
267 stats: BackendStats::default(),
268 })
269 }
270
271 pub fn index(&self) -> &VectorIndex {
273 &self.index
274 }
275
276 pub fn index_mut(&mut self) -> &mut VectorIndex {
278 &mut self.index
279 }
280}
281
282#[async_trait]
283impl VectorBackend for IpfrsBackend {
284 fn insert(&mut self, cid: Cid, vector: &[f32], metadata: Option<Metadata>) -> Result<()> {
285 self.index.insert(&cid, vector)?;
286 self.vector_store.insert(cid, vector.to_vec());
287 if let Some(meta) = metadata {
288 self.metadata_store.insert(cid, meta);
289 }
290 self.stats.insertions += 1;
291 self.stats.vector_count = self.index.len();
292 Ok(())
293 }
294
295 fn insert_batch(&mut self, items: &[(Cid, Vec<f32>, Option<Metadata>)]) -> Result<()> {
296 for (cid, vector, metadata) in items {
297 self.index.insert(cid, vector)?;
298 self.vector_store.insert(*cid, vector.clone());
299 if let Some(meta) = metadata {
300 self.metadata_store.insert(*cid, meta.clone());
301 }
302 self.stats.insertions += 1;
303 }
304 self.stats.vector_count = self.index.len();
305 Ok(())
306 }
307
308 fn search(
309 &mut self,
310 query: &[f32],
311 k: usize,
312 filter: Option<&MetadataFilter>,
313 ) -> Result<Vec<BackendSearchResult>> {
314 let ef_search = 50; let raw_results = self.index.search(query, k * 2, ef_search)?; self.stats.searches += 1;
317
318 let mut results = Vec::new();
319 for result in raw_results {
320 if let Some(filter) = filter {
322 if let Some(metadata) = self.metadata_store.get(&result.cid) {
323 if !filter.matches(metadata) {
324 continue;
325 }
326 } else {
327 continue;
328 }
329 }
330
331 results.push(BackendSearchResult {
332 cid: result.cid,
333 score: result.score,
334 metadata: self.metadata_store.get(&result.cid).cloned(),
335 });
336
337 if results.len() >= k {
338 break;
339 }
340 }
341
342 Ok(results)
343 }
344
345 fn delete(&mut self, cid: &Cid) -> Result<()> {
346 self.index.delete(cid)?;
347 self.vector_store.remove(cid);
348 self.metadata_store.remove(cid);
349 self.stats.vector_count = self.index.len();
350 Ok(())
351 }
352
353 fn get(&self, cid: &Cid) -> Result<Option<(Vec<f32>, Option<Metadata>)>> {
354 if let Some(vector) = self.vector_store.get(cid) {
355 let metadata = self.metadata_store.get(cid).cloned();
356 Ok(Some((vector.clone(), metadata)))
357 } else {
358 Ok(None)
359 }
360 }
361
362 fn count(&self) -> Result<usize> {
363 Ok(self.index.len())
364 }
365
366 fn clear(&mut self) -> Result<()> {
367 self.index = VectorIndex::new(self.config.dimension, self.config.metric, 16, 200)?;
369 self.vector_store.clear();
370 self.metadata_store.clear();
371 self.stats = BackendStats::default();
372 Ok(())
373 }
374
375 fn backend_name(&self) -> &str {
376 "ipfrs-hnsw"
377 }
378
379 fn stats(&self) -> BackendStats {
380 self.stats.clone()
381 }
382}
383
384pub struct BackendMigration;
386
387impl BackendMigration {
388 #[allow(dead_code)]
390 pub fn migrate(
391 _source: &dyn VectorBackend,
392 _dest: &mut dyn VectorBackend,
393 ) -> Result<MigrationStats> {
394 let stats = MigrationStats::default();
395
396 Ok(stats)
401 }
402
403 pub fn migrate_cids(
405 source: &dyn VectorBackend,
406 dest: &mut dyn VectorBackend,
407 cids: &[Cid],
408 ) -> Result<MigrationStats> {
409 Self::migrate_cids_with_progress(source, dest, cids, |_, _| {})
410 }
411
412 pub fn migrate_cids_with_progress<F>(
431 source: &dyn VectorBackend,
432 dest: &mut dyn VectorBackend,
433 cids: &[Cid],
434 mut progress_callback: F,
435 ) -> Result<MigrationStats>
436 where
437 F: FnMut(usize, usize),
438 {
439 let mut stats = MigrationStats::default();
440 let total = cids.len();
441
442 for (index, cid) in cids.iter().enumerate() {
443 if let Some((vector, metadata)) = source.get(cid)? {
444 dest.insert(*cid, &vector, metadata)?;
445 stats.migrated += 1;
446 } else {
447 stats.not_found += 1;
448 }
449
450 progress_callback(index + 1, total);
452 }
453
454 Ok(stats)
455 }
456
457 pub fn export_to_json(backend: &dyn VectorBackend, cids: &[Cid]) -> Result<String> {
459 let mut exports = Vec::new();
460
461 for cid in cids {
462 if let Some((vector, metadata)) = backend.get(cid)? {
463 let export = ExportedVector {
464 cid: cid.to_string(),
465 vector,
466 metadata,
467 };
468 exports.push(export);
469 }
470 }
471
472 serde_json::to_string_pretty(&exports)
473 .map_err(|e| Error::Serialization(format!("JSON export failed: {}", e)))
474 }
475
476 pub fn import_from_json(backend: &mut dyn VectorBackend, json: &str) -> Result<usize> {
478 let exports: Vec<ExportedVector> = serde_json::from_str(json)
479 .map_err(|e| Error::Serialization(format!("JSON import failed: {}", e)))?;
480
481 let mut count = 0;
482 for export in exports {
483 let cid: Cid = export
484 .cid
485 .parse()
486 .map_err(|e| Error::InvalidInput(format!("Invalid CID: {}", e)))?;
487 backend.insert(cid, &export.vector, export.metadata)?;
488 count += 1;
489 }
490
491 Ok(count)
492 }
493}
494
495#[derive(Debug, Clone, Default)]
497pub struct MigrationStats {
498 pub migrated: usize,
500 pub not_found: usize,
502 pub errors: usize,
504}
505
506#[derive(Debug, Clone, Serialize, Deserialize)]
508struct ExportedVector {
509 cid: String,
510 vector: Vec<f32>,
511 metadata: Option<Metadata>,
512}
513
514pub struct BackendRegistry {
516 backends: HashMap<String, Box<dyn VectorBackend>>,
517 default_backend: Option<String>,
518}
519
520impl BackendRegistry {
521 pub fn new() -> Self {
523 Self {
524 backends: HashMap::new(),
525 default_backend: None,
526 }
527 }
528
529 pub fn register(&mut self, name: String, backend: Box<dyn VectorBackend>) {
531 if self.default_backend.is_none() {
532 self.default_backend = Some(name.clone());
533 }
534 self.backends.insert(name, backend);
535 }
536
537 pub fn get(&self, name: &str) -> Option<&dyn VectorBackend> {
539 self.backends.get(name).map(|b| b.as_ref())
540 }
541
542 pub fn get_mut(&mut self, name: &str) -> Option<&mut (dyn VectorBackend + '_)> {
544 match self.backends.get_mut(name) {
545 Some(backend) => Some(backend.as_mut()),
546 None => None,
547 }
548 }
549
550 pub fn get_default(&self) -> Option<&dyn VectorBackend> {
552 self.default_backend
553 .as_ref()
554 .and_then(|name| self.get(name))
555 }
556
557 pub fn get_default_mut(&mut self) -> Option<&mut (dyn VectorBackend + '_)> {
559 if let Some(name) = self.default_backend.clone() {
560 self.get_mut(&name)
561 } else {
562 None
563 }
564 }
565
566 pub fn set_default(&mut self, name: String) -> Result<()> {
568 if self.backends.contains_key(&name) {
569 self.default_backend = Some(name);
570 Ok(())
571 } else {
572 Err(Error::NotFound(format!("Backend '{}' not found", name)))
573 }
574 }
575
576 pub fn list_backends(&self) -> Vec<String> {
578 self.backends.keys().cloned().collect()
579 }
580}
581
582impl Default for BackendRegistry {
583 fn default() -> Self {
584 Self::new()
585 }
586}
587
588#[cfg(test)]
589mod tests {
590 use super::*;
591
592 #[test]
593 fn test_ipfrs_backend_creation() {
594 let config = BackendConfig::default();
595 let backend = IpfrsBackend::new(config);
596 assert!(backend.is_ok());
597 }
598
599 #[test]
600 fn test_insert_and_search() {
601 let config = BackendConfig {
602 dimension: 4,
603 ..Default::default()
604 };
605 let mut backend = IpfrsBackend::new(config).unwrap();
606
607 let cid = Cid::default();
608 let vector = vec![1.0, 2.0, 3.0, 4.0];
609 backend.insert(cid, &vector, None).unwrap();
610
611 let query = vec![1.1, 2.1, 3.1, 4.1];
612 let results = backend.search(&query, 1, None).unwrap();
613
614 assert_eq!(results.len(), 1);
615 assert_eq!(results[0].cid, cid);
616 }
617
618 #[test]
619 fn test_insert_with_metadata() {
620 use crate::metadata::MetadataValue;
621
622 let config = BackendConfig {
623 dimension: 3,
624 ..Default::default()
625 };
626 let mut backend = IpfrsBackend::new(config).unwrap();
627
628 let cid = Cid::default();
629 let vector = vec![1.0, 2.0, 3.0];
630 let mut metadata = Metadata::new();
631 metadata.set("key", MetadataValue::String("value".to_string()));
632
633 backend.insert(cid, &vector, Some(metadata)).unwrap();
634
635 let retrieved = backend.get(&cid).unwrap();
636 assert!(retrieved.is_some());
637 let (_, meta) = retrieved.unwrap();
638 assert!(meta.is_some());
639 }
640
641 #[test]
642 fn test_batch_insert() {
643 use multihash_codetable::{Code, MultihashDigest};
644
645 let config = BackendConfig {
646 dimension: 2,
647 ..Default::default()
648 };
649 let mut backend = IpfrsBackend::new(config).unwrap();
650
651 let cid1 = Cid::new_v1(0x55, Code::Sha2_256.digest(b"test_batch_1"));
653 let cid2 = Cid::new_v1(0x55, Code::Sha2_256.digest(b"test_batch_2"));
654 let cid3 = Cid::new_v1(0x55, Code::Sha2_256.digest(b"test_batch_3"));
655
656 let items = vec![
657 (cid1, vec![1.0, 2.0], None),
658 (cid2, vec![3.0, 4.0], None),
659 (cid3, vec![5.0, 6.0], None),
660 ];
661
662 backend.insert_batch(&items).unwrap();
663 assert_eq!(backend.count().unwrap(), 3);
664 }
665
666 #[test]
667 fn test_delete() {
668 let config = BackendConfig {
669 dimension: 2,
670 ..Default::default()
671 };
672 let mut backend = IpfrsBackend::new(config).unwrap();
673
674 let cid = Cid::default();
675 let vector = vec![1.0, 2.0];
676 backend.insert(cid, &vector, None).unwrap();
677
678 assert_eq!(backend.count().unwrap(), 1);
679
680 backend.delete(&cid).unwrap();
681 assert_eq!(backend.count().unwrap(), 0);
682 }
683
684 #[test]
685 fn test_update() {
686 let config = BackendConfig {
687 dimension: 2,
688 ..Default::default()
689 };
690 let mut backend = IpfrsBackend::new(config).unwrap();
691
692 let cid = Cid::default();
693 let vector1 = vec![1.0, 2.0];
694 backend.insert(cid, &vector1, None).unwrap();
695
696 let vector2 = vec![3.0, 4.0];
697 backend.update(&cid, &vector2, None).unwrap();
698
699 let retrieved = backend.get(&cid).unwrap().unwrap();
700 assert_eq!(retrieved.0, vector2);
701 }
702
703 #[test]
704 fn test_clear() {
705 use multihash_codetable::{Code, MultihashDigest};
706
707 let config = BackendConfig {
708 dimension: 2,
709 ..Default::default()
710 };
711 let mut backend = IpfrsBackend::new(config).unwrap();
712
713 let cid1 = Cid::new_v1(0x55, Code::Sha2_256.digest(b"test_clear_1"));
715 let cid2 = Cid::new_v1(0x55, Code::Sha2_256.digest(b"test_clear_2"));
716
717 backend.insert(cid1, &[1.0, 2.0], None).unwrap();
718 backend.insert(cid2, &[3.0, 4.0], None).unwrap();
719
720 assert_eq!(backend.count().unwrap(), 2);
721
722 backend.clear().unwrap();
723 assert_eq!(backend.count().unwrap(), 0);
724 }
725
726 #[test]
727 fn test_stats() {
728 let config = BackendConfig {
729 dimension: 2,
730 ..Default::default()
731 };
732 let mut backend = IpfrsBackend::new(config).unwrap();
733
734 backend.insert(Cid::default(), &[1.0, 2.0], None).unwrap();
735 backend.search(&[1.0, 2.0], 1, None).unwrap();
736
737 let stats = backend.stats();
738 assert_eq!(stats.insertions, 1);
739 assert_eq!(stats.searches, 1);
740 }
741
742 #[test]
743 fn test_backend_registry() {
744 let mut registry = BackendRegistry::new();
745
746 let config = BackendConfig {
747 dimension: 2,
748 ..Default::default()
749 };
750 let backend = IpfrsBackend::new(config).unwrap();
751
752 registry.register("test".to_string(), Box::new(backend));
753
754 assert!(registry.get("test").is_some());
755 assert_eq!(registry.list_backends().len(), 1);
756 }
757
758 #[test]
759 fn test_migration_stats() {
760 let stats = MigrationStats::default();
761 assert_eq!(stats.migrated, 0);
762 assert_eq!(stats.not_found, 0);
763 assert_eq!(stats.errors, 0);
764 }
765
766 #[test]
767 fn test_export_import() {
768 let config = BackendConfig {
769 dimension: 3,
770 ..Default::default()
771 };
772 let mut backend = IpfrsBackend::new(config.clone()).unwrap();
773
774 let cid = Cid::default();
775 let vector = vec![1.0, 2.0, 3.0];
776 backend.insert(cid, &vector, None).unwrap();
777
778 let json = BackendMigration::export_to_json(&backend, &[cid]).unwrap();
780 assert!(!json.is_empty());
781
782 let mut backend2 = IpfrsBackend::new(config).unwrap();
784 let count = BackendMigration::import_from_json(&mut backend2, &json).unwrap();
785 assert_eq!(count, 1);
786 assert_eq!(backend2.count().unwrap(), 1);
787 }
788}