seyk_core/cache/
image_codes.rs

1//! Derived Photo Anchor rows. All writes happen during an opted-in sync.
2use super::*;
3use seyk_embed::quant::PhotoAnchorQ4;
4
5/// Pins the codec format, upstream revision, manifest and packing/rounding rules.
6pub const IMAGE_CODEC_ID: &str = "photo-anchor-q4:v1:be9e066984345c92424c86bfcce3b17f3e3e8484:4dbde655fdd44a1feb4b010685fe065ffbe9c1f996331b4f4a40c9a5b2580c1e:hi-nibble:ties-even";
7
8pub(super) const CREATE: &str = "CREATE TABLE IF NOT EXISTS image_codes (
9    profile TEXT NOT NULL,
10    input_hash TEXT NOT NULL,
11    codec_id TEXT NOT NULL,
12    codes BLOB NOT NULL,
13    PRIMARY KEY(profile,input_hash),
14    FOREIGN KEY(profile,input_hash) REFERENCES embeddings(profile,input_hash) ON DELETE CASCADE
15);";
16
17/// Search representation for one segment. Only standalone images use codes.
18#[derive(Debug, Clone)]
19pub enum ScanData {
20    Exact(Vec<f32>),
21    PhotoAnchorQ4(Vec<u8>),
22}
23
24#[derive(Debug, Clone)]
25pub struct ScanVector {
26    pub resource: ResourceId,
27    pub ordinal: usize,
28    pub segment: CachedSegment,
29    pub data: ScanData,
30}
31
32pub(crate) fn check_profile(profile: &CacheProfile) -> Result<()> {
33    let model = &profile.model;
34    if model.model_id != "google/embeddinggemma-2"
35        || !matches!(model.dimension.get(), 256 | 768)
36        || model.normalization != "mean-pool;truncate;l2"
37        || !matches!(
38            model.variant.as_str(),
39            "onnx-q8" | "onnx-fp32" | "onnx-q8;text-literals-v1" | "onnx-fp32;text-literals-v1"
40        )
41    {
42        return Err(Error::Config("q4 image scan requires EmbeddingGemma 2, d256/d768 and mean-pool;truncate;l2 normalization".into()));
43    }
44    Ok(())
45}
46
47impl SqliteCache {
48    fn image_dimension(&self, id: &str) -> Result<usize> {
49        let (json, dimension): (String, u32) = self.conn.query_row(
50            "SELECT profile,dimension FROM profiles WHERE id=?1",
51            [id],
52            |r| Ok((r.get(0)?, r.get(1)?)),
53        )?;
54        let profile: CacheProfile = serde_json::from_str(&json).map_err(json_err)?;
55        check_profile(&profile)?;
56        if profile.id() != id || profile.model.dimension.get() != dimension as usize {
57            return Err(Error::Config(
58                "corrupt image profile identifier or dimension".into(),
59            ));
60        }
61        Ok(dimension as usize)
62    }
63    pub(super) fn scan_rows(
64        &self,
65        profile: &str,
66        paths: &HashSet<ResourceId>,
67    ) -> Result<Vec<ScanVector>> {
68        let dimension = self.image_dimension(profile)?;
69        let whole = serde_json::to_string(&ResourceLocator::WholeResource).map_err(json_err)?;
70        let mut stmt = self.conn.prepare_cached(
71            "SELECT s.path,s.ordinal,s.locator,s.heading,s.input_hash,
72                CASE WHEN r.modality='image' AND s.locator=?3 AND c.codec_id=?2 THEN c.codes END,
73                CASE WHEN r.modality='image' AND s.locator=?3 AND c.codec_id=?2 THEN NULL ELSE e.vector END
74             FROM segments s JOIN resources r ON r.profile=s.profile AND r.path=s.path
75             JOIN embeddings e ON e.profile=s.profile AND e.input_hash=s.input_hash
76             LEFT JOIN image_codes c ON c.profile=s.profile AND c.input_hash=s.input_hash
77             WHERE s.profile=?1 ORDER BY s.path,s.ordinal")?;
78        let mut rows = stmt.query(params![profile, IMAGE_CODEC_ID, whole])?;
79        let mut out = Vec::new();
80        while let Some(r) = rows.next()? {
81            let resource = key_resource(r.get(0)?)?;
82            if !paths.contains(&resource) {
83                continue;
84            }
85            let codes: Option<Vec<u8>> = r.get(5)?;
86            let data = match codes {
87                Some(codes) => {
88                    if !matches!(dimension, 256 | 768) || codes.len() != dimension / 2 {
89                        return Err(Error::Config(
90                            "corrupt current-codec image_codes row".into(),
91                        ));
92                    }
93                    ScanData::PhotoAnchorQ4(codes)
94                }
95                None => {
96                    let vector = from_bytes(&r.get::<_, Vec<u8>>(6)?)
97                        .map_err(|_| Error::Config("corrupt f32 vector bytes".into()))?;
98                    if vector.len() != dimension || vector.iter().any(|v| !v.is_finite()) {
99                        return Err(Error::Config(
100                            "corrupt f32 vector dimension or values".into(),
101                        ));
102                    }
103                    ScanData::Exact(vector)
104                }
105            };
106            out.push(ScanVector {
107                resource,
108                ordinal: r.get::<_, u32>(1)? as usize,
109                segment: CachedSegment {
110                    locator: serde_json::from_str(&r.get::<_, String>(2)?).map_err(json_err)?,
111                    heading: r.get(3)?,
112                    input_hash: r.get(4)?,
113                },
114                data,
115            });
116        }
117        Ok(out)
118    }
119
120    pub(super) fn backfill_codes(
121        &mut self,
122        profile: &str,
123        paths: &HashSet<ResourceId>,
124    ) -> Result<()> {
125        let codec = PhotoAnchorQ4::new(self.image_dimension(profile)?)
126            .map_err(|e| Error::Config(e.to_string()))?;
127        let tx = self
128            .conn
129            .transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
130        let mut items = HashMap::new();
131        {
132            let whole = serde_json::to_string(&ResourceLocator::WholeResource).map_err(json_err)?;
133            let mut stmt=tx.prepare("SELECT s.path,s.input_hash,CASE WHEN c.codec_id=?3 THEN NULL ELSE e.vector END,c.codec_id,c.codes FROM segments s JOIN resources r ON r.profile=s.profile AND r.path=s.path JOIN embeddings e ON e.profile=s.profile AND e.input_hash=s.input_hash LEFT JOIN image_codes c ON c.profile=s.profile AND c.input_hash=s.input_hash WHERE s.profile=?1 AND r.modality='image' AND s.locator=?2")?;
134            let mut rows = stmt.query(params![profile, whole, IMAGE_CODEC_ID])?;
135            while let Some(r) = rows.next()? {
136                if !paths.contains(&key_resource(r.get(0)?)?) {
137                    continue;
138                }
139                let current: Option<String> = r.get(3)?;
140                if current.as_deref() == Some(IMAGE_CODEC_ID) {
141                    let codes: Vec<u8> = r.get(4)?;
142                    if codes.len() != codec.dim() / 2 {
143                        return Err(Error::Config(
144                            "corrupt current-codec image_codes row".into(),
145                        ));
146                    }
147                    continue;
148                }
149                let hash: String = r.get(1)?;
150                if items.contains_key(&hash) {
151                    continue;
152                }
153                let vector = from_bytes(&r.get::<_, Vec<u8>>(2)?)
154                    .map_err(|_| Error::Config("corrupt f32 vector bytes".into()))?;
155                let embedding = crate::Embedding::from_preserved(vector)
156                    .map_err(|e| Error::Config(e.to_string()))?;
157                let codes = codec
158                    .encode(&embedding)
159                    .map_err(|e| Error::Config(e.to_string()))?;
160                items.insert(hash, codes);
161            }
162        }
163        for (hash, codes) in &items {
164            tx.execute("INSERT OR REPLACE INTO image_codes(profile,input_hash,codec_id,codes) VALUES(?1,?2,?3,?4)",params![profile,hash,IMAGE_CODEC_ID,codes])?;
165        }
166        let generation = if items.is_empty() {
167            None
168        } else {
169            Some(mark_cleanup(&tx, None)?)
170        };
171        tx.commit()?;
172        self.record_write(!items.is_empty(), generation);
173        Ok(())
174    }
175
176    pub(super) fn selected_vectors(
177        &self,
178        profile: &str,
179        hashes: &HashSet<String>,
180    ) -> Result<HashMap<String, Vec<f32>>> {
181        let dimension = self.image_dimension(profile)?;
182        let mut stmt = self
183            .conn
184            .prepare_cached("SELECT vector FROM embeddings WHERE profile=?1 AND input_hash=?2")?;
185        let mut out = HashMap::with_capacity(hashes.len());
186        for hash in hashes {
187            let bytes: Vec<u8> = stmt.query_row(params![profile, hash], |r| r.get(0))?;
188            let values = from_bytes(&bytes)
189                .map_err(|_| Error::Config("corrupt selected f32 vector".into()))?;
190            if values.len() != dimension || values.iter().any(|v| !v.is_finite()) {
191                return Err(Error::Config("corrupt selected f32 vector".into()));
192            }
193            out.insert(hash.clone(), values);
194        }
195        Ok(out)
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use crate::{Dimension, Embedding};
203
204    fn fixture(path: &Path) -> (SqliteCache, String, HashSet<ResourceId>, Vec<f32>) {
205        fixture_with_dimension(path, Dimension::D256)
206    }
207
208    fn fixture_with_dimension(
209        path: &Path,
210        dimension: Dimension,
211    ) -> (SqliteCache, String, HashSet<ResourceId>, Vec<f32>) {
212        let mut cache = SqliteCache::open(path).unwrap();
213        let profile = CacheProfile {
214            model: ModelIdentity {
215                model_id: "google/embeddinggemma-2".into(),
216                revision: "synthetic".into(),
217                variant: "onnx-q8;text-literals-v1".into(),
218                dimension,
219                normalization: "mean-pool;truncate;l2".into(),
220                query_prompt: String::new(),
221                document_prompt: String::new(),
222                max_tokens: 512,
223            },
224            segmenter_version: "test".into(),
225        };
226        let pid = profile.id();
227        cache.register_profile(&pid, &profile).unwrap();
228        let vector = Embedding::from_raw(vec![1.0; dimension.get()], dimension)
229            .unwrap()
230            .as_slice()
231            .to_vec();
232        cache
233            .put_embeddings(&pid, &[("shared".into(), vector.clone())])
234            .unwrap();
235        let mut paths = HashSet::new();
236        for name in ["first", "second"] {
237            let id: ResourceId = format!("test://fixture/{name}").parse().unwrap();
238            let metadata = ResourceMetadata {
239                modality: Modality::Image,
240                media_type: Some("image/jpeg".into()),
241                size_bytes: 1,
242                modified_ns: None,
243                content_hash: "hash".into(),
244            };
245            let segment = CachedSegment {
246                locator: ResourceLocator::WholeResource,
247                heading: None,
248                input_hash: "shared".into(),
249            };
250            cache
251                .put_resource(
252                    &pid,
253                    &id,
254                    &ResourceWrite {
255                        metadata: &metadata,
256                        version: &ResourceVersion::unknown(),
257                        descriptor_hash: Some("descriptor"),
258                        segments: vec![SegmentWrite {
259                            segment: &segment,
260                            lexical_body: "",
261                        }],
262                    },
263                )
264                .unwrap();
265            paths.insert(id);
266        }
267        (cache, pid, paths, vector)
268    }
269
270    #[test]
271    fn backfilled_codes_match_serde_restoration_byte_for_byte() {
272        for dim in [256, 768] {
273            let temp = tempfile::tempdir().unwrap();
274            let (mut cache, pid, paths, _) = fixture_with_dimension(
275                &temp.path().join("cache.sqlite"),
276                Dimension::try_from(dim).unwrap(),
277            );
278            for phase in [0.0, 0.37, 1.5] {
279                let mut vector = Embedding::from_raw(
280                    (0..dim)
281                        .map(|i| {
282                            if i < 2 {
283                                0.0
284                            } else {
285                                (i as f32 * 0.11 + phase).sin()
286                            }
287                        })
288                        .collect(),
289                    Dimension::try_from(dim).unwrap(),
290                )
291                .unwrap()
292                .into_vec();
293                vector[0] = -0.0;
294                vector[1] = f32::from_bits(1);
295                let legacy: Embedding =
296                    serde_json::from_value(serde_json::to_value(&vector).unwrap()).unwrap();
297                let expected = PhotoAnchorQ4::new(dim).unwrap().encode(&legacy).unwrap();
298                cache
299                    .put_embeddings(&pid, &[("shared".into(), vector.clone())])
300                    .unwrap();
301                cache.sync_image_codes(&pid, &paths).unwrap();
302                let actual: Vec<u8> = cache
303                    .conn
304                    .query_row("SELECT codes FROM image_codes", [], |r| r.get(0))
305                    .unwrap();
306                assert_eq!(actual, expected, "dim={dim}, phase={phase}");
307                let retained = cache
308                    .selected_image_vectors(&pid, &HashSet::from(["shared".into()]))
309                    .unwrap();
310                assert_eq!(to_bytes(&retained["shared"]), to_bytes(&vector));
311            }
312        }
313    }
314
315    #[test]
316    fn codes_lifecycle_preserves_f32_and_cached_only_reads_do_not_backfill() {
317        let temp = tempfile::tempdir().unwrap();
318        let path = temp.path().join("cache.sqlite");
319        let (mut cache, pid, paths, vector) = fixture(&path);
320        let original = to_bytes(&vector);
321        assert!(
322            cache
323                .image_scan_rows(&pid, &paths)
324                .unwrap()
325                .iter()
326                .all(|r| matches!(r.data, ScanData::Exact(_)))
327        );
328        assert_eq!(cache.stats().unwrap().profiles[0].image_codes, 0);
329        cache.sync_image_codes(&pid, &paths).unwrap();
330        assert_eq!(cache.stats().unwrap().profiles[0].image_codes, 1);
331        assert!(
332            cache
333                .image_scan_rows(&pid, &paths)
334                .unwrap()
335                .iter()
336                .all(|r| matches!(r.data, ScanData::PhotoAnchorQ4(_)))
337        );
338        cache
339            .put_embeddings(&pid, &[("shared".into(), vector.clone())])
340            .unwrap();
341        assert_eq!(cache.stats().unwrap().profiles[0].image_codes, 1);
342        let selected = cache
343            .selected_image_vectors(&pid, &HashSet::from(["shared".into()]))
344            .unwrap();
345        assert_eq!(to_bytes(&selected["shared"]), original);
346        let one = paths.iter().next().unwrap().clone();
347        cache.remove_resources(&pid, &[one]).unwrap();
348        cache.prune_orphans(&pid).unwrap();
349        assert_eq!(cache.stats().unwrap().profiles[0].image_codes, 1);
350        cache
351            .put_embeddings(
352                &pid,
353                &[("shared".into(), vector.iter().map(|v| -v).collect())],
354            )
355            .unwrap();
356        assert_eq!(cache.stats().unwrap().profiles[0].image_codes, 0);
357        cache.sync_image_codes(&pid, &paths).unwrap();
358        let code: Vec<u8> = cache
359            .conn
360            .query_row("SELECT codes FROM image_codes", [], |r| r.get(0))
361            .unwrap();
362        drop(cache);
363        let mut cache = SqliteCache::open(&path).unwrap();
364        assert_eq!(cache.stats().unwrap().profiles[0].image_codes, 1);
365        cache.clear().unwrap();
366        drop(cache);
367        for suffix in ["", "-wal"] {
368            if let Ok(bytes) = std::fs::read(format!("{}{suffix}", path.display())) {
369                assert!(!bytes.windows(code.len()).any(|w| w == code));
370            }
371        }
372    }
373
374    #[test]
375    fn v6_migration_keeps_vectors_segments_descriptors_and_postings() {
376        let temp = tempfile::tempdir().unwrap();
377        let path = temp.path().join("cache.sqlite");
378        let (cache, pid, paths, vector) = fixture(&path);
379        cache
380            .conn
381            .execute_batch(
382                "DROP TABLE image_codes; UPDATE meta SET value='6' WHERE key='schema_version';",
383            )
384            .unwrap();
385        let before: Vec<(i64, String)> = cache
386            .conn
387            .prepare("SELECT id,locator FROM segments ORDER BY id")
388            .unwrap()
389            .query_map([], |r| Ok((r.get(0)?, r.get(1)?)))
390            .unwrap()
391            .collect::<rusqlite::Result<_>>()
392            .unwrap();
393        drop(cache);
394        let cache = SqliteCache::open(&path).unwrap();
395        let after: Vec<(i64, String)> = cache
396            .conn
397            .prepare("SELECT id,locator FROM segments ORDER BY id")
398            .unwrap()
399            .query_map([], |r| Ok((r.get(0)?, r.get(1)?)))
400            .unwrap()
401            .collect::<rusqlite::Result<_>>()
402            .unwrap();
403        assert_eq!(before, after);
404        assert_eq!(cache.vectors_for(&pid, &paths).unwrap()[0].vector, vector);
405        for id in paths {
406            assert_eq!(
407                cache
408                    .resource(&pid, &id)
409                    .unwrap()
410                    .unwrap()
411                    .descriptor_hash
412                    .as_deref(),
413                Some("descriptor")
414            );
415        }
416        assert_eq!(cache.stats().unwrap().profiles[0].image_codes, 0);
417    }
418
419    #[test]
420    fn unknown_codec_falls_back_but_malformed_current_codec_fails() {
421        let temp = tempfile::tempdir().unwrap();
422        let (mut cache, pid, paths, _) = fixture(&temp.path().join("cache.sqlite"));
423        cache.sync_image_codes(&pid, &paths).unwrap();
424        cache
425            .conn
426            .execute("UPDATE image_codes SET codec_id='future',codes=x'01'", [])
427            .unwrap();
428        assert!(matches!(
429            cache.image_scan_rows(&pid, &paths).unwrap()[0].data,
430            ScanData::Exact(_)
431        ));
432        cache.sync_image_codes(&pid, &paths).unwrap();
433        cache
434            .conn
435            .execute("UPDATE image_codes SET codes=x'01'", [])
436            .unwrap();
437        assert!(cache.image_scan_rows(&pid, &paths).is_err());
438        assert!(cache.sync_image_codes(&pid, &paths).is_err());
439    }
440    #[test]
441    fn q4_profile_corruption_is_rejected_on_reads_and_writes() {
442        let temp = tempfile::tempdir().unwrap();
443        let (mut cache, pid, paths, _) = fixture(&temp.path().join("cache.sqlite"));
444        cache
445            .conn
446            .execute("UPDATE profiles SET dimension=768 WHERE id=?1", [&pid])
447            .unwrap();
448        assert!(cache.image_scan_rows(&pid, &paths).is_err());
449        assert!(cache.sync_image_codes(&pid, &paths).is_err());
450        assert!(
451            cache
452                .selected_image_vectors(&pid, &HashSet::from(["shared".into()]))
453                .is_err()
454        );
455    }
456
457    #[test]
458    fn pdf_pages_with_shared_image_inputs_keep_exact_f32() {
459        let temp = tempfile::tempdir().unwrap();
460        let (mut cache, pid, mut paths, _) = fixture(&temp.path().join("cache.sqlite"));
461        let page: ResourceId = "test://fixture/document.pdf".parse().unwrap();
462        let metadata = ResourceMetadata {
463            modality: Modality::Document,
464            media_type: Some("application/pdf".into()),
465            size_bytes: 1,
466            modified_ns: None,
467            content_hash: "hash".into(),
468        };
469        let segment = CachedSegment {
470            locator: ResourceLocator::Page { page: 1 },
471            heading: None,
472            input_hash: "shared".into(),
473        };
474        cache
475            .put_resource(
476                &pid,
477                &page,
478                &ResourceWrite {
479                    metadata: &metadata,
480                    version: &ResourceVersion::unknown(),
481                    descriptor_hash: Some("pdf"),
482                    segments: vec![SegmentWrite {
483                        segment: &segment,
484                        lexical_body: "",
485                    }],
486                },
487            )
488            .unwrap();
489        paths.insert(page.clone());
490        cache.sync_image_codes(&pid, &paths).unwrap();
491        let rows = cache.image_scan_rows(&pid, &paths).unwrap();
492        assert!(
493            rows.iter()
494                .filter(|r| r.resource == page)
495                .all(|r| matches!(r.data, ScanData::Exact(_)))
496        );
497        assert_eq!(
498            rows.iter()
499                .filter(|r| matches!(r.data, ScanData::PhotoAnchorQ4(_)))
500                .count(),
501            2
502        );
503        assert_eq!(cache.stats().unwrap().profiles[0].image_codes, 1);
504    }
505
506    #[test]
507    fn removing_the_final_segment_reference_prunes_derived_codes() {
508        let temp = tempfile::tempdir().unwrap();
509        let (mut cache, pid, paths, _) = fixture(&temp.path().join("cache.sqlite"));
510        cache.sync_image_codes(&pid, &paths).unwrap();
511        cache
512            .remove_resources(&pid, &paths.into_iter().collect::<Vec<_>>())
513            .unwrap();
514        assert_eq!(cache.prune_orphans(&pid).unwrap(), 1);
515        let stats = cache.stats().unwrap();
516        assert_eq!(stats.profiles[0].embeddings, 0);
517        assert_eq!(stats.profiles[0].image_codes, 0);
518    }
519}
seyk_core/cache/image_codes.rs

seyk_core/cache/
image_codes.rs
