scone_core/index/
vectors.rs

1//! Vector index over chunk embeddings (usearch HNSW, cosine).
2//!
3//! Derived data: rebuildable from the embedding blobs in SQLite. The
4//! dimension is pinned at open; a mismatch is a typed error pointing at
5//! `doctor --rebuild` (spec §9), never a silent re-index.
6
7use std::path::{Path, PathBuf};
8
9use usearch::{Index, IndexOptions, MetricKind, ScalarKind};
10
11use crate::error::{Result, SconeError};
12
13const FILE_NAME: &str = "vectors.usearch";
14
15pub struct VectorIndex {
16    index: Index,
17    path: PathBuf,
18    dim: usize,
19}
20
21fn ix(e: impl std::fmt::Display) -> SconeError {
22    SconeError::Index(format!("vectors: {e}"))
23}
24
25fn new_index(dim: usize) -> Result<Index> {
26    let options = IndexOptions {
27        dimensions: dim,
28        metric: MetricKind::Cos,
29        quantization: ScalarKind::F32,
30        ..Default::default()
31    };
32    usearch::new_index(&options).map_err(ix)
33}
34
35impl std::fmt::Debug for VectorIndex {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        f.debug_struct("VectorIndex")
38            .field("path", &self.path)
39            .field("dim", &self.dim)
40            .field("size", &self.index.size())
41            .finish()
42    }
43}
44
45impl VectorIndex {
46    /// Open, resetting the stored file on dimension mismatch. Only the
47    /// repair path uses this; normal opens refuse mismatches loudly.
48    pub fn open_or_reset(dir: &Path, dim: usize) -> Result<Self> {
49        match Self::open(dir, dim) {
50            Err(SconeError::Index(_)) => {
51                std::fs::remove_file(dir.join(FILE_NAME))?;
52                Self::open(dir, dim)
53            }
54            other => other,
55        }
56    }
57
58    pub fn open(dir: &Path, dim: usize) -> Result<Self> {
59        std::fs::create_dir_all(dir)?;
60        let path = dir.join(FILE_NAME);
61        let index = new_index(dim)?;
62        if path.exists() {
63            let p = path
64                .to_str()
65                .ok_or_else(|| ix("index path is not valid UTF-8"))?;
66            index.load(p).map_err(ix)?;
67            let stored = index.dimensions();
68            if stored != dim {
69                return Err(SconeError::Index(format!(
70                    "vectors: stored dimension {stored} != embedder dimension {dim}; \
71                     run `scone doctor --rebuild`"
72                )));
73            }
74        }
75        Ok(Self { index, path, dim })
76    }
77
78    /// Stage vectors in memory (idempotent: existing keys are skipped, so
79    /// catch-up re-adds are safe). Durable only after [`VectorIndex::flush`]
80    /// — saving the whole file per note was the 96ms/ingest bottleneck.
81    pub fn add(&mut self, rows: &[(u64, &[f32])]) -> Result<()> {
82        let needed = self.index.size() + rows.len();
83        if self.index.capacity() < needed {
84            self.index.reserve(needed).map_err(ix)?;
85        }
86        for (key, vector) in rows {
87            if vector.len() != self.dim {
88                return Err(SconeError::Index(format!(
89                    "vectors: vector for key {key} has dimension {} != {}",
90                    vector.len(),
91                    self.dim
92                )));
93            }
94            if self.index.contains(*key) {
95                continue;
96            }
97            self.index.add(*key, vector).map_err(ix)?;
98        }
99        Ok(())
100    }
101
102    /// Persist the staged index to disk.
103    pub fn flush(&self) -> Result<()> {
104        self.save()
105    }
106
107    /// Top-`k` `(chunk_id, cosine_similarity)`, unfiltered by space —
108    /// callers post-filter against truth (memory/bugs.md P-3).
109    pub fn search(&self, query: &[f32], k: usize) -> Result<Vec<(u64, f32)>> {
110        if self.index.size() == 0 {
111            return Ok(Vec::new());
112        }
113        let matches = self.index.search(query, k.max(1)).map_err(ix)?;
114        Ok(matches
115            .keys
116            .iter()
117            .zip(&matches.distances)
118            .map(|(key, dist)| (*key, 1.0 - dist))
119            .collect())
120    }
121
122    pub fn wipe(&mut self) -> Result<()> {
123        self.index = new_index(self.dim)?;
124        self.save()
125    }
126
127    fn save(&self) -> Result<()> {
128        let p = self
129            .path
130            .to_str()
131            .ok_or_else(|| ix("index path is not valid UTF-8"))?;
132        self.index.save(p).map_err(ix)
133    }
134}
scone_core/index/vectors.rs

scone_core/index/
vectors.rs
