feat(nav): selectable embedding model + dynamic dimension

crates/rpg-core/src/config.rs

@@ -57,6 +57,11 @@ pub struct EncodingConfig {
 pub struct NavigationConfig {
     /// Maximum number of search results returned.
     pub search_result_limit: usize,
+    /// Embedding model used for semantic search. Resolved by the navigation layer
+    /// to a concrete model + dimension. Known values: "bge-small-en" (default, 384d,
+    /// general-text), "jina-code" (768d, code-specialized), "bge-base-en" (768d),
+    /// "bge-large-en" (1024d). Unknown values fall back to the default.
+    pub embedding_model: String,
 }
 
 impl Default for EncodingConfig {
@@ -78,6 +83,7 @@ impl Default for NavigationConfig {
     fn default() -> Self {
         Self {
             search_result_limit: 10,
+            embedding_model: "bge-small-en".to_string(),
         }
     }
 }
@@ -127,6 +133,7 @@ impl RpgConfig {
             "RPG_SEARCH_LIMIT",
             &mut config.navigation.search_result_limit,
         );
+        env_override("RPG_EMBED_MODEL", &mut config.navigation.embedding_model);
 
         // Validate drift thresholds
         if config.encoding.drift_ignore_threshold >= config.encoding.drift_auto_threshold {

crates/rpg-nav/src/embeddings.rs

@@ -1,8 +1,9 @@
 //! Embedding-based semantic search using fastembed.
 //!
 //! Stores per-feature embeddings for each entity, enabling max-cosine similarity
-//! search that preserves multi-role entity semantics. Uses BGE-small-en-v1.5
-//! (384 dimensions) via the fastembed crate.
+//! search that preserves multi-role entity semantics. The embedding model is
+//! configurable via `navigation.embedding_model` (default BGE-small-en-v1.5, 384d;
+//! "jina-code" selects a 768d code-specialized model) — see `resolve_embedding_model`.
 
 use anyhow::{Context, Result, ensure};
 use fastembed::{EmbeddingModel, TextEmbedding};
@@ -14,7 +15,10 @@ use std::path::{Path, PathBuf};
 /// Magic bytes for the binary embedding file format.
 const MAGIC: u32 = 0x5250_4745; // "RPGE"
 const FORMAT_VERSION: u32 = 1;
-const DIMENSION: usize = 384;
+/// Dimension of the default model (BGE-small-en-v1.5). Other models report their
+/// own dimension via `resolve_embedding_model`; the active dimension is stored on
+/// the index and in the binary header.
+const DEFAULT_DIMENSION: usize = 384;
 
 /// Loaded embeddings: entity map + fingerprints from disk.
 type LoadedEmbeddings = (HashMap<String, EntityEmbeddings>, BTreeMap<String, String>);
@@ -41,6 +45,10 @@ struct EntityEmbeddings {
 /// In-memory embedding index for semantic search.
 pub struct EmbeddingIndex {
     model: TextEmbedding,
+    /// Canonical name of the active embedding model (e.g. "BAAI/bge-small-en-v1.5").
+    model_name: String,
+    /// Vector dimension of the active model.
+    dimension: usize,
     /// Map from entity_id → feature-level embeddings.
     entities: HashMap<String, EntityEmbeddings>,
     /// Path to the .rpg directory for persistence.
@@ -51,6 +59,31 @@ pub struct EmbeddingIndex {
     fingerprints: BTreeMap<String, String>,
 }
 
+/// Resolve a configured model name to its fastembed model, canonical name, and
+/// vector dimension. Unknown names fall back to the default (BGE-small-en-v1.5).
+fn resolve_embedding_model(name: &str) -> (EmbeddingModel, &'static str, usize) {
+    match name.trim().to_lowercase().as_str() {
+        "jina-code" | "jinaai/jina-embeddings-v2-base-code" => (
+            EmbeddingModel::JinaEmbeddingsV2BaseCode,
+            "jinaai/jina-embeddings-v2-base-code",
+            768,
+        ),
+        "bge-base-en" | "baai/bge-base-en-v1.5" => {
+            (EmbeddingModel::BGEBaseENV15, "BAAI/bge-base-en-v1.5", 768)
+        }
+        "bge-large-en" | "baai/bge-large-en-v1.5" => (
+            EmbeddingModel::BGELargeENV15,
+            "BAAI/bge-large-en-v1.5",
+            1024,
+        ),
+        _ => (
+            EmbeddingModel::BGESmallENV15,
+            "BAAI/bge-small-en-v1.5",
+            DEFAULT_DIMENSION,
+        ),
+    }
+}
+
 /// Statistics from an incremental embedding sync.
 #[derive(Debug, Default)]
 pub struct SyncStats {
@@ -69,25 +102,27 @@ impl EmbeddingIndex {
     /// Fingerprints are loaded from meta for incremental sync support.
     pub fn load_or_init(project_root: &Path, graph_updated_at: &str) -> Result<Self> {
         let rpg_dir = project_root.join(".rpg");
-        let model = init_model(&rpg_dir)?;
+        let (model, model_name, dimension) = init_model(&rpg_dir)?;
 
         let embeddings_path = rpg_dir.join("embeddings.bin");
         let meta_path = rpg_dir.join("embeddings.meta.json");
 
         // Try loading existing index (resilient to corruption)
         if embeddings_path.exists() && meta_path.exists() {
-            match Self::try_load_existing(&meta_path, &embeddings_path) {
+            match Self::try_load_existing(&meta_path, &embeddings_path, &model_name, dimension) {
                 Ok(Some((entities, fingerprints))) => {
                     return Ok(Self {
                         model,
+                        model_name,
+                        dimension,
                         entities,
                         rpg_dir,
                         graph_updated_at: graph_updated_at.to_string(),
                         fingerprints,
                     });
                 }
                 Ok(None) => {
-                    // Model/dimension mismatch — start fresh
+                    // Model/dimension mismatch (e.g. embedding_model changed) — rebuild fresh
                 }
                 Err(e) => {
                     // Corrupt on-disk data — delete and start fresh
@@ -101,6 +136,8 @@ impl EmbeddingIndex {
         // No valid index — start fresh
         Ok(Self {
             model,
+            model_name,
+            dimension,
             entities: HashMap::new(),
             rpg_dir,
             graph_updated_at: graph_updated_at.to_string(),
@@ -109,22 +146,25 @@ impl EmbeddingIndex {
     }
 
     /// Try to load existing embedding data. Returns Ok(Some((entities, fingerprints)))
-    /// if valid, Ok(None) if model/dimension mismatch, Err if corrupt.
+    /// if valid, Ok(None) if the stored model/dimension differs from the active one
+    /// (forcing a rebuild), Err if corrupt.
     fn try_load_existing(
         meta_path: &Path,
         embeddings_path: &Path,
+        expected_model: &str,
+        expected_dim: usize,
     ) -> Result<Option<LoadedEmbeddings>> {
         let meta_json =
             std::fs::read_to_string(meta_path).context("failed to read embeddings meta")?;
         let meta: EmbeddingMeta =
             serde_json::from_str(&meta_json).context("failed to parse embeddings meta")?;
 
-        // Only reject on model/dimension mismatch — fingerprints handle staleness
-        if meta.model != "BAAI/bge-small-en-v1.5" || meta.dimension != DIMENSION as u32 {
+        // Reject on model/dimension mismatch — fingerprints handle within-model staleness
+        if meta.model != expected_model || meta.dimension != expected_dim as u32 {
             return Ok(None);
         }
 
-        let entities = load_binary(embeddings_path)?;
+        let entities = load_binary(embeddings_path, expected_dim)?;
         Ok(Some((entities, meta.entity_fingerprints)))
     }
 
@@ -294,11 +334,15 @@ impl EmbeddingIndex {
     /// Save the index to disk (binary + meta sidecar with fingerprints).
     pub fn save(&self) -> Result<()> {
         std::fs::create_dir_all(&self.rpg_dir)?;
-        save_binary(&self.rpg_dir.join("embeddings.bin"), &self.entities)?;
+        save_binary(
+            &self.rpg_dir.join("embeddings.bin"),
+            &self.entities,
+            self.dimension,
+        )?;
 
         let meta = EmbeddingMeta {
-            model: "BAAI/bge-small-en-v1.5".to_string(),
-            dimension: DIMENSION as u32,
+            model: self.model_name.clone(),
+            dimension: self.dimension as u32,
             version: FORMAT_VERSION,
             graph_updated_at: self.graph_updated_at.clone(),
             entity_fingerprints: self.fingerprints.clone(),
@@ -328,19 +372,26 @@ fn compute_fingerprint(features: &[String]) -> String {
     format!("{:016x}", hasher.finish())
 }
 
-/// Initialize the fastembed model with cache in .rpg/models/.
-fn init_model(rpg_dir: &Path) -> Result<TextEmbedding> {
+/// Initialize the fastembed model with cache in .rpg/models/. The model is chosen
+/// from `navigation.embedding_model` (config.toml / RPG_EMBED_MODEL / default).
+/// Returns the model plus its canonical name and dimension.
+fn init_model(rpg_dir: &Path) -> Result<(TextEmbedding, String, usize)> {
     let cache_dir = rpg_dir.join("models");
     std::fs::create_dir_all(&cache_dir)?;
 
-    let options = fastembed::TextInitOptions::new(EmbeddingModel::BGESmallENV15)
+    let project_root = rpg_dir.parent().unwrap_or(rpg_dir);
+    let cfg = rpg_core::config::RpgConfig::load(project_root).unwrap_or_default();
+    let (model_kind, canonical, dimension) =
+        resolve_embedding_model(&cfg.navigation.embedding_model);
+
+    let options = fastembed::TextInitOptions::new(model_kind)
         .with_show_download_progress(true)
         .with_cache_dir(cache_dir);
 
     let model = TextEmbedding::try_new(options)
-        .context("failed to initialize embedding model (BGE-small-en-v1.5)")?;
+        .with_context(|| format!("failed to initialize embedding model ({canonical})"))?;
 
-    Ok(model)
+    Ok((model, canonical.to_string(), dimension))
 }
 
 /// Cosine similarity between two vectors.
@@ -365,14 +416,19 @@ fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 {
     dot / denom
 }
 
-/// Save entity embeddings to binary format.
-fn save_binary(path: &Path, entities: &HashMap<String, EntityEmbeddings>) -> Result<()> {
+/// Save entity embeddings to binary format. `dimension` is written to the header
+/// (it must match the active model's dimension).
+fn save_binary(
+    path: &Path,
+    entities: &HashMap<String, EntityEmbeddings>,
+    dimension: usize,
+) -> Result<()> {
     let mut buf: Vec<u8> = Vec::new();
 
     // Header (16 bytes)
     buf.write_all(&MAGIC.to_le_bytes())?;
     buf.write_all(&FORMAT_VERSION.to_le_bytes())?;
-    buf.write_all(&(DIMENSION as u32).to_le_bytes())?;
+    buf.write_all(&(dimension as u32).to_le_bytes())?;
     buf.write_all(&(entities.len() as u32).to_le_bytes())?;
 
     // Per entity
@@ -403,8 +459,9 @@ fn save_binary(path: &Path, entities: &HashMap<String, EntityEmbeddings>) -> Res
     Ok(())
 }
 
-/// Load entity embeddings from binary format.
-fn load_binary(path: &Path) -> Result<HashMap<String, EntityEmbeddings>> {
+/// Load entity embeddings from binary format. `expected_dim` is the active model's
+/// dimension; a header mismatch means the file was written by a different model.
+fn load_binary(path: &Path, expected_dim: usize) -> Result<HashMap<String, EntityEmbeddings>> {
     let data = std::fs::read(path).context("failed to read embeddings.bin")?;
     let mut cursor = &data[..];
 
@@ -414,7 +471,7 @@ fn load_binary(path: &Path) -> Result<HashMap<String, EntityEmbeddings>> {
     let version = read_u32(&mut cursor)?;
     anyhow::ensure!(version == FORMAT_VERSION, "unsupported embeddings version");
     let dimension = read_u32(&mut cursor)? as usize;
-    anyhow::ensure!(dimension == DIMENSION, "dimension mismatch");
+    anyhow::ensure!(dimension == expected_dim, "dimension mismatch");
     let entity_count = read_u32(&mut cursor)? as usize;
 
     let mut entities = HashMap::with_capacity(entity_count);
@@ -601,22 +658,57 @@ mod tests {
         entities.insert(
             "test:func".to_string(),
             EntityEmbeddings {
-                vectors: vec![vec![0.1; DIMENSION], vec![0.2; DIMENSION]],
+                vectors: vec![vec![0.1; DEFAULT_DIMENSION], vec![0.2; DEFAULT_DIMENSION]],
             },
         );
 
         let dir = tempfile::tempdir().unwrap();
         let path = dir.path().join("test.bin");
 
-        save_binary(&path, &entities).unwrap();
-        let loaded = load_binary(&path).unwrap();
+        save_binary(&path, &entities, DEFAULT_DIMENSION).unwrap();
+        let loaded = load_binary(&path, DEFAULT_DIMENSION).unwrap();
 
         assert_eq!(loaded.len(), 1);
         assert!(loaded.contains_key("test:func"));
         assert_eq!(loaded["test:func"].vectors.len(), 2);
         assert!((loaded["test:func"].vectors[0][0] - 0.1).abs() < 1e-6);
     }
 
+    #[test]
+    fn test_binary_dimension_mismatch_rejected() {
+        // A file written at 768d must not be read as the default 384d.
+        let mut entities = HashMap::new();
+        entities.insert(
+            "x".to_string(),
+            EntityEmbeddings {
+                vectors: vec![vec![0.5; 768]],
+            },
+        );
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("e.bin");
+        save_binary(&path, &entities, 768).unwrap();
+        assert!(load_binary(&path, DEFAULT_DIMENSION).is_err());
+        assert!(load_binary(&path, 768).is_ok());
+    }
+
+    #[test]
+    fn test_resolve_embedding_model() {
+        assert_eq!(resolve_embedding_model("bge-small-en").2, 384);
+        assert_eq!(resolve_embedding_model("jina-code").2, 768);
+        assert_eq!(
+            resolve_embedding_model("jina-code").1,
+            "jinaai/jina-embeddings-v2-base-code"
+        );
+        assert_eq!(resolve_embedding_model("bge-large-en").2, 1024);
+        // unknown -> default
+        assert_eq!(
+            resolve_embedding_model("nonexistent").1,
+            "BAAI/bge-small-en-v1.5"
+        );
+        // case-insensitive
+        assert_eq!(resolve_embedding_model("JINA-CODE").2, 768);
+    }
+
     #[test]
     fn test_fingerprint_deterministic() {
         let features = vec!["validate input".to_string(), "return result".to_string()];