77 lines
2.0 KiB
Rust
77 lines
2.0 KiB
Rust
use mem_store::{VectorStore, VectorRecord};
|
|
|
|
#[test]
|
|
fn a1_insert_and_search() {
|
|
let mut store = VectorStore::new();
|
|
|
|
// Insert two similar vectors
|
|
let v1 = vec![1.0, 0.0, 0.0];
|
|
let v2 = vec![0.99, 0.1, 0.0];
|
|
let v3 = vec![0.0, 0.0, 1.0]; // orthogonal
|
|
|
|
store.insert(VectorRecord {
|
|
id: "r1".to_string(),
|
|
chunk_id: "c1".to_string(),
|
|
kind: "text".to_string(),
|
|
embedding: v1,
|
|
tokens: 100,
|
|
}).unwrap();
|
|
|
|
store.insert(VectorRecord {
|
|
id: "r2".to_string(),
|
|
chunk_id: "c2".to_string(),
|
|
kind: "text".to_string(),
|
|
embedding: v2,
|
|
tokens: 100,
|
|
}).unwrap();
|
|
|
|
store.insert(VectorRecord {
|
|
id: "r3".to_string(),
|
|
chunk_id: "c3".to_string(),
|
|
kind: "text".to_string(),
|
|
embedding: v3,
|
|
tokens: 100,
|
|
}).unwrap();
|
|
|
|
// Search for vectors similar to v1
|
|
let results = store.search(&[1.0, 0.0, 0.0], 3, 0.0).unwrap();
|
|
|
|
// r1 should be first (identical)
|
|
assert_eq!(results[0].0, "r1");
|
|
assert!((results[0].1 - 1.0).abs() < 0.01);
|
|
|
|
// r2 should be second (similar)
|
|
assert_eq!(results[1].0, "r2");
|
|
assert!(results[1].1 > 0.9);
|
|
|
|
// r3 should be last (orthogonal)
|
|
assert_eq!(results[2].0, "r3");
|
|
assert!(results[2].1 < 0.1);
|
|
}
|
|
|
|
#[test]
|
|
fn a2_min_score_filter() {
|
|
let mut store = VectorStore::new();
|
|
|
|
store.insert(VectorRecord {
|
|
id: "r1".to_string(),
|
|
chunk_id: "c1".to_string(),
|
|
kind: "text".to_string(),
|
|
embedding: vec![1.0, 0.0],
|
|
tokens: 100,
|
|
}).unwrap();
|
|
|
|
store.insert(VectorRecord {
|
|
id: "r2".to_string(),
|
|
chunk_id: "c2".to_string(),
|
|
kind: "text".to_string(),
|
|
embedding: vec![0.0, 1.0],
|
|
tokens: 100,
|
|
}).unwrap();
|
|
|
|
// Search with high threshold - only perfect match
|
|
let results = store.search(&[1.0, 0.0], 10, 0.99).unwrap();
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].0, "r1");
|
|
}
|