diff --git a/crates/rpg-core/src/config.rs b/crates/rpg-core/src/config.rs index a97d60f..3488323 100644 --- a/crates/rpg-core/src/config.rs +++ b/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 { diff --git a/crates/rpg-nav/src/embeddings.rs b/crates/rpg-nav/src/embeddings.rs index 08ad761..9429a07 100644 --- a/crates/rpg-nav/src/embeddings.rs +++ b/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, BTreeMap); @@ -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, /// Path to the .rpg directory for persistence. @@ -51,6 +59,31 @@ pub struct EmbeddingIndex { fingerprints: BTreeMap, } +/// 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,17 +102,19 @@ impl EmbeddingIndex { /// Fingerprints are loaded from meta for incremental sync support. pub fn load_or_init(project_root: &Path, graph_updated_at: &str) -> Result { 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(), @@ -87,7 +122,7 @@ impl EmbeddingIndex { }); } 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> { 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 { +/// 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) -> 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, + dimension: usize, +) -> Result<()> { let mut buf: Vec = 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) -> Res Ok(()) } -/// Load entity embeddings from binary format. -fn load_binary(path: &Path) -> Result> { +/// 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> { 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> { 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,15 +658,15 @@ 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")); @@ -617,6 +674,41 @@ mod tests { 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()];