remem/retrieval/
vector.rs

1use anyhow::{Context, Result};
2use rusqlite::{params, Connection, OptionalExtension, Statement};
3
4use super::embedding::TextEmbedding;
5pub use super::vector_candidates::VECTOR_SEARCH_CANDIDATE_LIMIT;
6
7mod backfill;
8mod coverage;
9mod reindex;
10mod vec_index;
11
12pub(crate) use vec_index::ensure_vec_index;
13
14pub use super::embedding::{
15    LOCAL_EMBEDDING_DIMENSIONS as EMBEDDING_DIMENSIONS,
16    LOCAL_EMBEDDING_MODEL as DEFAULT_EMBEDDING_MODEL,
17};
18pub use backfill::{
19    backfill_missing_memory_embeddings, pending_memory_embedding_count,
20    pending_memory_embedding_reindex_count, pending_memory_embedding_reindex_count_for_target,
21    reindex_memory_embeddings, reindex_memory_embeddings_with_report, EmbeddingReindexReport,
22};
23pub(crate) use backfill::{
24    reindex_memory_embeddings_with_session_report, EmbeddingBackfillSession,
25};
26pub use coverage::{
27    active_embedding_coverage, active_embedding_coverage_for_status,
28    active_embedding_coverage_for_target, prune_inactive_memory_embeddings,
29    ActiveEmbeddingCoverage, InactiveEmbeddingPruneReport,
30};
31
32const EMBEDDING_REINDEX_WRITE_BATCH_SIZE: usize = 512;
33const UPSERT_EMBEDDING_SQL: &str = "INSERT INTO memory_embeddings
34         (memory_id, embedding, dimensions, model, content_hash, updated_at_epoch)
35         VALUES (?1, ?2, ?3, ?4, ?5, ?6)
36         ON CONFLICT(memory_id, model, dimensions) DO UPDATE SET
37             embedding = excluded.embedding,
38             content_hash = excluded.content_hash,
39             updated_at_epoch = excluded.updated_at_epoch";
40
41#[derive(Debug, Clone, PartialEq)]
42pub struct VectorHit {
43    pub memory_id: i64,
44    pub distance: f32,
45}
46
47#[derive(Debug, Clone, PartialEq)]
48pub struct VectorSearchOutcome {
49    pub hits: Vec<VectorHit>,
50    pub disabled_reason: Option<String>,
51    pub candidates_scanned: usize,
52    pub timings: Vec<crate::perf::PhaseTiming>,
53}
54
55impl VectorSearchOutcome {
56    pub fn disabled(reason: impl Into<String>) -> Self {
57        Self {
58            hits: vec![],
59            disabled_reason: Some(reason.into()),
60            candidates_scanned: 0,
61            timings: vec![],
62        }
63    }
64
65    fn disabled_with_timings(
66        reason: impl Into<String>,
67        timings: Vec<crate::perf::PhaseTiming>,
68    ) -> Self {
69        Self {
70            hits: vec![],
71            disabled_reason: Some(reason.into()),
72            candidates_scanned: 0,
73            timings,
74        }
75    }
76
77    pub fn ready(hits: Vec<VectorHit>) -> Self {
78        let candidates_scanned = hits.len();
79        Self::ready_with_scan_count(hits, candidates_scanned)
80    }
81
82    pub fn ready_with_scan_count(hits: Vec<VectorHit>, candidates_scanned: usize) -> Self {
83        Self::ready_with_scan_count_and_timings(hits, candidates_scanned, vec![])
84    }
85
86    fn ready_with_scan_count_and_timings(
87        hits: Vec<VectorHit>,
88        candidates_scanned: usize,
89        timings: Vec<crate::perf::PhaseTiming>,
90    ) -> Self {
91        Self {
92            hits,
93            disabled_reason: None,
94            candidates_scanned,
95            timings,
96        }
97    }
98}
99
100#[derive(Debug, Clone, Copy, Default)]
101pub struct VectorSearchFilters<'a> {
102    pub project: Option<&'a str>,
103    pub memory_type: Option<&'a str>,
104    pub branch: Option<&'a str>,
105    pub include_stale: bool,
106}
107
108/// Register the statically linked sqlite-vec extension on this connection
109/// (GH-957). The crate compiles `sqlite-vec.c` with `SQLITE_CORE`, so direct
110/// initialization with a null API pointer is the supported pattern and no
111/// dynamic loading is involved. An init failure logs at error level and
112/// leaves the connection on the portable brute-force cosine scan; the search
113/// contract is unchanged either way.
114pub fn load_vec_extension(conn: &Connection) -> Result<()> {
115    if vec_extension_loaded(conn) {
116        return Ok(());
117    }
118    // The crate declares the entry point with a zero-arg signature (intended
119    // for transmute into `sqlite3_auto_extension`, which would only serve
120    // connections opened after registration). Rebind the same symbol with the
121    // true SQLite extension entry-point signature so every connection is
122    // initialized deterministically at open time.
123    type SqliteVecInit = unsafe extern "C" fn(
124        db: *mut rusqlite::ffi::sqlite3,
125        pz_err_msg: *mut *mut std::ffi::c_char,
126        p_api: *const std::ffi::c_void,
127    ) -> std::ffi::c_int;
128    // SAFETY: `sqlite-vec` is compiled with SQLITE_CORE (see its build.rs), so
129    // the p_api argument is unused and null is valid; the transmute only
130    // restores the entry point's real signature (the crate's own test uses the
131    // same address for auto-extension registration). The handle comes from a
132    // live rusqlite connection on this thread.
133    let init: SqliteVecInit =
134        unsafe { std::mem::transmute(sqlite_vec::sqlite3_vec_init as *const ()) };
135    let rc = unsafe { init(conn.handle(), std::ptr::null_mut(), std::ptr::null()) };
136    if rc != rusqlite::ffi::SQLITE_OK {
137        crate::log::error(
138            "retrieval",
139            &format!(
140                "sqlite-vec init failed with rc={rc}; vector index disabled, brute-force scan remains"
141            ),
142        );
143        return Ok(());
144    }
145    if !vec_extension_loaded(conn) {
146        crate::log::error(
147            "retrieval",
148            "sqlite-vec init reported success but vec_version() is unavailable; brute-force scan remains",
149        );
150    }
151    Ok(())
152}
153
154/// True when sqlite-vec functions answer on this connection.
155pub(crate) fn vec_extension_loaded(conn: &Connection) -> bool {
156    conn.query_row("SELECT vec_version()", [], |row| row.get::<_, String>(0))
157        .is_ok()
158}
159
160pub fn ensure_vec_table(conn: &Connection) -> Result<()> {
161    create_embedding_table(conn)
162}
163
164pub fn upsert_embedding(conn: &Connection, memory_id: i64, embedding: &[f32]) -> Result<()> {
165    if super::embedding::provider_disabled_or_error()? {
166        return Ok(());
167    }
168    upsert_embedding_with_metadata(
169        conn,
170        memory_id,
171        DEFAULT_EMBEDDING_MODEL,
172        "",
173        embedding,
174        chrono::Utc::now().timestamp(),
175    )
176}
177
178/// Upsert the index embedding for one memory row from the authoritative
179/// passage (canonical fields + index-only `search_context`). The stored
180/// `content_hash` is the versioned `memory_index_hash` of the same snapshot.
181pub fn upsert_memory_embedding(
182    conn: &Connection,
183    memory_id: i64,
184    title: &str,
185    content: &str,
186    memory_type: &str,
187    topic_key: Option<&str>,
188    search_context: &str,
189) -> Result<()> {
190    if super::embedding::provider_disabled_or_error()? {
191        return Ok(());
192    }
193    let embedding = match super::embedding::embed_memory_index(
194        title,
195        content,
196        memory_type,
197        topic_key,
198        search_context,
199    ) {
200        Ok(embedding) => embedding,
201        Err(error) if super::embedding::is_embedding_provider_off_error(&error) => return Ok(()),
202        Err(error) if super::embedding::is_local_embedding_model_unavailable_error(&error) => {
203            crate::log::error(
204                "embedding",
205                &format!("memory embedding deferred for memory id={memory_id}: {error}"),
206            );
207            return Ok(());
208        }
209        Err(error) => return Err(error),
210    };
211    let content_hash =
212        super::embedding::memory_index_hash(title, content, memory_type, topic_key, search_context);
213    upsert_embedding_with_metadata(
214        conn,
215        memory_id,
216        embedding.model(),
217        &content_hash,
218        embedding.values(),
219        chrono::Utc::now().timestamp(),
220    )
221    .with_context(|| format!("memory embedding upsert failed for memory id={memory_id}"))
222}
223
224/// Row-atomic upsert used by the enrichment success CAS: the caller already
225/// prepared the vector and index hash for the committed snapshot.
226pub(crate) fn upsert_index_embedding(
227    conn: &Connection,
228    memory_id: i64,
229    model: &str,
230    index_hash: &str,
231    values: &[f32],
232) -> Result<()> {
233    upsert_embedding_with_metadata(
234        conn,
235        memory_id,
236        model,
237        index_hash,
238        values,
239        chrono::Utc::now().timestamp(),
240    )
241    .with_context(|| format!("index embedding upsert failed for memory id={memory_id}"))
242}
243
244pub fn upsert_memory_embedding_for_row(conn: &Connection, memory_id: i64) -> Result<()> {
245    let row: (Option<String>, String, String, String, Option<String>) = conn
246        .query_row(
247            "SELECT topic_key, title, content, memory_type,
248                    CASE WHEN search_context_source_hash IS NOT NULL
249                         THEN search_context ELSE '' END
250             FROM memories
251             WHERE id = ?1",
252            [memory_id],
253            |row| {
254                Ok((
255                    row.get(0)?,
256                    row.get(1)?,
257                    row.get(2)?,
258                    row.get(3)?,
259                    row.get(4)?,
260                ))
261            },
262        )
263        .with_context(|| format!("load memory row for embedding id={memory_id}"))?;
264    let (topic_key, title, content, memory_type, search_context) = row;
265    upsert_memory_embedding(
266        conn,
267        memory_id,
268        &title,
269        &content,
270        &memory_type,
271        topic_key.as_deref(),
272        search_context.as_deref().unwrap_or(""),
273    )
274}
275
276pub fn embedding_count(conn: &Connection) -> Result<i64> {
277    if !table_exists(conn, "memory_embeddings")? {
278        return Ok(0);
279    }
280    Ok(
281        conn.query_row("SELECT COUNT(*) FROM memory_embeddings", [], |row| {
282            row.get(0)
283        })?,
284    )
285}
286
287pub fn embed_query_text(query: &str) -> Vec<f32> {
288    super::embedding::embed_query_text_local(query)
289}
290
291pub fn embed_memory_text(
292    title: &str,
293    content: &str,
294    memory_type: &str,
295    topic_key: Option<&str>,
296) -> Vec<f32> {
297    super::embedding::embed_memory_text_local(title, content, memory_type, topic_key)
298}
299
300pub fn vector_search(
301    conn: &Connection,
302    query_embedding: &[f32],
303    limit: usize,
304) -> Result<Vec<(i64, f32)>> {
305    Ok(
306        vector_search_filtered(conn, query_embedding, VectorSearchFilters::default(), limit)?
307            .hits
308            .into_iter()
309            .map(|hit| (hit.memory_id, hit.distance))
310            .collect(),
311    )
312}
313
314pub fn vector_search_filtered(
315    conn: &Connection,
316    query_embedding: &[f32],
317    filters: VectorSearchFilters<'_>,
318    limit: usize,
319) -> Result<VectorSearchOutcome> {
320    if query_embedding.len() != EMBEDDING_DIMENSIONS {
321        anyhow::bail!(
322            "query embedding must be {} dimensions, got {}",
323            EMBEDDING_DIMENSIONS,
324            query_embedding.len()
325        );
326    }
327    let embedding = TextEmbedding::new(DEFAULT_EMBEDDING_MODEL, query_embedding.to_vec())?;
328    vector_search_embedding_filtered(conn, &embedding, filters, limit)
329}
330
331pub fn vector_search_embedding_filtered(
332    conn: &Connection,
333    query_embedding: &TextEmbedding,
334    filters: VectorSearchFilters<'_>,
335    limit: usize,
336) -> Result<VectorSearchOutcome> {
337    if limit == 0 {
338        return Ok(VectorSearchOutcome::ready(vec![]));
339    }
340    crate::memory::retrieval_enrichment::ensure_retrieval_open(conn)?;
341    if super::embedding::provider_disabled_or_error()? {
342        return Ok(VectorSearchOutcome::disabled("embedding provider is off"));
343    }
344    if !table_exists(conn, "memory_embeddings")? {
345        return Ok(VectorSearchOutcome::disabled(
346            "memory_embeddings table is missing; run migrations/backfill",
347        ));
348    }
349    let mut timings = Vec::new();
350    let profile = query_embedding.profile();
351    let knn_hits = crate::perf::time_result(&mut timings, "vector_knn_index", || {
352        vec_index::knn_candidates(
353            conn,
354            query_embedding.values(),
355            profile,
356            filters,
357            super::vector_candidates::vector_candidate_limit(limit),
358        )
359    })?;
360    if let Some(mut hits) = knn_hits {
361        // An empty KNN answer falls through to the portable path so the
362        // caller still gets its empty-store / missing-profile diagnostics.
363        if !hits.is_empty() {
364            let candidates_scanned = hits.len();
365            hits.truncate(limit);
366            return Ok(VectorSearchOutcome::ready_with_scan_count_and_timings(
367                hits,
368                candidates_scanned,
369                timings,
370            ));
371        }
372    }
373    let candidate_ids = crate::perf::time_result(&mut timings, "vector_select_candidates", || {
374        super::vector_candidates::select_candidate_ids(conn, filters, profile, limit)
375    })?;
376    let candidates_scanned = candidate_ids.len();
377    if candidate_ids.is_empty() {
378        if super::vector_candidates::matching_memory_count(conn, filters)? > 0 {
379            if embedding_count(conn)? == 0 {
380                return Ok(VectorSearchOutcome::disabled_with_timings(
381                    "memory_embeddings table is empty; run `remem reindex-embeddings --limit 1000`",
382                    timings,
383                ));
384            }
385            return Ok(VectorSearchOutcome::disabled_with_timings(
386                format!(
387                    "memory_embeddings has no rows for model={} dimensions={}; run `remem reindex-embeddings --limit 1000`",
388                    profile.model, profile.dimensions
389                ),
390                timings,
391            ));
392        }
393        return Ok(VectorSearchOutcome::ready_with_scan_count_and_timings(
394            vec![],
395            0,
396            timings,
397        ));
398    }
399    let placeholders = std::iter::repeat_n("?", candidate_ids.len())
400        .collect::<Vec<_>>()
401        .join(", ");
402    let sql = format!(
403        "SELECT memory_id, embedding, dimensions
404         FROM memory_embeddings INDEXED BY idx_memory_embeddings_profile_memory_id
405         WHERE model = ?
406           AND dimensions = ?
407           AND memory_id IN ({placeholders})"
408    );
409    let mut param_values: Vec<Box<dyn rusqlite::types::ToSql>> = vec![
410        Box::new(profile.model.to_string()),
411        Box::new(profile.dimensions as i64),
412    ];
413    param_values.extend(
414        candidate_ids
415            .iter()
416            .map(|id| Box::new(*id) as Box<dyn rusqlite::types::ToSql>),
417    );
418    let candidates = crate::perf::time_result(&mut timings, "vector_load_embeddings", || {
419        let refs = crate::db::to_sql_refs(&param_values);
420        let mut stmt = conn.prepare(&sql)?;
421        let rows = stmt.query_map(refs.as_slice(), |row| {
422            Ok((
423                row.get::<_, i64>(0)?,
424                row.get::<_, Vec<u8>>(1)?,
425                row.get::<_, i64>(2)?,
426            ))
427        })?;
428        crate::db::query::collect_rows(rows)
429    })?;
430    let mut hits = crate::perf::time_result(&mut timings, "vector_decode_cosine", || {
431        let mut hits = Vec::new();
432        for (memory_id, blob, dimensions) in candidates {
433            let embedding = decode_embedding(&blob, dimensions)
434                .with_context(|| format!("invalid embedding blob for memory id={memory_id}"))?;
435            let distance = cosine_distance(query_embedding.values(), &embedding)?;
436            hits.push(VectorHit {
437                memory_id,
438                distance,
439            });
440        }
441        Ok(hits)
442    })?;
443    crate::perf::time_value(&mut timings, "vector_sort_truncate", || {
444        hits.sort_by(|a, b| {
445            a.distance
446                .partial_cmp(&b.distance)
447                .unwrap_or(std::cmp::Ordering::Equal)
448                .then_with(|| a.memory_id.cmp(&b.memory_id))
449        });
450        hits.truncate(limit);
451    });
452    Ok(VectorSearchOutcome::ready_with_scan_count_and_timings(
453        hits,
454        candidates_scanned,
455        timings,
456    ))
457}
458
459pub fn find_similar_observations(
460    conn: &Connection,
461    query_embedding: &[f32],
462    threshold: f32,
463    limit: usize,
464) -> Result<Vec<i64>> {
465    let candidates = vector_search(conn, query_embedding, limit)?;
466    let distance_threshold = 1.0 - threshold;
467    let similar: Vec<i64> = candidates
468        .into_iter()
469        .filter(|(_, dist)| *dist < distance_threshold)
470        .map(|(id, _)| id)
471        .collect();
472
473    Ok(similar)
474}
475
476fn create_embedding_table(conn: &Connection) -> Result<()> {
477    conn.execute_batch(
478        "CREATE TABLE IF NOT EXISTS memory_embeddings (
479             memory_id INTEGER NOT NULL,
480             embedding BLOB NOT NULL,
481             dimensions INTEGER NOT NULL,
482             model TEXT NOT NULL,
483             content_hash TEXT NOT NULL,
484             updated_at_epoch INTEGER NOT NULL,
485             PRIMARY KEY(memory_id, model, dimensions),
486             FOREIGN KEY(memory_id) REFERENCES memories(id) ON DELETE CASCADE
487         );
488         CREATE INDEX IF NOT EXISTS idx_memory_embeddings_model
489             ON memory_embeddings(model, updated_at_epoch);
490         CREATE INDEX IF NOT EXISTS idx_memory_embeddings_profile_memory_id
491             ON memory_embeddings(model, dimensions, memory_id);",
492    )?;
493    Ok(())
494}
495
496fn upsert_embedding_with_metadata(
497    conn: &Connection,
498    memory_id: i64,
499    model: &str,
500    content_hash: &str,
501    embedding: &[f32],
502    updated_at_epoch: i64,
503) -> Result<()> {
504    let mut stmt = conn.prepare(UPSERT_EMBEDDING_SQL)?;
505    execute_embedding_upsert(
506        &mut stmt,
507        memory_id,
508        model,
509        content_hash,
510        embedding,
511        updated_at_epoch,
512    )?;
513    vec_index::sync_vec_upsert(conn, memory_id, model, embedding.len())
514}
515
516fn execute_embedding_upsert(
517    stmt: &mut Statement<'_>,
518    memory_id: i64,
519    model: &str,
520    content_hash: &str,
521    embedding: &[f32],
522    updated_at_epoch: i64,
523) -> Result<()> {
524    if model.trim().is_empty() {
525        anyhow::bail!("embedding model must not be empty");
526    }
527    if embedding.is_empty() {
528        anyhow::bail!("embedding vector must not be empty");
529    }
530    if embedding.iter().any(|value| !value.is_finite()) {
531        anyhow::bail!("embedding vector contains non-finite values");
532    }
533    let blob = encode_embedding(embedding);
534    let dimensions = embedding.len() as i64;
535    stmt.execute(params![
536        memory_id,
537        blob,
538        dimensions,
539        model,
540        content_hash,
541        updated_at_epoch
542    ])?;
543    Ok(())
544}
545
546fn encode_embedding(embedding: &[f32]) -> Vec<u8> {
547    let mut out = Vec::with_capacity(std::mem::size_of_val(embedding));
548    for value in embedding {
549        out.extend_from_slice(&value.to_le_bytes());
550    }
551    out
552}
553
554pub(crate) fn decode_embedding(blob: &[u8], dimensions: i64) -> Result<Vec<f32>> {
555    if dimensions <= 0 {
556        anyhow::bail!("embedding dimensions must be positive, got {dimensions}");
557    }
558    let dimensions = dimensions as usize;
559    let expected_bytes = dimensions * std::mem::size_of::<f32>();
560    if blob.len() != expected_bytes {
561        anyhow::bail!(
562            "embedding blob must be {} bytes, got {}",
563            expected_bytes,
564            blob.len()
565        );
566    }
567    Ok(blob
568        .as_chunks::<{ std::mem::size_of::<f32>() }>()
569        .0
570        .iter()
571        .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
572        .collect())
573}
574
575pub(crate) fn cosine_distance(a: &[f32], b: &[f32]) -> Result<f32> {
576    if a.len() != b.len() {
577        anyhow::bail!(
578            "embedding dimensions differ: query={} stored={}",
579            a.len(),
580            b.len()
581        );
582    }
583    let mut dot = 0.0f32;
584    let mut a_norm = 0.0f32;
585    let mut b_norm = 0.0f32;
586    for (left, right) in a.iter().zip(b) {
587        dot += left * right;
588        a_norm += left * left;
589        b_norm += right * right;
590    }
591    if a_norm == 0.0 || b_norm == 0.0 {
592        return Ok(1.0);
593    }
594    Ok((1.0 - dot / (a_norm.sqrt() * b_norm.sqrt())).clamp(0.0, 2.0))
595}
596
597fn table_exists(conn: &Connection, table: &str) -> Result<bool> {
598    Ok(conn
599        .query_row(
600            "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?1 LIMIT 1",
601            params![table],
602            |_| Ok(()),
603        )
604        .optional()?
605        .is_some())
606}
607
608#[cfg(test)]
609mod tests;
remem/retrieval/vector.rs

remem/retrieval/
vector.rs
