feat: cost-guarded resumable index build (sample smoke + full)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
+41
-2
@@ -52,9 +52,48 @@ fn main() {
|
|||||||
let cfg = Config::load(&cli.config).expect("config");
|
let cfg = Config::load(&cli.config).expect("config");
|
||||||
match cli.cmd {
|
match cli.cmd {
|
||||||
Cmd::Index { sample, full, confirm } => {
|
Cmd::Index { sample, full, confirm } => {
|
||||||
|
use alpha_id::embed::{EmbeddingStore, build_corpus_embeddings, estimate};
|
||||||
|
use alpha_id::ionos::IonosClient;
|
||||||
let p = Pipeline::load(&cfg).expect("load");
|
let p = Pipeline::load(&cfg).expect("load");
|
||||||
eprintln!("Loaded {} entries, {} ICD metas. (sample={:?} full={} confirm={})",
|
let texts: Vec<String> = p.entries.iter().map(|e| e.text.clone()).collect();
|
||||||
p.entries.len(), p.meta.len(), sample, full, confirm);
|
let client = IonosClient::new(&cfg.ionos_base_url, &cfg.token_path)
|
||||||
|
.expect("ionos client (token?)");
|
||||||
|
let store = EmbeddingStore::open(&cfg.index_dir, &cfg.embed_model).expect("store");
|
||||||
|
|
||||||
|
let subset: Vec<String> = match sample {
|
||||||
|
Some(n) => texts.iter().take(n).cloned().collect(),
|
||||||
|
None => texts.clone(),
|
||||||
|
};
|
||||||
|
|
||||||
|
if sample.is_some() {
|
||||||
|
let est = estimate(&subset);
|
||||||
|
eprintln!("SMOKE: embedding {} sample texts ({} chars)…", est.items, est.chars);
|
||||||
|
let n = build_corpus_embeddings(&client, &store, &cfg.embed_model, &subset, 32)
|
||||||
|
.expect("sample embed failed");
|
||||||
|
let dim = store.get(&subset[0]).map(|v| v.len()).unwrap_or(0);
|
||||||
|
eprintln!("SMOKE OK: embedded {n}, vector dim = {dim}. \
|
||||||
|
Inspect, then run `index --full --confirm`.");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
if full {
|
||||||
|
let pending: Vec<String> = texts.iter().filter(|t| !store.has(t)).cloned().collect();
|
||||||
|
let est = estimate(&pending);
|
||||||
|
eprintln!("FULL RUN: {} of {} texts not yet embedded ({} chars). \
|
||||||
|
This calls IONOS and costs money.", est.items, texts.len(), est.chars);
|
||||||
|
if !confirm {
|
||||||
|
eprintln!("Refusing without --confirm. Re-run: \
|
||||||
|
`alpha-id index --full --confirm`");
|
||||||
|
std::process::exit(2);
|
||||||
|
}
|
||||||
|
let n = build_corpus_embeddings(&client, &store, &cfg.embed_model, &texts, 32)
|
||||||
|
.expect("full embed failed");
|
||||||
|
eprintln!("FULL RUN done: {n} newly embedded, {} total cached.",
|
||||||
|
texts.iter().filter(|t| store.has(t)).count());
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
eprintln!("Specify --sample N (smoke test first) or --full --confirm.");
|
||||||
}
|
}
|
||||||
Cmd::Suggest { input, mode, top, billable_only, valid_only, exclude_exotic, chapters, json } => {
|
Cmd::Suggest { input, mode, top, billable_only, valid_only, exclude_exotic, chapters, json } => {
|
||||||
let p = Pipeline::load(&cfg).expect("load");
|
let p = Pipeline::load(&cfg).expect("load");
|
||||||
|
|||||||
@@ -62,6 +62,23 @@ fn hex(b: &[u8]) -> String {
|
|||||||
b.iter().map(|x| format!("{:02x}", x)).collect()
|
b.iter().map(|x| format!("{:02x}", x)).collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Embed all `texts` not yet cached, in batches. Resumable: a crash
|
||||||
|
/// leaves cached vectors intact; rerun continues from the gap.
|
||||||
|
pub fn build_corpus_embeddings(
|
||||||
|
client: &IonosClient, store: &EmbeddingStore, model: &str,
|
||||||
|
texts: &[String], batch: usize,
|
||||||
|
) -> Result<usize, AppError> {
|
||||||
|
let pending: Vec<String> = texts.iter()
|
||||||
|
.filter(|t| !store.has(t))
|
||||||
|
.cloned()
|
||||||
|
.collect();
|
||||||
|
let mut done = 0;
|
||||||
|
for chunk in pending.chunks(batch.max(1)) {
|
||||||
|
done += embed_batch(client, store, model, chunk)?;
|
||||||
|
}
|
||||||
|
Ok(done)
|
||||||
|
}
|
||||||
|
|
||||||
/// Embed a batch via IONOS, persisting each result. Only `texts` not
|
/// Embed a batch via IONOS, persisting each result. Only `texts` not
|
||||||
/// already cached should be passed (caller filters via `has`).
|
/// already cached should be passed (caller filters via `has`).
|
||||||
pub fn embed_batch(
|
pub fn embed_batch(
|
||||||
|
|||||||
@@ -0,0 +1,32 @@
|
|||||||
|
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();
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user