diff --git a/src/bin/alpha_id.rs b/src/bin/alpha_id.rs index b83be50..a57a0ee 100644 --- a/src/bin/alpha_id.rs +++ b/src/bin/alpha_id.rs @@ -85,34 +85,24 @@ fn main() { use alpha_id::ionos::IonosClient; if pack { - use alpha_id::packed::{corpus_sha256, packed_path, write_packed, PackedHeader}; + use alpha_id::packed::{pack_from_store, packed_path, PackOutcome}; let entries = alpha_id::corpus::load(&cfg.alpha_id_path).expect("corpus"); - let store = EmbeddingStore::open(&cfg.index_dir, &cfg.embed_model).expect("store"); - let cached = entries.iter().filter(|e| store.has(&e.text)).count(); - if cached != entries.len() { - eprintln!( - "error: store incomplete: {} of {} entries embedded; \ - run `index --full --confirm` first", - cached, entries.len() - ); - std::process::exit(2); + match pack_from_store(&cfg.index_dir, &cfg.embed_model, &cfg.alpha_id_path, &entries) + .expect("pack") + { + PackOutcome::Incomplete { cached, total } => { + eprintln!( + "error: store incomplete: {} of {} entries embedded; \ + run `index --full --confirm` first", + cached, total + ); + std::process::exit(2); + } + PackOutcome::Packed(header) => { + let path = packed_path(&cfg.index_dir, &cfg.embed_model); + eprintln!("packed {} rows × {} dims → {}", header.n_rows, header.dim, path.display()); + } } - let rows: Vec> = - entries.iter().map(|e| store.get(&e.text).expect("cached vector")).collect(); - let dim = rows.first().map(|r| r.len()).unwrap_or(0) as u32; - let corpus_name = std::path::Path::new(&cfg.alpha_id_path) - .file_name().and_then(|n| n.to_str()) - .unwrap_or(&cfg.alpha_id_path).to_string(); - let header = PackedHeader { - n_rows: entries.len() as u32, - dim, - embed_model: cfg.embed_model.clone(), - corpus_sha256: corpus_sha256(&cfg.alpha_id_path).expect("corpus hash"), - corpus_name, - }; - let path = packed_path(&cfg.index_dir, &cfg.embed_model); - write_packed(&path, &header, &rows).expect("write packed"); - eprintln!("packed {} rows × {} dims → {}", header.n_rows, header.dim, path.display()); return; } diff --git a/src/packed.rs b/src/packed.rs index 1e6b302..99770e4 100644 --- a/src/packed.rs +++ b/src/packed.rs @@ -1,4 +1,5 @@ -use crate::model::AppError; +use crate::embed::EmbeddingStore; +use crate::model::{AlphaIdEntry, AppError}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use std::fs; @@ -35,6 +36,55 @@ pub fn corpus_sha256(path: &str) -> Result { Ok(h.finalize().iter().map(|x| format!("{:02x}", x)).collect()) } +/// Result of a pack attempt: either the header that was written, or a report +/// that the per-file store is missing some corpus entry's embedding (in which +/// case nothing is written, and the caller refuses loudly). +pub enum PackOutcome { + Packed(PackedHeader), + Incomplete { cached: usize, total: usize }, +} + +/// Pack a complete per-file store into the single packed matrix file at +/// `packed_path(index_dir, embed_model)`, returning the header written. Rows +/// are collected in `entries` order (row `i` = `entries[i]`), so the packed +/// matrix is corpus-aligned. If any entry's embedding is not cached, returns +/// `Incomplete { cached, total }` and writes nothing. +pub fn pack_from_store( + index_dir: &str, + embed_model: &str, + alpha_id_path: &str, + entries: &[AlphaIdEntry], +) -> Result { + let store = EmbeddingStore::open(index_dir, embed_model)?; + let cached = entries.iter().filter(|e| store.has(&e.text)).count(); + if cached != entries.len() { + return Ok(PackOutcome::Incomplete { cached, total: entries.len() }); + } + let rows: Vec> = entries + .iter() + .map(|e| { + store + .get(&e.text) + .ok_or_else(|| AppError::Io("cached vector vanished during pack".to_string())) + }) + .collect::, _>>()?; + let dim = rows.first().map(|r| r.len()).unwrap_or(0) as u32; + let corpus_name = Path::new(alpha_id_path) + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or(alpha_id_path) + .to_string(); + let header = PackedHeader { + n_rows: entries.len() as u32, + dim, + embed_model: embed_model.to_string(), + corpus_sha256: corpus_sha256(alpha_id_path)?, + corpus_name, + }; + write_packed(&packed_path(index_dir, embed_model), &header, &rows)?; + Ok(PackOutcome::Packed(header)) +} + /// Write the packed matrix: magic + format_version + header_len + JSON header /// + pad-to-4 + row-major little-endian f32 payload. Atomic via temp+rename. pub fn write_packed(path: &Path, header: &PackedHeader, rows: &[Vec]) -> Result<(), AppError> { diff --git a/tests/packed_tests.rs b/tests/packed_tests.rs index 435fa31..40a7d22 100644 --- a/tests/packed_tests.rs +++ b/tests/packed_tests.rs @@ -1,6 +1,10 @@ use alpha_id::packed::{corpus_sha256, read_packed, write_packed, PackedHeader}; use std::io::Write; +use alpha_id::embed::EmbeddingStore; +use alpha_id::model::AlphaIdEntry; +use alpha_id::packed::{pack_from_store, packed_path, PackOutcome}; + fn header(n_rows: u32, dim: u32) -> PackedHeader { PackedHeader { n_rows, @@ -44,3 +48,79 @@ fn read_packed_rejects_bad_magic() { f.write_all(b"XXXX____________").unwrap(); assert!(read_packed(&path).is_err()); } + +fn mk_entry(text: &str) -> AlphaIdEntry { + AlphaIdEntry { + alpha_id: "A0001".to_string(), + valid: true, + icd_primary: "A00".to_string(), + icd_star: String::new(), + icd_addon: String::new(), + icd_primary2: String::new(), + orpha: String::new(), + text: text.to_string(), + } +} + +#[test] +fn pack_from_store_complete_builds_header_and_writes_file() { + let dir = tempfile::tempdir().unwrap(); + let index_dir = dir.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + let corpus = dir.path().join("corpus.txt"); + std::fs::write(&corpus, b"corpus bytes").unwrap(); + + let store = EmbeddingStore::open(index_dir.to_str().unwrap(), "test/model").unwrap(); + store.put("alpha", &[1.0f32, 0.0, 0.0]).unwrap(); + store.put("beta", &[0.0f32, 1.0, 0.0]).unwrap(); + let entries = vec![mk_entry("alpha"), mk_entry("beta")]; + + let outcome = pack_from_store( + index_dir.to_str().unwrap(), + "test/model", + corpus.to_str().unwrap(), + &entries, + ) + .unwrap(); + + match outcome { + PackOutcome::Packed(h) => { + assert_eq!(h.n_rows, 2); + assert_eq!(h.dim, 3); + assert_eq!(h.embed_model, "test/model"); + assert_eq!( + h.corpus_sha256, + alpha_id::packed::corpus_sha256(corpus.to_str().unwrap()).unwrap() + ); + } + PackOutcome::Incomplete { .. } => panic!("a complete store must pack"), + } + assert!(packed_path(index_dir.to_str().unwrap(), "test/model").exists()); +} + +#[test] +fn pack_from_store_incomplete_reports_counts_and_writes_nothing() { + let dir = tempfile::tempdir().unwrap(); + let index_dir = dir.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + let corpus = dir.path().join("corpus.txt"); + std::fs::write(&corpus, b"corpus bytes").unwrap(); + + let store = EmbeddingStore::open(index_dir.to_str().unwrap(), "test/model").unwrap(); + store.put("alpha", &[1.0f32, 0.0]).unwrap(); // only 1 of 2 entries cached + let entries = vec![mk_entry("alpha"), mk_entry("beta")]; + + let outcome = pack_from_store( + index_dir.to_str().unwrap(), + "test/model", + corpus.to_str().unwrap(), + &entries, + ) + .unwrap(); + + assert!(matches!(outcome, PackOutcome::Incomplete { cached: 1, total: 2 })); + assert!( + !packed_path(index_dir.to_str().unwrap(), "test/model").exists(), + "an incomplete store must not write a packed file" + ); +}