use std::collections::HashMap;
use std::io::{BufWriter, Write};
use std::path::Path;
use serde::{Deserialize, Serialize};
use crate::error::{M1ndError, M1ndResult};
pub const EMBED_CACHE_VERSION: u32 = 1;
#[derive(Serialize, Deserialize)]
pub struct EmbeddingCache {
pub version: u32,
pub model_id: String,
pub dim: u32,
pub entries: HashMap<u64, Box<[f32]>>,
}
impl EmbeddingCache {
pub fn new(model_id: String, dim: u32) -> Self {
Self {
version: EMBED_CACHE_VERSION,
model_id,
dim,
entries: HashMap::new(),
}
}
pub fn load_compatible(path: &Path, model_id: &str, dim: u32) -> Option<Self> {
let bytes = std::fs::read(path).ok()?;
let cache: EmbeddingCache = bincode::deserialize(&bytes).ok()?;
if cache.version != EMBED_CACHE_VERSION || cache.model_id != model_id || cache.dim != dim {
return None;
}
Some(cache)
}
pub fn save(&self, path: &Path) -> M1ndResult<()> {
let bytes =
bincode::serialize(self).map_err(|e| M1ndError::PersistenceFailed(e.to_string()))?;
let temp_path = path.with_extension("tmp");
{
let file = std::fs::File::create(&temp_path)?;
let mut writer = BufWriter::new(file);
writer.write_all(&bytes)?;
writer.flush()?;
}
std::fs::rename(&temp_path, path)?;
Ok(())
}
}
pub fn content_key(model_id: &str, text: &str) -> u64 {
const OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
const PRIME: u64 = 0x0000_0100_0000_01b3;
let mut h = OFFSET;
for &b in model_id.as_bytes() {
h ^= b as u64;
h = h.wrapping_mul(PRIME);
}
h ^= 0xff;
h = h.wrapping_mul(PRIME);
for &b in text.as_bytes() {
h ^= b as u64;
h = h.wrapping_mul(PRIME);
}
h
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn content_key_is_deterministic_and_discriminating() {
let a = content_key("model-x", "graceful shutdown on cancel");
let b = content_key("model-x", "graceful shutdown on cancel");
let c = content_key("model-x", "strawberry cheesecake");
let d = content_key("model-y", "graceful shutdown on cancel");
assert_eq!(
a, b,
"same (model, text) must hash identically across calls"
);
assert_ne!(a, c, "different text must (practically) differ");
assert_ne!(a, d, "different model must (practically) differ");
}
#[test]
fn cache_round_trips_and_validates_identity() {
let tmp = std::env::temp_dir().join(format!(
"m1nd_embcache_test_{}.bin",
content_key("t", "round-trip")
));
let mut cache = EmbeddingCache::new("local/potion-base-8M".to_string(), 256);
let k = content_key("local/potion-base-8M", "hello world");
let vec: Box<[f32]> = vec![0.1f32; 256].into_boxed_slice();
cache.entries.insert(k, vec.clone());
cache.save(&tmp).expect("save");
let loaded = EmbeddingCache::load_compatible(&tmp, "local/potion-base-8M", 256)
.expect("compatible load");
assert_eq!(loaded.entries.len(), 1);
assert_eq!(loaded.entries.get(&k).map(|v| &v[..]), Some(&vec[..]));
assert!(
EmbeddingCache::load_compatible(&tmp, "other/model", 256).is_none(),
"model mismatch must invalidate the cache"
);
assert!(
EmbeddingCache::load_compatible(&tmp, "local/potion-base-8M", 512).is_none(),
"dim mismatch must invalidate the cache"
);
let _ = std::fs::remove_file(&tmp);
}
}