mod builder;
mod pgvector;
pub use builder::*;
pub use pgvector::*;
const PG_LOCK_ID_EMBEDDING_TABLE: i64 = 1573678846307946494;
const PG_LOCK_ID_COLLECTION_TABLE: i64 = 1573678846307946495;
const PG_LOCKID_EXTENSION: i64 = 1573678846307946496;
#[cfg(test)]
mod tests {
use crate::{
schemas::Document,
test_utils::FakeEmbedder,
vectorstore::{pgvector::StoreBuilder, VecStoreOptions, VectorStore},
};
#[tokio::test]
#[ignore]
async fn pgvector_similarity_search_returns_closest_first() {
let url = match std::env::var("PGVECTOR_TEST_URL") {
Ok(url) => url,
Err(_) => {
let out = std::process::Command::new("docker")
.args([
"run",
"-d",
"--rm",
"-e",
"POSTGRES_PASSWORD=pw",
"-p",
"0:5432",
"pgvector/pgvector:pg17",
])
.output()
.expect("docker");
let id = String::from_utf8(out.stdout).unwrap().trim().to_string();
let port = std::process::Command::new("docker")
.args(["port", &id, "5432/tcp"])
.output()
.unwrap();
let port = String::from_utf8(port.stdout).unwrap();
let port = port
.lines()
.next()
.unwrap()
.rsplit(':')
.next()
.unwrap()
.to_string();
format!("postgres://postgres:pw@127.0.0.1:{port}/postgres")
}
};
let embedder = FakeEmbedder::new(
3,
&[
("near", vec![1.0, 0.0, 0.0]),
("mid", vec![1.0, 1.0, 0.0]),
("far", vec![-1.0, 0.0, 0.0]),
("query", vec![1.0, 0.05, 0.0]),
],
);
let mut store = None;
for _ in 0..60 {
match StoreBuilder::new()
.embedder(embedder.clone())
.connection_url(&url)
.vector_dimensions(3)
.collection_name("order_test")
.pre_delete_collection(true)
.build()
.await
{
Ok(s) => {
store = Some(s);
break;
}
Err(_) => tokio::time::sleep(std::time::Duration::from_millis(500)).await,
}
}
let store = store.expect("postgres did not come up");
let docs = ["far", "mid", "near"].map(Document::new);
store
.add_documents(&docs, &VecStoreOptions::default())
.await
.unwrap();
let found = store
.similarity_search("query", 2, &VecStoreOptions::default())
.await
.unwrap();
let contents: Vec<_> = found.iter().map(|d| d.page_content.as_str()).collect();
assert_eq!(contents, ["near", "mid"]);
assert!(found[0].score <= found[1].score, "score is cosine distance");
}
}