feat: Dynamically detect RAG embedding model dimension for any given model

src/rag/mod.rs

@@ -151,7 +151,7 @@ impl Rag {
         println!("⚙ Initializing RAG...");
         let mut data = Self::resolve_init_data(app, config)?;
         data.driver = config.driver.clone().unwrap_or_else(|| "yaml".to_string());
-        let mut rag = Self::create(app, name, save_path, data)?;
+        let mut rag = Self::create(app, name, save_path, data).await?;
         let loaders = app.document_loaders.clone();
         let (spinner, spinner_rx) = Spinner::create("");
         abortable_run_with_spinner_rx(
@@ -298,7 +298,7 @@ impl Rag {
             },
         );
         data.driver = driver;
-        let mut rag = Self::create(app, name, save_path, data)?;
+        let mut rag = Self::create(app, name, save_path, data).await?;
         let mut paths = doc_paths.to_vec();
         if paths.is_empty() {
             paths = add_documents()?;
@@ -317,12 +317,12 @@ impl Rag {
         Ok(rag)
     }
 
-    pub fn load(app: &AppConfig, name: &str, path: &Path) -> Result<Self> {
+    pub async fn load(app: &AppConfig, name: &str, path: &Path) -> Result<Self> {
         let err = || format!("Failed to load rag '{name}' at '{}'", path.display());
         let content = fs::read_to_string(path).with_context(err)?;
         let data: RagData = serde_yaml::from_str(&content).with_context(err)?;
         data.validate().with_context(err)?;
-        Self::create(app, name, path, data)
+        Self::create(app, name, path, data).await
     }
 
     /// Loads a RAG from a YAML file. External drivers need an async constructor
@@ -372,7 +372,7 @@ impl Rag {
                     last_sources: RwLock::new(None),
                 })
             }
-            _ => Self::load(app, name, path),
+            _ => Self::load(app, name, path).await,
         }
     }
 
@@ -537,14 +537,22 @@ impl Rag {
         Ok(rag)
     }
 
-    pub fn create(app: &AppConfig, name: &str, path: &Path, mut data: RagData) -> Result<Self> {
+    pub async fn create(
+        app: &AppConfig,
+        name: &str,
+        path: &Path,
+        mut data: RagData,
+    ) -> Result<Self> {
         // Deliberately does NOT call rebuild_indexes: both callers construct the Rag
         // before any documents are added, so rebuilding empty data would be a no-op.
         // Actual population happens later via sync_documents.
         let (provider, bm25): (Box<dyn RagProvider>, _) = match data.driver.as_str() {
             "duckdb" => {
                 let db_path = providers::duckdb_path_from_yaml(path);
-                let dim = embedding_dim_for_model(&data.embedding_model);
+                let dim = match DuckDbProvider::introspect_dim(&db_path)? {
+                    Some(existing) => existing,
+                    None => probe_embedding_dim(app, &data.embedding_model).await?,
+                };
                 let duck = DuckDbProvider::open(&db_path, dim)?;
                 // HYDRATE — mandatory, not an optimization. The YAML file for a duckdb
                 // RAG deliberately omits `vectors`, so `data.vectors` arrives empty from
@@ -2119,21 +2127,25 @@ fn reciprocal_rank_fusion(
         .collect()
 }
 
-/// Map an embedding model id to its vector dimension.
-///
-/// The DuckDB `FLOAT[N]` column type and its HNSW index are fixed at schema-creation
-/// time, so this value must be decided before the first insert. An unrecognized model
-/// falls back to 1536; if that is wrong, DuckDB raises a dimension-mismatch error on
-/// the first insert rather than silently corrupting the schema, and the recovery is to
-/// delete the sidecar and re-ingest from source.
-fn embedding_dim_for_model(model_id: &str) -> usize {
-    match model_id {
-        m if m.contains("3-large") => 3072,
-        m if m.contains("3-small") || m.contains("ada-002") => 1536,
-        m if m.contains("nomic-embed-text") || m.contains("all-minilm") => 768,
-        m if m.contains("jina-embeddings-v2") => 1024,
-        _ => 1536,
+async fn probe_embedding_dim(app: &AppConfig, model_id: &str) -> Result<usize> {
+    let model = Model::retrieve_model(app, model_id, ModelType::Embedding)?;
+    let client = init_client(&Arc::new(app.clone()), model)?;
+    let out = client
+        .embeddings(&EmbeddingsData::new(vec!["dimension probe".into()], false))
+        .await
+        .with_context(|| {
+            format!(
+                "Failed to probe the embedding dimension of model '{model_id}'. \
+                 Creating a duckdb RAG requires one call to the embedding endpoint."
+            )
+        })?;
+    let dim = out.first().map(|v| v.len()).unwrap_or(0);
+
+    if dim == 0 {
+        bail!("Embedding model '{model_id}' returned an empty vector during the dimension probe");
     }
+
+    Ok(dim)
 }
 
 /// True only for "the vault does not hold this key".
@@ -2673,18 +2685,6 @@ mod tests {
         assert_eq!(data.attached_source_label(), "[external collection]");
     }
 
-    #[test]
-    fn embedding_dim_for_model_maps_known_models() {
-        assert_eq!(embedding_dim_for_model("text-embedding-3-large"), 3072);
-        assert_eq!(embedding_dim_for_model("text-embedding-3-small"), 1536);
-        assert_eq!(embedding_dim_for_model("text-embedding-ada-002"), 1536);
-        assert_eq!(embedding_dim_for_model("nomic-embed-text"), 768);
-        assert_eq!(embedding_dim_for_model("all-minilm"), 768);
-        assert_eq!(embedding_dim_for_model("jina-embeddings-v2-base-en"), 1024);
-        // Unknown models fall back to the OpenAI-compatible default.
-        assert_eq!(embedding_dim_for_model("some-unknown-model"), 1536);
-    }
-
     #[test]
     fn document_id_round_trip() {
         let id = DocumentId::new(5, 17);

src/rag/providers/duckdb.rs

@@ -5,7 +5,7 @@ use std::collections::HashMap;
 use anyhow::{Context, Result, anyhow, bail};
 use async_trait::async_trait;
 use duckdb::types::Value;
-use duckdb::{AccessMode, Config, Connection};
+use duckdb::{AccessMode, Config, Connection, OptionalExt};
 use indexmap::IndexMap;
 use log::warn;
 use std::path::{Path, PathBuf};
@@ -30,6 +30,11 @@ pub(crate) fn duckdb_path_from_yaml(yaml_path: &Path) -> PathBuf {
     yaml_path.with_extension("duckdb")
 }
 
+fn parse_float_array_dim(data_type: &str) -> Option<usize> {
+    let inner = data_type.strip_prefix("FLOAT[")?.strip_suffix(']')?;
+    inner.parse::<usize>().ok().filter(|&n| n > 0)
+}
+
 /// The shared connection together with the access mode it was opened with.
 ///
 /// `conn` is an `Option` only so that an upgrade can DROP the read-only connection
@@ -42,6 +47,11 @@ struct ConnHandle {
     /// `duplicate()` clone sharing the `Arc` observes an upgrade performed through any
     /// other handle instead of keeping its own stale copy of the mode.
     writable: bool,
+    /// Embedding dimension the `FLOAT[N]` column was opened (or last rebuilt) with.
+    /// Shared for the same reason as `writable`: a self-healing rebuild through one
+    /// handle updates the width, and every `duplicate()` clone must cast with the new
+    /// width instead of erroring on a healthy store with its stale copy.
+    dim: usize,
 }
 
 impl ConnHandle {
@@ -68,8 +78,6 @@ impl ConnHandle {
 pub struct DuckDbProvider {
     path: PathBuf,
     conn: Arc<Mutex<ConnHandle>>,
-    /// Embedding dimension; fixed at open time because the `FLOAT[N]` column depends on it.
-    dim: usize,
     /// True once an FTS index has been built on `documents`. Until then
     /// `fts_main_documents.match_bm25` does not exist and any keyword query would
     /// fail with a DuckDB catalog error. Backs `has_native_keyword_search`.
@@ -96,8 +104,8 @@ impl DuckDbProvider {
             conn: Arc::new(Mutex::new(ConnHandle {
                 conn: Some(conn),
                 writable,
+                dim,
             })),
-            dim,
             fts_ready: AtomicBool::new(fts_exists),
         })
     }
@@ -164,6 +172,35 @@ impl DuckDbProvider {
         Ok(conn)
     }
 
+    pub fn introspect_dim(db_path: &Path) -> Result<Option<usize>> {
+        if !db_path.exists() {
+            return Ok(None);
+        }
+        let conn = Self::open_read_only(db_path).with_context(|| {
+            format!(
+                "Cannot inspect the existing RAG store at '{}'",
+                db_path.display()
+            )
+        })?;
+        let ty: Option<String> = conn
+            .query_row(
+                "SELECT data_type FROM duckdb_columns() \
+                 WHERE table_name = 'vectors' AND column_name = 'embedding'",
+                [],
+                |r| r.get(0),
+            )
+            .optional()
+            .with_context(|| {
+                format!(
+                    "Failed to introspect the embedding dimension of the DuckDB store \
+                     at '{}'",
+                    db_path.display()
+                )
+            })?;
+
+        Ok(ty.and_then(|t| parse_float_array_dim(&t)))
+    }
+
     /// Open the store read-write and make sure its schema exists. Exactly one process
     /// may hold such a handle, and no reader from another process may hold it meanwhile.
     fn open_read_write(db_path: &Path, dim: usize) -> Result<Connection> {
@@ -245,7 +282,7 @@ impl DuckDbProvider {
             return Ok(());
         }
         drop(handle.conn.take());
-        match Self::open_read_write(&self.path, self.dim) {
+        match Self::open_read_write(&self.path, handle.dim) {
             Ok(conn) => {
                 handle.conn = Some(conn);
                 handle.writable = true;
@@ -428,13 +465,33 @@ impl RagProvider for DuckDbProvider {
         if embedding.iter().any(|f| !f.is_finite()) {
             bail!("Query embedding contains a non-finite value (NaN or infinity)");
         }
+        let handle = self.lock_conn()?;
+        let dim = handle.dim;
+        if embedding.len() != dim {
+            let rows: i64 = handle
+                .conn()?
+                .query_row("SELECT count(*) FROM vectors", [], |r| r.get(0))
+                .context("Failed to count vectors before a dimension-mismatch query")?;
+            if rows == 0 {
+                // A never-synced store answers "nothing", not a cast error.
+                return Ok(Vec::new());
+            }
+
+            bail!(
+                "RAG store at '{}' was built with {dim}-dim embeddings, but the \
+                 embedding model now returns {}-dim vectors. The embedding model \
+                 changed since ingestion. Re-embed the documents, or delete the \
+                 sidecar file and re-ingest.",
+                self.path.display(),
+                embedding.len()
+            );
+        }
+
         let vals: String = embedding
             .iter()
             .map(|f| f.to_string())
             .collect::<Vec<_>>()
             .join(", ");
-        let dim = self.dim;
-        let handle = self.lock_conn()?;
         // array_cosine_distance requires a FLOAT[N] ARRAY, not the LIST type FLOAT[].
         // ORDER BY distance ASC is required for the planner to use hnsw_idx; the
         // similarity form (DESC) does NOT trigger the ANN index. Distance is converted
@@ -554,7 +611,28 @@ impl RagProvider for DuckDbProvider {
                 );
             }
         }
-        let dim = self.dim;
+        let dim = match data.vectors.first() {
+            None => self.lock_conn()?.dim,
+            Some((_, first)) => {
+                let dim = first.len();
+                if let Some((doc_id, other)) = data.vectors.iter().find(|(_, e)| e.len() != dim) {
+                    let matching = data.vectors.values().filter(|e| e.len() == dim).count();
+
+                    bail!(
+                        "Refusing to rebuild the RAG store at '{}': the rebuild batch \
+                         mixes {dim}-dim and {}-dim vectors ({matching} vs {} vectors; \
+                         first mismatch: document {}). Re-embed the documents, or \
+                         delete the sidecar file and re-ingest.",
+                        self.path.display(),
+                        other.len(),
+                        data.vectors.len() - matching,
+                        doc_id.0
+                    );
+                }
+
+                dim
+            }
+        };
         // THE write path. Everything above this line only reads, so the upgrade happens
         // here, after both guards have had their say: a rebuild that is going to be
         // refused must not first take the exclusive lock away from other processes.
@@ -664,6 +742,8 @@ impl RagProvider for DuckDbProvider {
         // fall back to local BM25, which is also empty, and therefore correct.
         self.fts_ready.store(doc_count > 0, Ordering::Relaxed);
 
+        handle.dim = dim;
+
         Ok(())
     }
 
@@ -729,10 +809,11 @@ impl RagProvider for DuckDbProvider {
         // Sharing the Arc also shares the ACCESS MODE, which lives inside the ConnHandle
         // rather than beside it: when one handle upgrades itself to read-write, every
         // clone is upgraded with it and none is left holding a stale "read-only" belief.
+        // The embedding dimension lives there too, so a self-healing rebuild through
+        // one handle updates the width every clone casts with.
         Box::new(DuckDbProvider {
             path: self.path.clone(),
             conn: Arc::clone(&self.conn),
-            dim: self.dim,
             fts_ready: AtomicBool::new(self.fts_ready.load(Ordering::Relaxed)),
         })
     }
@@ -1074,6 +1155,127 @@ mod tests {
             .expect("a fresh RAG with nothing indexed must rebuild cleanly");
     }
 
+    #[test]
+    fn parse_float_array_dim_handles_arrays_lists_and_scalars() {
+        assert_eq!(parse_float_array_dim("FLOAT[768]"), Some(768));
+        assert_eq!(parse_float_array_dim("FLOAT[]"), None);
+        assert_eq!(parse_float_array_dim("VARCHAR"), None);
+    }
+
+    #[test]
+    fn introspect_dim_round_trips_the_open_dim() {
+        let db = TempDb::new("introspect");
+
+        assert_eq!(
+            DuckDbProvider::introspect_dim(&db.path).unwrap(),
+            None,
+            "a file that does not exist has no dim"
+        );
+        {
+            let _provider = DuckDbProvider::open(&db.path, 5).unwrap();
+        }
+        assert_eq!(DuckDbProvider::introspect_dim(&db.path).unwrap(), Some(5));
+    }
+
+    #[test]
+    fn introspect_dim_propagates_an_unopenable_existing_file() {
+        let db = TempDb::new("introspectgarbage");
+        fs::write(&db.path, b"not a duckdb database").unwrap();
+
+        let err = DuckDbProvider::introspect_dim(&db.path).unwrap_err();
+
+        assert!(
+            format!("{err:#}").contains(&format!(
+                "Cannot inspect the existing RAG store at '{}'",
+                db.path.display()
+            )),
+            "an existing-but-unopenable file must be an error naming the path; got: {err:#}"
+        );
+    }
+
+    #[tokio::test]
+    async fn rebuild_indexes_self_heals_dim_from_the_vectors_it_writes() {
+        let db = TempDb::new("selfheal");
+        let mut provider = DuckDbProvider::open(&db.path, 5).unwrap();
+        let mut data = minimal_rag_data();
+        data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
+
+        provider.rebuild_indexes(&data, true).await.unwrap();
+
+        let results = provider
+            .vector_search(&[0.1, 0.2, 0.3], 5, 0.0)
+            .await
+            .unwrap();
+        assert_eq!(results.len(), 1);
+
+        drop(provider);
+        assert_eq!(DuckDbProvider::introspect_dim(&db.path).unwrap(), Some(3));
+    }
+
+    #[tokio::test]
+    async fn a_self_healed_dim_is_visible_through_duplicate_clones() {
+        let db = TempDb::new("dimdup");
+        let mut provider = DuckDbProvider::open(&db.path, 5).unwrap();
+        let dup = provider.duplicate(&minimal_rag_data());
+
+        let mut data = minimal_rag_data();
+        data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
+        provider.rebuild_indexes(&data, true).await.unwrap();
+
+        let results = dup.vector_search(&[0.1, 0.2, 0.3], 5, 0.0).await.unwrap();
+        assert_eq!(
+            results.len(),
+            1,
+            "a duplicate() clone must observe the dim written by a rebuild through \
+             the original"
+        );
+    }
+
+    #[tokio::test]
+    async fn rebuild_indexes_rejects_mixed_dim_vectors() {
+        let db = TempDb::new("mixeddim");
+        let mut provider = DuckDbProvider::open(&db.path, 3).unwrap();
+        let mut data = minimal_rag_data();
+        data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
+        data.vectors.insert(DocumentId(1), vec![0.1, 0.2, 0.3, 0.4]);
+
+        let err = provider.rebuild_indexes(&data, true).await.unwrap_err();
+
+        assert!(
+            err.to_string().contains("rebuild batch mixes"),
+            "got: {err}"
+        );
+    }
+
+    #[tokio::test]
+    async fn vector_search_dim_mismatch_on_empty_store_returns_nothing() {
+        let db = TempDb::new("dimempty");
+        let provider = DuckDbProvider::open(&db.path, 3).unwrap();
+
+        let results = provider.vector_search(&[0.1, 0.2], 5, 0.0).await.unwrap();
+
+        assert!(
+            results.is_empty(),
+            "a never-synced store must answer 'nothing', not a cast error"
+        );
+    }
+
+    #[tokio::test]
+    async fn vector_search_dim_mismatch_on_populated_store_errors() {
+        let db = TempDb::new("dimfull");
+        let mut provider = DuckDbProvider::open(&db.path, 3).unwrap();
+        let mut data = minimal_rag_data();
+        data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
+        provider.rebuild_indexes(&data, true).await.unwrap();
+
+        let err = provider
+            .vector_search(&[0.1, 0.2], 5, 0.0)
+            .await
+            .unwrap_err();
+
+        assert!(err.to_string().contains("was built with"), "got: {err}");
+    }
+
     #[tokio::test]
     async fn duplicate_shares_the_same_connection() {
         let db = TempDb::new("dup");