use git_vdb::{open, Condition, Document, Embedder, Error, Filter, Result, TextQuery};
use serde_json::json;
use tempfile::TempDir;
#[derive(Clone, Debug)]
struct CompassEmbedder(&'static str);
impl Embedder for CompassEmbedder {
fn model_id(&self) -> &str {
self.0
}
fn embed(&self, input: &[String]) -> Result<Vec<Vec<f32>>> {
Ok(input
.iter()
.map(|text| {
if text.to_ascii_lowercase().contains("north") {
vec![0.0, 1.0]
} else {
vec![1.0, 0.0]
}
})
.collect())
}
}
#[test]
fn documents_are_embedded_stored_and_queried_offline() {
let temp = TempDir::new().unwrap();
let store = open(temp.path().join("documents.git")).unwrap();
let documents = store
.text_collection("docs", CompassEmbedder("test/compass@1"))
.unwrap();
documents
.upsert_documents([
Document::new("east", "A document about the east")
.with_metadata(json!({"kind": "guide"}))
.unwrap(),
Document::new("north", "A document about the north"),
])
.unwrap();
let hits = documents.search_text("north wind", 1).unwrap();
assert_eq!(hits[0].id.to_string(), "north");
assert_eq!(
hits[0].payload.as_ref().unwrap()["document"],
"A document about the north"
);
assert_eq!(
documents
.vectors()
.advanced()
.unwrap()
.info()
.unwrap()
.config
.vector_space
.as_deref(),
Some("test/compass@1")
);
}
#[test]
fn typed_document_queries_support_filters_and_batches() {
let temp = TempDir::new().unwrap();
let documents = open(temp.path().join("documents.git"))
.unwrap()
.text_collection("docs", CompassEmbedder("test/compass@1"))
.unwrap();
documents
.upsert_documents([
Document::new("east-guide", "A guide about the east")
.with_metadata(json!({"kind": "guide"}))
.unwrap(),
Document::new("north-guide", "A guide about the north")
.with_metadata(json!({"kind": "guide"}))
.unwrap(),
Document::new("north-note", "A note about the north")
.with_metadata(json!({"kind": "note"}))
.unwrap(),
])
.unwrap();
let hits = documents
.query(
TextQuery::new("north wind")
.limit(5)
.with_filter(Filter::must([Condition::matches("kind", "guide")])),
)
.unwrap();
assert_eq!(hits.len(), 2);
assert_eq!(hits[0].id.to_string(), "north-guide");
assert_eq!(hits[0].document, "A guide about the north");
assert_eq!(hits[0].metadata["kind"], "guide");
assert!(!hits[0].metadata.contains_key("document"));
let batches = documents
.query_batch([
TextQuery::new("east").limit(1),
TextQuery::new("north").limit(1),
])
.unwrap();
assert_eq!(batches.len(), 2);
assert_eq!(batches[0][0].id.to_string(), "east-guide");
assert!(batches[1][0].document.contains("north"));
}
#[test]
fn model_identity_is_required_and_cannot_change() {
let temp = TempDir::new().unwrap();
let store = open(temp.path().join("documents.git")).unwrap();
assert!(matches!(
store.text_collection("docs", CompassEmbedder("")),
Err(Error::Invalid(_))
));
store
.text_collection("docs", CompassEmbedder("model/a@1"))
.unwrap()
.upsert_documents([Document::new("id", "east")])
.unwrap();
assert!(matches!(
store.text_collection("docs", CompassEmbedder("model/b@1")),
Err(Error::Invalid(_))
));
}
#[derive(Clone, Debug)]
struct BrokenEmbedder;
impl Embedder for BrokenEmbedder {
fn model_id(&self) -> &str {
"broken@1"
}
fn embed(&self, _: &[String]) -> Result<Vec<Vec<f32>>> {
Ok(Vec::new())
}
}
#[test]
fn malformed_embedder_output_is_rejected_without_creating_a_collection() {
let temp = TempDir::new().unwrap();
let store = open(temp.path().join("documents.git")).unwrap();
let documents = store.text_collection("docs", BrokenEmbedder).unwrap();
assert!(documents
.upsert_documents([Document::new("id", "text")])
.is_err());
assert!(documents.vectors().advanced().is_err());
}