diff --git a/src/bin/alpha_id.rs b/src/bin/alpha_id.rs index 961de10..178e8b0 100644 --- a/src/bin/alpha_id.rs +++ b/src/bin/alpha_id.rs @@ -15,7 +15,7 @@ struct Cli { #[derive(Subcommand)] enum Cmd { /// Build indexes (lexical now; embeddings added in Phase A) - Index { #[arg(long)] sample: Option, #[arg(long)] full: bool, #[arg(long)] confirm: bool }, + Index { #[arg(long)] sample: Option, #[arg(long)] full: bool, #[arg(long)] confirm: bool, #[arg(long)] pack: bool }, /// Suggest ICD codes for a dictation (FILE or - for stdin) Suggest { input: String, @@ -80,9 +80,42 @@ fn main() { let cli = Cli::parse(); let cfg = Config::load(&cli.config).expect("config"); match cli.cmd { - Cmd::Index { sample, full, confirm } => { + Cmd::Index { sample, full, confirm, pack } => { use alpha_id::embed::{EmbeddingStore, build_corpus_embeddings, estimate}; use alpha_id::ionos::IonosClient; + + if pack { + use alpha_id::packed::{corpus_sha256, packed_path, write_packed, PackedHeader}; + 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); + } + 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; + } + let p = Pipeline::load(&cfg).expect("load"); let texts: Vec = p.entries.iter().map(|e| e.text.clone()).collect(); let client = IonosClient::new(&cfg.ionos_base_url, &cfg.token_path) diff --git a/src/lib.rs b/src/lib.rs index 136c336..5d1b4a5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,5 +10,6 @@ pub mod ionos; pub mod embed; pub mod vector; pub mod rerank; +pub mod packed; pub mod pipeline; pub mod eval; diff --git a/src/model.rs b/src/model.rs index 79cb065..217642c 100644 --- a/src/model.rs +++ b/src/model.rs @@ -80,6 +80,14 @@ impl Filter { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Mode { Lexical, Hybrid } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Default)] +pub enum IndexStatus { + Ok, // packed file present and structurally valid + #[default] + Absent, // no usable index — intentional lexical-only + Mismatch(String), // packed file present but structurally invalid; reason carried +} + impl FromStr for Mode { type Err = String; fn from_str(s: &str) -> Result { @@ -118,6 +126,7 @@ pub struct Diagnostics { pub segment_count: usize, pub ionos_calls: usize, pub millis: u128, + pub index_status: IndexStatus, } #[derive(Debug, Serialize)] diff --git a/src/packed.rs b/src/packed.rs new file mode 100644 index 0000000..1e6b302 --- /dev/null +++ b/src/packed.rs @@ -0,0 +1,115 @@ +use crate::model::AppError; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::fs; +use std::path::{Path, PathBuf}; + +const MAGIC: &[u8; 4] = b"AIPK"; +pub const PACKED_FORMAT_VERSION: u32 = 1; + +/// Manifest written into the packed file header. `embed_model` and +/// `corpus_sha256` bind the matrix to the corpus snapshot and model it was +/// built from; iteration 1 records them, iteration 2 enforces them. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PackedHeader { + pub n_rows: u32, + pub dim: u32, + pub embed_model: String, + pub corpus_sha256: String, + pub corpus_name: String, +} + +/// Derived packed-file path `/packed..bin`, using the +/// same `/`→`_` model sanitization `EmbeddingStore::open` applies, so `--pack` +/// and `load` compute the identical path from the same config. +pub fn packed_path(index_dir: &str, model: &str) -> PathBuf { + let safe = model.replace('/', "_"); + PathBuf::from(index_dir).join(format!("packed.{safe}.bin")) +} + +/// SHA-256 of the raw bytes of the corpus file at `path`, lower-case hex. +pub fn corpus_sha256(path: &str) -> Result { + let bytes = fs::read(path).map_err(|e| AppError::Io(format!("{path}: {e}")))?; + let mut h = Sha256::new(); + h.update(&bytes); + Ok(h.finalize().iter().map(|x| format!("{:02x}", x)).collect()) +} + +/// 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> { + let json = serde_json::to_vec(header) + .map_err(|e| AppError::Parse(format!("packed header: {e}")))?; + let mut out = Vec::new(); + out.extend_from_slice(MAGIC); + out.extend_from_slice(&PACKED_FORMAT_VERSION.to_le_bytes()); + out.extend_from_slice(&(json.len() as u32).to_le_bytes()); + out.extend_from_slice(&json); + while out.len() % 4 != 0 { + out.push(0); + } + for row in rows { + for f in row { + out.extend_from_slice(&f.to_le_bytes()); + } + } + let tmp = path.with_extension("tmp"); + fs::write(&tmp, &out).map_err(|e| AppError::Io(format!("{}: {e}", tmp.display())))?; + fs::rename(&tmp, path).map_err(|e| AppError::Io(format!("{}: {e}", path.display())))?; + Ok(()) +} + +/// Read a packed matrix back into its header and row vectors. Fails on a +/// wrong magic, an unsupported (newer) format version, a malformed or +/// truncated header, or a payload length inconsistent with `n_rows × dim`. +pub fn read_packed(path: &Path) -> Result<(PackedHeader, Vec>), AppError> { + let bytes = fs::read(path).map_err(|e| AppError::Io(format!("{}: {e}", path.display())))?; + if bytes.len() < 12 || &bytes[0..4] != MAGIC { + return Err(AppError::Parse(format!("{}: bad magic", path.display()))); + } + let version = u32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]); + if version > PACKED_FORMAT_VERSION { + return Err(AppError::Parse(format!( + "{}: unsupported packed format version {version}", + path.display() + ))); + } + let header_len = u32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]) as usize; + let header_end = 12 + header_len; + if bytes.len() < header_end { + return Err(AppError::Parse(format!("{}: truncated header", path.display()))); + } + let header: PackedHeader = serde_json::from_slice(&bytes[12..header_end]) + .map_err(|e| AppError::Parse(format!("{}: header: {e}", path.display())))?; + let mut payload_start = header_end; + while payload_start % 4 != 0 { + payload_start += 1; + } + let n = header.n_rows as usize; + let dim = header.dim as usize; + let expected = n * dim * 4; + if bytes.len() - payload_start != expected { + return Err(AppError::Parse(format!( + "{}: payload length {} != n_rows×dim×4 {}", + path.display(), + bytes.len() - payload_start, + expected + ))); + } + let mut rows = Vec::with_capacity(n); + let mut off = payload_start; + for _ in 0..n { + let mut row = Vec::with_capacity(dim); + for _ in 0..dim { + row.push(f32::from_le_bytes([ + bytes[off], + bytes[off + 1], + bytes[off + 2], + bytes[off + 3], + ])); + off += 4; + } + rows.push(row); + } + Ok((header, rows)) +} diff --git a/src/pipeline.rs b/src/pipeline.rs index 3274aa3..d634623 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -5,6 +5,7 @@ use crate::fusion::{cross_segment_dedupe, rrf_merge}; use crate::ionos::IonosClient; use crate::lexical::LexicalIndex; use crate::model::*; +use crate::packed; use crate::rerank; use crate::segment::split_segments; use crate::tags::tags_for; @@ -20,6 +21,49 @@ pub struct Pipeline { pub valid_by_code: HashMap, cfg: Config, vector: Option, + index_status: IndexStatus, +} + +/// Build the semantic index, preferring the packed matrix file and falling +/// back to the per-file store. Returns the index (`None` when none is usable) +/// and an `IndexStatus` describing the path taken. +/// +/// Iteration-1 scope: the packed file is validated *structurally* only +/// (`read_packed` checks magic/version/payload length; here we additionally +/// require `n_rows == entries.len()`). The header's `corpus_sha256` and +/// `embed_model` are NOT yet compared against the config — that enforcement +/// is iteration 2. +pub fn load_packed_or_store(cfg: &Config, entries: &[AlphaIdEntry]) -> (Option, IndexStatus) { + let packed = packed::packed_path(&cfg.index_dir, &cfg.embed_model); + if packed.exists() { + match packed::read_packed(&packed) { + Ok((header, rows)) => { + if header.n_rows as usize != entries.len() { + return ( + None, + IndexStatus::Mismatch(format!( + "packed n_rows {} != corpus entries {}", + header.n_rows, + entries.len() + )), + ); + } + return (Some(VectorIndex::from_vectors(rows)), IndexStatus::Ok); + } + Err(e) => return (None, IndexStatus::Mismatch(format!("{e}"))), + } + } + // Fallback: per-file store, all-or-nothing (today's behaviour). + match EmbeddingStore::open(&cfg.index_dir, &cfg.embed_model) { + Ok(store) if entries.iter().all(|e| store.has(&e.text)) => { + let vecs: Option>> = entries.iter().map(|e| store.get(&e.text)).collect(); + match vecs { + Some(v) => (Some(VectorIndex::from_vectors(v)), IndexStatus::Ok), + None => (None, IndexStatus::Absent), + } + } + _ => (None, IndexStatus::Absent), + } } impl Pipeline { @@ -37,17 +81,7 @@ impl Pipeline { } m }; - // Build the semantic index only if EVERY corpus entry already has a - // cached embedding; otherwise leave it `None` so Mode Hybrid degrades to - // Lexical rather than mispairing partial vectors. - let vector = match EmbeddingStore::open(&cfg.index_dir, &cfg.embed_model) { - Ok(store) if entries.iter().all(|e| store.has(&e.text)) => { - let vecs: Option>> = - entries.iter().map(|e| store.get(&e.text)).collect(); - vecs.map(VectorIndex::from_vectors) - } - _ => None, - }; + let (vector, index_status) = load_packed_or_store(cfg, &entries); Ok(Self { entries, meta, @@ -56,6 +90,7 @@ impl Pipeline { valid_by_code, cfg: cfg.clone(), vector, + index_status, }) } @@ -232,6 +267,7 @@ impl Pipeline { segment_count: segments.len(), ionos_calls, millis: t0.elapsed().as_millis(), + index_status: self.index_status.clone(), }, } } diff --git a/tests/pack_cli_tests.rs b/tests/pack_cli_tests.rs new file mode 100644 index 0000000..fe2b84f --- /dev/null +++ b/tests/pack_cli_tests.rs @@ -0,0 +1,122 @@ +use alpha_id::model::{AlphaIdEntry, Config, IndexStatus}; +use alpha_id::pipeline::load_packed_or_store; +use std::path::Path; +use std::process::Command; + +// Write a minimal config TOML pointing the corpus + index dir at temp paths. +// The `--pack` branch only reads `alpha_id_path` (corpus) and `index_dir` +// (store); `claml_path` etc. are present for parse but unused on this path. +fn write_config(dir: &Path, corpus: &Path, index_dir: &Path) -> std::path::PathBuf { + let toml = format!( + r#"ionos_base_url = "http://127.0.0.1:1" +token_path = "/nonexistent" +embed_model = "test/model" +rerank_model = "test/rerank" +alpha_id_path = "{corpus}" +claml_path = "/nonexistent.xml" +index_dir = "{index}" +pool_size = 10 +top_k = 10 +"#, + corpus = corpus.to_str().unwrap(), + index = index_dir.to_str().unwrap(), + ); + let p = dir.join("config.toml"); + std::fs::write(&p, toml).unwrap(); + p +} + +fn corpus_line(text: &str) -> String { + // Alpha-ID format: valid|alpha_id|primary|star|addon|primary2|orpha|text + format!("1|A0001|A00|||||{text}") +} + +#[test] +fn pack_refuses_incomplete_store_exit_2() { + let dir = tempfile::tempdir().unwrap(); + let corpus = dir.path().join("corpus.txt"); + std::fs::write(&corpus, format!("{}\n{}\n", corpus_line("alpha"), corpus_line("beta"))).unwrap(); + let index_dir = dir.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + let config = write_config(dir.path(), &corpus, &index_dir); + + let out = Command::new(env!("CARGO_BIN_EXE_alpha-id")) + .args(["--config", config.to_str().unwrap(), "index", "--pack"]) + .output() + .unwrap(); + assert_eq!(out.status.code(), Some(2), "incomplete store must exit 2"); + let stderr = String::from_utf8_lossy(&out.stderr); + assert!(stderr.contains("store incomplete"), "stderr was: {stderr}"); +} + +#[test] +fn pack_writes_file_from_complete_store() { + let dir = tempfile::tempdir().unwrap(); + let corpus = dir.path().join("corpus.txt"); + std::fs::write(&corpus, format!("{}\n{}\n", corpus_line("alpha"), corpus_line("beta"))).unwrap(); + let index_dir = dir.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + + // Populate the per-file store for both entry texts so the pack succeeds. + let store = alpha_id::embed::EmbeddingStore::open(index_dir.to_str().unwrap(), "test/model").unwrap(); + store.put("alpha", &[1.0f32, 0.0]).unwrap(); + store.put("beta", &[0.0f32, 1.0]).unwrap(); + + let config = write_config(dir.path(), &corpus, &index_dir); + let out = Command::new(env!("CARGO_BIN_EXE_alpha-id")) + .args(["--config", config.to_str().unwrap(), "index", "--pack"]) + .output() + .unwrap(); + assert!(out.status.success(), "stderr: {}", String::from_utf8_lossy(&out.stderr)); + let packed = index_dir.join("packed.test_model.bin"); + assert!(packed.exists(), "packed file should be written at {}", packed.display()); +} + +fn 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(), + } +} + +/// End-to-end seam: a packed file produced by the `index --pack` CLI loads back +/// through the production `load_packed_or_store` path as `IndexStatus::Ok` with a +/// usable index, with no per-file store consulted. This pins the cross-boundary +/// round trip (CLI writer ↔ pipeline reader) that the two halves' unit tests +/// each cover only one side of. +#[test] +fn cli_packed_file_loads_back_as_ok_index() { + let dir = tempfile::tempdir().unwrap(); + let corpus = dir.path().join("corpus.txt"); + std::fs::write(&corpus, format!("{}\n{}\n", corpus_line("alpha"), corpus_line("beta"))).unwrap(); + let index_dir = dir.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + + let store = alpha_id::embed::EmbeddingStore::open(index_dir.to_str().unwrap(), "test/model").unwrap(); + store.put("alpha", &[1.0f32, 0.0]).unwrap(); + store.put("beta", &[0.0f32, 1.0]).unwrap(); + + let config = write_config(dir.path(), &corpus, &index_dir); + let out = Command::new(env!("CARGO_BIN_EXE_alpha-id")) + .args(["--config", config.to_str().unwrap(), "index", "--pack"]) + .output() + .unwrap(); + assert!(out.status.success(), "stderr: {}", String::from_utf8_lossy(&out.stderr)); + + // Load the CLI-written packed file through the production path. The store + // directory is removed first to prove the index comes from the packed file + // alone, not the per-file store. + std::fs::remove_dir_all(index_dir.join("embeddings")).unwrap(); + let mut cfg = Config::load(config.to_str().unwrap()).unwrap(); + cfg.index_dir = index_dir.to_str().unwrap().to_string(); + let entries = vec![entry("alpha"), entry("beta")]; + let (vector, status) = load_packed_or_store(&cfg, &entries); + assert!(vector.is_some(), "packed file written by the CLI must yield an index"); + assert_eq!(status, IndexStatus::Ok); +} diff --git a/tests/packed_load_tests.rs b/tests/packed_load_tests.rs new file mode 100644 index 0000000..b66e0fb --- /dev/null +++ b/tests/packed_load_tests.rs @@ -0,0 +1,73 @@ +use alpha_id::model::{AlphaIdEntry, Config, IndexStatus}; +use alpha_id::packed::{packed_path, write_packed, PackedHeader}; +use alpha_id::pipeline::load_packed_or_store; + +fn entry(text: &str) -> AlphaIdEntry { + AlphaIdEntry { + alpha_id: "X".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(), + } +} + +fn cfg_in(dir: &std::path::Path) -> Config { + let mut cfg = Config::load("config/default.toml").unwrap(); + cfg.index_dir = dir.to_str().unwrap().to_string(); + cfg.embed_model = "test/model".to_string(); + cfg +} + +#[test] +fn index_status_defaults_to_absent() { + assert_eq!(IndexStatus::default(), IndexStatus::Absent); +} + +#[test] +fn index_status_serializes_externally_tagged() { + let j = serde_json::to_string(&IndexStatus::Mismatch("x".to_string())).unwrap(); + assert!(j.contains("Mismatch"), "got {j}"); +} + +#[test] +fn packed_beats_store_builds_without_store() { + let dir = tempfile::tempdir().unwrap(); + let cfg = cfg_in(dir.path()); + let entries = vec![entry("a"), entry("b")]; + let rows = vec![vec![1.0f32, 0.0], vec![0.0f32, 1.0]]; + let h = PackedHeader { + n_rows: 2, + dim: 2, + embed_model: cfg.embed_model.clone(), + corpus_sha256: "irrelevant-in-iter1".to_string(), + corpus_name: "c.txt".to_string(), + }; + write_packed(&packed_path(&cfg.index_dir, &cfg.embed_model), &h, &rows).unwrap(); + // No per-file store populated in `dir`; the index must come from the packed file. + let (vector, status) = load_packed_or_store(&cfg, &entries); + assert!(vector.is_some()); + assert_eq!(status, IndexStatus::Ok); +} + +#[test] +fn n_rows_mismatch_yields_mismatch_none() { + let dir = tempfile::tempdir().unwrap(); + let cfg = cfg_in(dir.path()); + let entries = vec![entry("a"), entry("b"), entry("c")]; // 3 entries + let rows = vec![vec![1.0f32, 0.0], vec![0.0f32, 1.0]]; // packed claims 2 rows + let h = PackedHeader { + n_rows: 2, + dim: 2, + embed_model: cfg.embed_model.clone(), + corpus_sha256: "x".to_string(), + corpus_name: "c.txt".to_string(), + }; + write_packed(&packed_path(&cfg.index_dir, &cfg.embed_model), &h, &rows).unwrap(); + let (vector, status) = load_packed_or_store(&cfg, &entries); + assert!(vector.is_none()); + assert!(matches!(status, IndexStatus::Mismatch(_))); +} diff --git a/tests/packed_tests.rs b/tests/packed_tests.rs new file mode 100644 index 0000000..435fa31 --- /dev/null +++ b/tests/packed_tests.rs @@ -0,0 +1,46 @@ +use alpha_id::packed::{corpus_sha256, read_packed, write_packed, PackedHeader}; +use std::io::Write; + +fn header(n_rows: u32, dim: u32) -> PackedHeader { + PackedHeader { + n_rows, + dim, + embed_model: "test/model".to_string(), + corpus_sha256: "deadbeef".to_string(), + corpus_name: "corpus.txt".to_string(), + } +} + +#[test] +fn write_read_roundtrip_preserves_header_and_rows() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("packed.bin"); + let rows = vec![vec![0.0f32, 1.5, -2.25], vec![3.0, 4.0, 5.0]]; + let h = header(2, 3); + write_packed(&path, &h, &rows).unwrap(); + let (got_h, got_rows) = read_packed(&path).unwrap(); + assert_eq!(got_h, h); + assert_eq!(got_rows, rows); +} + +#[test] +fn corpus_sha256_stable_and_sensitive() { + let dir = tempfile::tempdir().unwrap(); + let p = dir.path().join("c.txt"); + std::fs::write(&p, b"alpha|beta").unwrap(); + let a = corpus_sha256(p.to_str().unwrap()).unwrap(); + let b = corpus_sha256(p.to_str().unwrap()).unwrap(); + assert_eq!(a, b, "same bytes hash identically"); + std::fs::write(&p, b"alpha|betb").unwrap(); + let c = corpus_sha256(p.to_str().unwrap()).unwrap(); + assert_ne!(a, c, "one changed byte changes the hash"); +} + +#[test] +fn read_packed_rejects_bad_magic() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("bad.bin"); + let mut f = std::fs::File::create(&path).unwrap(); + f.write_all(b"XXXX____________").unwrap(); + assert!(read_packed(&path).is_err()); +}