tawasuyu/brahman

1//! Cache de embeddings keyed por sha256 del contenido + model_id.
2//!
3//! Razón de existir: el modelo real (`fastembed-allMiniLML6V2`) es
4//! caro (1-50 ms por archivo según tamaño y CPU). Cada vez que el
5//! daemon de nouser re-publica una Mónada o el watcher dispara un
6//! re-cluster por cambio de FS, todos los archivos pasan otra vez
7//! por embed. Para árboles de 1000 archivos, eso son segundos
8//! desperdiciados re-embedidando contenido que no cambió.
9//!
10//! ## Diseño
11//!
12//! - **Cache key**: `sha256(bytes que el modelo realmente vio)` +
13//!   `MODEL_ID` (string). Usar el sha de los bytes-vistos garantiza
14//!   que la cache no devuelva un embedding de contenido viejo
15//!   simplemente porque el path no cambió.
16//! - **Cache value**: el `Vec<f32>` serializado como bytes
17//!   little-endian (4 bytes por f32). Compacto, sin overhead de
18//!   bincode/postcard para datos numéricos puros.
19//! - **Backend**: sled, tree único `embed_cache_v1`. Path:
20//!   `$XDG_CACHE_HOME/brahman/nouser-nous-real-embed-cache.sled`.
21//!
22//! ## Versionado
23//!
24//! El nombre del tree (`embed_cache_v1`) es el "schema version" del
25//! formato value. Si bumpeamos a (p. ej.) almacenar también el
26//! tiempo de cómputo o el ONNX session id, creamos `embed_cache_v2`
27//! y el viejo queda como dato muerto que sled puede limpiar.
28//!
29//! El `MODEL_ID` viaja dentro del key, así que cambiar de modelo
30//! invalida implícitamente las entradas viejas (no se accede más
31//! a esos keys; sled las mantiene hasta GC manual).
32
33use std::path::PathBuf;
34
35/// Wrapper sobre sled::Db con la API justa que necesita `handle_file`.
36#[derive(Clone)]
37pub struct EmbedCache {
38    tree: sled::Tree,
39}
40
41impl EmbedCache {
42    /// Abre (o crea) la cache en su path canónico. El sled::Db queda
43    /// referenciado por el Tree; mientras `EmbedCache` viva, el DB
44    /// vive.
45    pub fn open() -> Result<Self, sled::Error> {
46        let path = default_path();
47        if let Some(parent) = path.parent() {
48            // best-effort: si no podemos crear el dir, sled falla con
49            // mensaje específico abajo.
50            let _ = std::fs::create_dir_all(parent);
51        }
52        let db = sled::open(&path)?;
53        let tree = db.open_tree("embed_cache_v1")?;
54        Ok(Self { tree })
55    }
56
57    /// Variante para tests: cache efímera bajo `dir`.
58    #[cfg(test)]
59    pub fn open_at(dir: &std::path::Path) -> Result<Self, sled::Error> {
60        let db = sled::open(dir)?;
61        let tree = db.open_tree("embed_cache_v1")?;
62        Ok(Self { tree })
63    }
64
65    /// Lookup. `None` si miss; `Some(vec)` si hit.
66    pub fn get(&self, file_sha: &[u8; 32], model_id: &str) -> Option<Vec<f32>> {
67        let key = build_key(file_sha, model_id);
68        let bytes = self.tree.get(&key).ok()??;
69        decode_embedding(&bytes)
70    }
71
72    /// Almacena. Errores se loggean pero no propagan — cache miss es
73    /// recuperable, no querés tirar el embed válido por fallo de I/O
74    /// de cache.
75    pub fn put(&self, file_sha: &[u8; 32], model_id: &str, embedding: &[f32]) {
76        let key = build_key(file_sha, model_id);
77        let bytes = encode_embedding(embedding);
78        if let Err(e) = self.tree.insert(key, bytes) {
79            tracing::warn!(error = %e, "embed-cache put falló (no-fatal)");
80        }
81    }
82
83    /// Cantidad actual de entradas (best-effort para logs).
84    pub fn len(&self) -> usize {
85        self.tree.len()
86    }
87}
88
89/// Path default. Honra `XDG_CACHE_HOME`, cae a `$HOME/.cache`, y de
90/// último recurso a `/tmp` (sin persistencia, pero al menos no
91/// crashea en entornos minimalistas como CI sin HOME).
92fn default_path() -> PathBuf {
93    if let Ok(p) = std::env::var("NOUSER_NOUS_REAL_CACHE") {
94        return PathBuf::from(p);
95    }
96    let base = std::env::var("XDG_CACHE_HOME")
97        .ok()
98        .map(PathBuf::from)
99        .or_else(|| {
100            std::env::var("HOME")
101                .ok()
102                .map(|h| PathBuf::from(h).join(".cache"))
103        })
104        .unwrap_or_else(std::env::temp_dir);
105    base.join("brahman").join("nouser-nous-real-embed-cache.sled")
106}
107
108fn build_key(file_sha: &[u8; 32], model_id: &str) -> Vec<u8> {
109    let mut k = Vec::with_capacity(32 + 1 + model_id.len());
110    k.extend_from_slice(file_sha);
111    // separator byte para que prefijos de model_id no choquen con
112    // bytes del sha (improbable pero barato).
113    k.push(0xff);
114    k.extend_from_slice(model_id.as_bytes());
115    k
116}
117
118fn encode_embedding(v: &[f32]) -> Vec<u8> {
119    let mut out = Vec::with_capacity(v.len() * 4);
120    for f in v {
121        out.extend_from_slice(&f.to_le_bytes());
122    }
123    out
124}
125
126fn decode_embedding(bytes: &[u8]) -> Option<Vec<f32>> {
127    if bytes.len() % 4 != 0 {
128        return None;
129    }
130    let mut out = Vec::with_capacity(bytes.len() / 4);
131    for chunk in bytes.chunks_exact(4) {
132        out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
133    }
134    Some(out)
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    fn sha(s: &[u8]) -> [u8; 32] {
142        ente_cas::sha256_of(s)
143    }
144
145    #[test]
146    fn roundtrip_returns_same_vector() {
147        let dir = tempfile::tempdir().unwrap();
148        let cache = EmbedCache::open_at(dir.path()).unwrap();
149        let key = sha(b"hello world");
150        let v = vec![0.1f32, -0.5, 1.0, 3.14159];
151        cache.put(&key, "real-fastembed-allMiniLML6V2-384d", &v);
152        let got = cache
153            .get(&key, "real-fastembed-allMiniLML6V2-384d")
154            .expect("hit esperado");
155        assert_eq!(got, v);
156    }
157
158    #[test]
159    fn miss_returns_none() {
160        let dir = tempfile::tempdir().unwrap();
161        let cache = EmbedCache::open_at(dir.path()).unwrap();
162        let key = sha(b"never stored");
163        assert!(cache.get(&key, "real-fastembed-allMiniLML6V2-384d").is_none());
164    }
165
166    #[test]
167    fn different_models_do_not_collide() {
168        let dir = tempfile::tempdir().unwrap();
169        let cache = EmbedCache::open_at(dir.path()).unwrap();
170        let key = sha(b"same content");
171        cache.put(&key, "model-a", &[1.0, 2.0]);
172        cache.put(&key, "model-b", &[7.0, 8.0]);
173        assert_eq!(cache.get(&key, "model-a").unwrap(), vec![1.0, 2.0]);
174        assert_eq!(cache.get(&key, "model-b").unwrap(), vec![7.0, 8.0]);
175    }
176
177    #[test]
178    fn different_content_different_keys() {
179        let dir = tempfile::tempdir().unwrap();
180        let cache = EmbedCache::open_at(dir.path()).unwrap();
181        let k1 = sha(b"abc");
182        let k2 = sha(b"abd");
183        cache.put(&k1, "m", &[1.0]);
184        assert!(cache.get(&k2, "m").is_none());
185    }
186
187    #[test]
188    fn corrupted_value_returns_none() {
189        // Si sled devuelve bytes con length no múltiplo de 4, decode
190        // debe fallar limpio (None) en vez de panicar.
191        let dir = tempfile::tempdir().unwrap();
192        let cache = EmbedCache::open_at(dir.path()).unwrap();
193        let key = sha(b"x");
194        // Insertamos manualmente bytes inválidos.
195        let raw_key = build_key(&key, "m");
196        cache.tree.insert(raw_key, &[1u8, 2, 3][..]).unwrap();
197        assert!(cache.get(&key, "m").is_none());
198    }
199}