diff --git a/src/embed.rs b/src/embed.rs index 4b0fd9d..c80f7e3 100644 --- a/src/embed.rs +++ b/src/embed.rs @@ -1 +1,84 @@ -// implemented in a later task +use crate::ionos::IonosClient; +use crate::model::AppError; +use sha2::{Digest, Sha256}; +use std::fs; +use std::path::PathBuf; + +pub struct EmbeddingStore { + dir: PathBuf, // /embeddings// + model: String, +} + +pub struct Estimate { pub items: usize, pub chars: usize } + +pub fn estimate(pending: &[String]) -> Estimate { + Estimate { items: pending.len(), chars: pending.iter().map(|s| s.len()).sum() } +} + +impl EmbeddingStore { + pub fn key(model: &str, text: &str) -> String { + let mut h = Sha256::new(); + h.update(model.as_bytes()); + h.update([0u8]); + h.update(text.as_bytes()); + hex(&h.finalize()) + } + + pub fn open(index_dir: &str, model: &str) -> Result { + let safe = model.replace('/', "_"); + let dir = PathBuf::from(index_dir).join("embeddings").join(safe); + fs::create_dir_all(&dir) + .map_err(|e| AppError::Io(format!("{}: {e}", dir.display())))?; + Ok(Self { dir, model: model.to_string() }) + } + + fn path(&self, text: &str) -> PathBuf { + self.dir.join(format!("{}.f32", Self::key(&self.model, text))) + } + + pub fn has(&self, text: &str) -> bool { self.path(text).exists() } + + pub fn get(&self, text: &str) -> Option> { + let bytes = fs::read(self.path(text)).ok()?; + Some(bytes.chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect()) + } + + /// Atomic write (temp + rename) so an interrupted run never leaves + /// a half-written vector — the next run simply re-embeds that key. + pub fn put(&self, text: &str, v: &[f32]) -> Result<(), AppError> { + let mut bytes = Vec::with_capacity(v.len() * 4); + for f in v { bytes.extend_from_slice(&f.to_le_bytes()); } + let final_path = self.path(text); + let tmp = final_path.with_extension("tmp"); + fs::write(&tmp, &bytes).map_err(|e| AppError::Io(format!("{e}")))?; + fs::rename(&tmp, &final_path).map_err(|e| AppError::Io(format!("{e}")))?; + Ok(()) + } +} + +fn hex(b: &[u8]) -> String { + b.iter().map(|x| format!("{:02x}", x)).collect() +} + +/// Embed a batch via IONOS, persisting each result. Only `texts` not +/// already cached should be passed (caller filters via `has`). +pub fn embed_batch( + client: &IonosClient, store: &EmbeddingStore, model: &str, texts: &[String], +) -> Result { + if texts.is_empty() { return Ok(0); } + let body = serde_json::json!({ "model": model, "input": texts }); + let resp: serde_json::Value = client.post_json("/embeddings", &body)?; + let data = resp["data"].as_array() + .ok_or_else(|| AppError::Ionos("embeddings: no data[]".into()))?; + let mut n = 0; + for (i, item) in data.iter().enumerate() { + let v: Vec = item["embedding"].as_array() + .ok_or_else(|| AppError::Ionos("embeddings: no embedding[]".into()))? + .iter().map(|x| x.as_f64().unwrap_or(0.0) as f32).collect(); + store.put(&texts[i], &v)?; + n += 1; + } + Ok(n) +} diff --git a/tests/embed_tests.rs b/tests/embed_tests.rs new file mode 100644 index 0000000..43bbf80 --- /dev/null +++ b/tests/embed_tests.rs @@ -0,0 +1,28 @@ +use alpha_id::embed::{EmbeddingStore, estimate}; + +#[test] +fn cache_key_is_stable_sha256() { + let k1 = EmbeddingStore::key("BAAI/bge-m3", "diabetes"); + let k2 = EmbeddingStore::key("BAAI/bge-m3", "diabetes"); + assert_eq!(k1, k2); + assert_ne!(k1, EmbeddingStore::key("BAAI/bge-m3", "hypertonie")); + assert_eq!(k1.len(), 64); // hex sha256 +} + +#[test] +fn store_persists_and_reloads_vectors_resumably() { + let dir = tempfile::tempdir().unwrap(); + let store = EmbeddingStore::open(dir.path().to_str().unwrap(), "BAAI/bge-m3").unwrap(); + assert!(store.get("diabetes").is_none()); + store.put("diabetes", &[0.1, 0.2, 0.3]).unwrap(); + let reopened = EmbeddingStore::open(dir.path().to_str().unwrap(), "BAAI/bge-m3").unwrap(); + assert_eq!(reopened.get("diabetes").unwrap(), vec![0.1, 0.2, 0.3]); + assert!(reopened.has("diabetes")); // resume: skip already-embedded +} + +#[test] +fn estimate_reports_pending_count_and_chars() { + let est = estimate(&["abc".to_string(), "de".to_string()]); + assert_eq!(est.items, 2); + assert_eq!(est.chars, 5); +}