33 lines
1.4 KiB
Rust
33 lines
1.4 KiB
Rust
use alpha_id::embed::{EmbeddingStore, build_corpus_embeddings};
|
|
use alpha_id::ionos::IonosClient;
|
|
use std::io::Write;
|
|
|
|
#[test]
|
|
fn build_embeds_only_missing_and_is_resumable() {
|
|
let server = tiny_http::Server::http("127.0.0.1:0").unwrap();
|
|
let url = format!("http://{}", server.server_addr());
|
|
let h = std::thread::spawn(move || {
|
|
// Only 1 HTTP request is made: both texts fit in a single batch (batch=8).
|
|
// The second build_corpus_embeddings call finds everything cached and
|
|
// makes zero requests, so the server only needs to serve once.
|
|
for _ in 0..1 {
|
|
if let Ok(req) = server.recv() {
|
|
let body = r#"{"data":[{"embedding":[0.1,0.2]},{"embedding":[0.3,0.4]}]}"#;
|
|
req.respond(tiny_http::Response::from_string(body)).ok();
|
|
}
|
|
}
|
|
});
|
|
let mut tf = tempfile::NamedTempFile::new().unwrap();
|
|
write!(tf, "tok").unwrap();
|
|
let client = IonosClient::new(&url, tf.path().to_str().unwrap()).unwrap();
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let store = EmbeddingStore::open(dir.path().to_str().unwrap(), "BAAI/bge-m3").unwrap();
|
|
let texts = vec!["a".to_string(), "b".to_string()];
|
|
|
|
let n1 = build_corpus_embeddings(&client, &store, "BAAI/bge-m3", &texts, 8).unwrap();
|
|
assert_eq!(n1, 2);
|
|
let n2 = build_corpus_embeddings(&client, &store, "BAAI/bge-m3", &texts, 8).unwrap();
|
|
assert_eq!(n2, 0);
|
|
h.join().ok();
|
|
}
|