diff --git a/docs/CLI-GUIDE.md b/docs/CLI-GUIDE.md index 7c1aa0b..72293fd 100644 --- a/docs/CLI-GUIDE.md +++ b/docs/CLI-GUIDE.md @@ -133,6 +133,40 @@ Everything under `/.leio-code/`: Secrets in `explain` are redacted unless `--show-secrets` on a TTY. +### Encoder model and width + +The encoder is described by two values that travel together, in `[embed]` of +`.leio-code/config.toml` (or `LEIO_CODE_EMBED_MODEL` / `LEIO_CODE_EMBED_DIM`, +which win): + +```toml +[embed] +url = "http://localhost:8080/embed" +model = "google/embeddinggemma-2" +dim = 768 +``` + +Defaults are `BAAI/bge-m3` / `1024`, the contract every store written before +`dim` existed used, so leaving `[embed]` alone changes nothing. Any width works: +`embeddinggemma-2` is 768 (Matryoshka cuts: 512/256/128), bge-m3 is 1024. + +Three rules hold the contract together: + +- **Ingest rejects a mismatch.** A response whose width is not `dim` is an + error on all three endpoint forms, never a stored vector that cannot be + compared. Gemini requests send `dim` as `outputDimensionality`. +- **Rows carry real provenance.** `embed_model` / `embed_dim` on each row are + the configured pair that produced it, not a compiled-in default. Rows that + fell through the graceful-degrade path stay `"none"` / `0`. +- **Search skips incompatible vectors.** Stored vectors of a different width + than the query vector are not scored; those candidates fall back to on-demand + embedding or lexical ranking. + +Changing `model` or `dim` makes existing stored vectors incomparable. Re-run +`export arrow-nodes` to rebuild them; a store is not migrated in place. Note +that vectors from two *different* models of the same width are not detected +by the width check alone — do not point one store at two encoders. + ## Knowledge wiki Markdown in the repo compiles to a local Arrow wiki — same idea as diff --git a/src/config.rs b/src/config.rs index fb6f419..37edad2 100644 --- a/src/config.rs +++ b/src/config.rs @@ -11,9 +11,11 @@ //! [`PLATFORM_CONTRACTS_ROOT_ENV`], [`PLATFORM_CARTRIDGE_ENV`]) add one //! private checkout beside the neutral defaults. Unset means only those //! defaults are recognized. -//! - `embed.url` / `embed.model` → `LEIO_CODE_EMBED_URL` / -//! `LEIO_CODE_EMBED_MODEL` (optional BGE-M3 encoder). Direct TEI URLs end in -//! `/embed`; LiteLLM/OpenAI-compatible bases receive `/v1/embeddings`. +//! - `embed.url` / `embed.model` / `embed.dim` → `LEIO_CODE_EMBED_URL` / +//! `LEIO_CODE_EMBED_MODEL` / `LEIO_CODE_EMBED_DIM` (optional dense encoder). +//! Direct TEI URLs end in `/embed`; LiteLLM/OpenAI-compatible bases receive +//! `/v1/embeddings`. Model and dimension travel together: a store embedded at +//! one pair must not be scored against another. //! //! [`repo_namespace`] resolves a stable repo-name slug for cross-graph URN //! generation; falls back to the working-dir basename. @@ -35,12 +37,22 @@ pub const PROFILE_EXAMPLE: &str = "example"; pub const DEFAULT_CODE_RDF_NAMESPACE: &str = "https://example.local/leio/code#"; /// Env var that overrides [`DEFAULT_CODE_RDF_NAMESPACE`] / `[rdf] namespace`. pub const CODE_RDF_NAMESPACE_ENV: &str = "LEIO_CODE_RDF_NAMESPACE"; -/// OpenAI-compatible or direct TEI URL for BGE-M3 (`[embed] url`). +/// OpenAI-compatible or direct TEI URL for the dense encoder (`[embed] url`). pub const EMBED_URL_ENV: &str = "LEIO_CODE_EMBED_URL"; -/// Encoder model id. Defaults to `BAAI/bge-m3`. +/// Encoder model id. Defaults to [`DEFAULT_EMBED_MODEL`]. pub const EMBED_MODEL_ENV: &str = "LEIO_CODE_EMBED_MODEL"; -/// Workspace-standard dense model. Must stay 1024-d BGE-M3. +/// Encoder output width. Defaults to [`DEFAULT_EMBED_DIM`]. +pub const EMBED_DIM_ENV: &str = "LEIO_CODE_EMBED_DIM"; +/// Workspace-standard dense model, kept for stores embedded before +/// `[embed] model` became load-bearing. pub const DEFAULT_EMBED_MODEL: &str = "BAAI/bge-m3"; +/// Dense width of [`DEFAULT_EMBED_MODEL`]. +/// +/// BGE-M3 emits 1024 floats. Every store written before `[embed] dim` existed +/// is this wide, so changing the default would silently reinterpret them; +/// a different encoder sets `[embed] dim` instead. Other widths are ordinary: +/// `google/embeddinggemma-2` is 768, and its Matryoshka cuts are 512/256/128. +pub const DEFAULT_EMBED_DIM: usize = 1024; /// Neutral API directory recognized for route mounting and Python modules. /// @@ -219,15 +231,29 @@ pub struct OrphanFilesConfig { pub surfaces: Option>, } -/// Remote BGE-M3 encoder under `[embed]` in `.leio-code/config.toml`. +/// Remote dense encoder under `[embed]` in `.leio-code/config.toml`. /// /// Direct TEI endpoints use the `/embed` suffix; LiteLLM/OpenAI-compatible /// endpoints may use their base URL and receive `/v1/embeddings` automatically. /// +/// `model` and `dim` describe one encoder and must agree with it: the ingest +/// paths reject any response of a different width rather than store a vector +/// that cannot be compared. +/// /// ```toml /// [embed] /// url = "http://tei-bge-m3.example:8080/embed" /// model = "BAAI/bge-m3" +/// dim = 1024 +/// ``` +/// +/// A locally served 768-d encoder looks the same with its own two values: +/// +/// ```toml +/// [embed] +/// url = "http://localhost:8080/embed" +/// model = "google/embeddinggemma-2" +/// dim = 768 /// ``` #[derive(Debug, Clone, Default, Deserialize)] pub struct EmbedConfig { @@ -235,6 +261,8 @@ pub struct EmbedConfig { pub url: Option, /// Model id posted to `/v1/embeddings`. Defaults to [`DEFAULT_EMBED_MODEL`]. pub model: Option, + /// Encoder output width. Defaults to [`DEFAULT_EMBED_DIM`]. + pub dim: Option, } pub fn load_repo_config(root: &Path) -> Option { @@ -299,6 +327,24 @@ pub fn embed_model(root: &Path) -> String { .unwrap_or_else(|| DEFAULT_EMBED_MODEL.to_string()) } +/// Encoder width (`LEIO_CODE_EMBED_DIM` / `[embed] dim` / [`DEFAULT_EMBED_DIM`]). +/// +/// Pairs with [`embed_model`]: both describe the same encoder, and ingest +/// rejects a response that does not match this width. Zero and unparseable +/// values fall back to the default rather than producing empty vectors. +pub fn embed_dim(root: &Path) -> usize { + std::env::var(EMBED_DIM_ENV) + .ok() + .and_then(|value| value.trim().parse::().ok()) + .filter(|dim| *dim > 0) + .or_else(|| { + load_repo_config(root) + .and_then(|config| config.embed.and_then(|embed| embed.dim)) + .filter(|dim| *dim > 0) + }) + .unwrap_or(DEFAULT_EMBED_DIM) +} + /// Adds `http://` when the value has no scheme. Empty → `None`. pub fn normalize_http_base(raw: &str) -> Option { let trimmed = raw.trim().trim_end_matches('/'); @@ -683,6 +729,40 @@ mod tests { } } + #[test] + fn embed_dim_reads_config_and_defaults() { + let dir = tempfile::tempdir().expect("tempdir"); + let cfg_dir = dir.path().join(".leio-code"); + std::fs::create_dir_all(&cfg_dir).expect("mkdir"); + + // No `[embed]` table at all: the BGE-M3 width every older store used. + std::fs::write(cfg_dir.join("config.toml"), "version = 1\n").expect("write"); + if std::env::var(EMBED_DIM_ENV).is_err() { + assert_eq!(embed_dim(dir.path()), DEFAULT_EMBED_DIM); + } + + // A 768-d encoder is configured per repository, not compiled in. + std::fs::write( + cfg_dir.join("config.toml"), + "version = 1\n[embed]\nmodel = \"google/embeddinggemma-2\"\ndim = 768\n", + ) + .expect("write"); + if std::env::var(EMBED_DIM_ENV).is_err() { + assert_eq!(embed_dim(dir.path()), 768); + assert_eq!(embed_model(dir.path()), "google/embeddinggemma-2"); + } + + // Zero would produce empty vectors, so it is treated as unset. + std::fs::write( + cfg_dir.join("config.toml"), + "version = 1\n[embed]\ndim = 0\n", + ) + .expect("write"); + if std::env::var(EMBED_DIM_ENV).is_err() { + assert_eq!(embed_dim(dir.path()), DEFAULT_EMBED_DIM); + } + } + #[test] fn code_namespace_defaults_without_override() { assert_eq!(resolve_code_namespace(None), DEFAULT_CODE_RDF_NAMESPACE); diff --git a/src/embed.rs b/src/embed.rs index e9ca073..3dca400 100644 --- a/src/embed.rs +++ b/src/embed.rs @@ -1,9 +1,14 @@ -//! Optional BGE-M3 embedding of LEIO node rows for the textual regime. +//! Optional dense embedding of LEIO node rows for the textual regime. //! -//! Produces real semantic embeddings via a remote GPU encoder when +//! Produces real semantic embeddings via a remote encoder when //! `LEIO_CODE_EMBED_URL` is set (direct TEI `/embed`, Gemini -//! `generativelanguage.googleapis.com`, or LiteLLM OpenAI `/v1/embeddings`, -//! 1024-d). Ingest may also read `[embed] url`. Query embedding uses env only. +//! `generativelanguage.googleapis.com`, or LiteLLM OpenAI `/v1/embeddings`). +//! Ingest may also read `[embed] url`. Query embedding uses env only. +//! +//! The encoder's model id and width come from `[embed] model` / `[embed] dim` +//! (see [`crate::config::embed_model`] and [`crate::config::embed_dim`]), not +//! from a compiled-in constant. Every embedded row is stamped with the pair +//! actually used, so a store never claims an encoder that did not produce it. //! //! Three distinct text *views* are embedded per node so the vector arms are //! meaningful instead of three identical probes: @@ -25,13 +30,7 @@ use std::time::Duration; use serde_json::{Value, json}; use ureq::Agent; -use crate::config::{embed_model, embed_url}; - -/// Canonical encoder model id stamped on embedded rows. -pub(crate) const EMBED_MODEL: &str = "BAAI/bge-m3"; - -/// Dense vector dimension of [`EMBED_MODEL`]. -pub(crate) const EMBED_DIM: usize = 1024; +use crate::config::{embed_dim, embed_model, embed_url}; /// Maximum characters of any single view text sent to the encoder. /// @@ -163,15 +162,16 @@ pub(crate) fn embed_node_entities( let base = embed_url(repo_root) .ok_or_else(|| "no [embed] url / LEIO_CODE_EMBED_URL configured".to_string())?; - let model = embed_model(repo_root); + let model = resolved_model(&base, &embed_model(repo_root)); + let dim = embed_dim(repo_root); let code_views: Vec = entities.iter().map(build_code_view).collect(); let semantic_views: Vec = entities.iter().map(candidate_embedding_text).collect(); let ontology_views: Vec = entities.iter().map(build_ontology_view).collect(); - let code_vecs = embed_texts_remote(&base, &model, &code_views)?; - let semantic_vecs = embed_texts_remote(&base, &model, &semantic_views)?; - let ontology_vecs = embed_texts_remote(&base, &model, &ontology_views)?; + let code_vecs = embed_texts_remote(&base, &model, dim, &code_views)?; + let semantic_vecs = embed_texts_remote(&base, &model, dim, &semantic_views)?; + let ontology_vecs = embed_texts_remote(&base, &model, dim, &ontology_views)?; if code_vecs.len() != entities.len() || semantic_vecs.len() != entities.len() @@ -200,8 +200,10 @@ pub(crate) fn embed_node_entities( "ontology_vec".to_string(), float_json(&ontology_vecs[index]), ); - object.insert("embed_model".to_string(), json!(EMBED_MODEL)); - object.insert("embed_dim".to_string(), json!(EMBED_DIM)); + // Provenance is the configured pair, never a compiled-in default: a row + // must not claim BGE-M3 while holding vectors from another encoder. + object.insert("embed_model".to_string(), json!(model)); + object.insert("embed_dim".to_string(), json!(dim)); embedded += 1; } @@ -226,15 +228,20 @@ pub(crate) fn embed_texts(repo_root: &Path, texts: &[String]) -> Result Result>, String> { +fn embed_texts_remote( + base: &str, + model: &str, + dim: usize, + texts: &[String], +) -> Result>, String> { if texts.is_empty() { return Ok(Vec::new()); } @@ -248,9 +255,9 @@ fn embed_texts_remote(base: &str, model: &str, texts: &[String]) -> Result embed_openai_chunk(&agent, url, model, chunk)?, - EmbeddingEndpoint::Tei(url) => embed_tei_chunk(&agent, url, chunk)?, - EmbeddingEndpoint::Gemini(url) => embed_gemini_chunk(&agent, url, chunk)?, + EmbeddingEndpoint::OpenAi(url) => embed_openai_chunk(&agent, url, model, dim, chunk)?, + EmbeddingEndpoint::Tei(url) => embed_tei_chunk(&agent, url, dim, chunk)?, + EmbeddingEndpoint::Gemini(url) => embed_gemini_chunk(&agent, url, model, dim, chunk)?, }; out.extend(vectors); } @@ -274,10 +281,26 @@ enum EmbeddingEndpoint { /// Host of Google's Generative Language API embedding endpoint. const GEMINI_EMBED_HOST: &str = "generativelanguage.googleapis.com"; -/// Model id sent for Gemini embedding requests. `outputDimensionality: 1024` -/// keeps vectors on the same contract as BGE-M3 so cosine scoring stays valid. +/// Fallback model id for Gemini requests when `[embed] model` names no Gemini +/// model. `outputDimensionality` is the configured width, so the response +/// matches whatever the rest of the store was embedded at. const GEMINI_EMBED_MODEL: &str = "models/gemini-embedding-2"; +/// The model id actually sent to `base`, which is what rows are stamped with. +/// +/// Gemini only accepts its own `models/…` ids, so a configured id it cannot +/// serve is replaced by [`GEMINI_EMBED_MODEL`]. Resolving this once — rather +/// than substituting inside the request builder — keeps the request and the +/// stamped `embed_model` from disagreeing. +fn resolved_model(base: &str, configured: &str) -> String { + match embedding_endpoint(base) { + EmbeddingEndpoint::Gemini(_) if !configured.starts_with("models/") => { + GEMINI_EMBED_MODEL.to_string() + } + _ => configured.to_string(), + } +} + fn embedding_endpoint(base: &str) -> EmbeddingEndpoint { let normalized = base.trim_end_matches('/'); if normalized.contains(GEMINI_EMBED_HOST) { @@ -306,6 +329,7 @@ fn embed_openai_chunk( agent: &Agent, url: &str, model: &str, + dim: usize, texts: &[String], ) -> Result>, String> { let body = json!({ @@ -326,10 +350,15 @@ fn embed_openai_chunk( .body_mut() .read_json() .map_err(|err| format!("remote embed {url} decode: {err}"))?; - parse_openai_embeddings(&payload, texts.len()) + parse_openai_embeddings(&payload, dim, texts.len()) } -fn embed_tei_chunk(agent: &Agent, url: &str, texts: &[String]) -> Result>, String> { +fn embed_tei_chunk( + agent: &Agent, + url: &str, + dim: usize, + texts: &[String], +) -> Result>, String> { let body = json!({"inputs": texts}); let mut request = agent.post(url).header("Content-Type", "application/json"); if let Ok(key) = std::env::var("LEIO_CODE_EMBED_API_KEY") { @@ -345,27 +374,36 @@ fn embed_tei_chunk(agent: &Agent, url: &str, texts: &[String]) -> Result Value { +/// +/// `model` is sent verbatim; [`resolved_model`] has already mapped it to an id +/// this endpoint accepts, so the request and the stamped provenance agree. +fn gemini_batch_request(model: &str, dim: usize, texts: &[String]) -> Value { json!({ "requests": texts .iter() .map(|text| { json!({ - "model": GEMINI_EMBED_MODEL, + "model": model, "content": {"parts": [{"text": text}]}, - "outputDimensionality": EMBED_DIM, + "outputDimensionality": dim, }) }) .collect::>(), }) } -fn embed_gemini_chunk(agent: &Agent, url: &str, texts: &[String]) -> Result>, String> { - let body = gemini_batch_request(texts); +fn embed_gemini_chunk( + agent: &Agent, + url: &str, + model: &str, + dim: usize, + texts: &[String], +) -> Result>, String> { + let body = gemini_batch_request(model, dim, texts); let mut request = agent.post(url).header("Content-Type", "application/json"); let key = std::env::var("LEIO_CODE_EMBED_API_KEY") .ok() @@ -380,10 +418,14 @@ fn embed_gemini_chunk(agent: &Agent, url: &str, texts: &[String]) -> Result Result>, String> { +fn parse_openai_embeddings( + payload: &Value, + dim: usize, + expected: usize, +) -> Result>, String> { let data = payload .get("data") .and_then(Value::as_array) @@ -417,13 +459,13 @@ fn parse_openai_embeddings(payload: &Value, expected: usize) -> Result Result Result>, String> { +fn parse_tei_embeddings( + payload: &Value, + dim: usize, + expected: usize, +) -> Result>, String> { let rows = payload .as_array() .ok_or_else(|| format!("remote TEI embed expected vector[]: {payload}"))?; @@ -455,9 +501,9 @@ fn parse_tei_embeddings(payload: &Value, expected: usize) -> Result let values = row .as_array() .ok_or_else(|| format!("remote TEI embed index {index} is not a vector"))?; - if values.len() != EMBED_DIM { + if values.len() != dim { return Err(format!( - "remote TEI embed index {index} dim {} (expected {EMBED_DIM})", + "remote TEI embed index {index} dim {} (expected {dim} for the configured encoder)", values.len() )); } @@ -475,7 +521,11 @@ fn parse_tei_embeddings(payload: &Value, expected: usize) -> Result /// Parse a Gemini `batchEmbedContents` response. Embeddings come back in /// request order under `embeddings[].values`. -fn parse_gemini_embeddings(payload: &Value, expected: usize) -> Result>, String> { +fn parse_gemini_embeddings( + payload: &Value, + dim: usize, + expected: usize, +) -> Result>, String> { let rows = payload .get("embeddings") .and_then(Value::as_array) @@ -501,9 +551,9 @@ fn parse_gemini_embeddings(payload: &Value, expected: usize) -> Result Value { #[cfg(test)] mod tests { use super::*; + use crate::config::{DEFAULT_EMBED_DIM, DEFAULT_EMBED_MODEL}; use serde_json::json; + /// Width these fixtures are built at. The parsers take it as an argument + /// now, so the value is a fixture detail rather than a global contract. + const EMBED_DIM: usize = DEFAULT_EMBED_DIM; + #[test] fn code_view_prefers_symbol_then_kind() { let entity = json!({ @@ -637,7 +692,7 @@ mod tests { {"index": 0, "embedding": vec![1.0; EMBED_DIM]}, ] }); - let vectors = parse_openai_embeddings(&payload, 2).expect("parse"); + let vectors = parse_openai_embeddings(&payload, EMBED_DIM, 2).expect("parse"); assert_eq!(vectors.len(), 2); assert!((vectors[0][0] - 1.0).abs() < f32::EPSILON); assert!(vectors[1][0].abs() < f32::EPSILON); @@ -646,7 +701,7 @@ mod tests { #[test] fn parse_tei_embeddings_preserves_input_order() { let payload = json!([vec![0.0; EMBED_DIM], vec![1.0; EMBED_DIM]]); - let vectors = parse_tei_embeddings(&payload, 2).expect("parse"); + let vectors = parse_tei_embeddings(&payload, EMBED_DIM, 2).expect("parse"); assert_eq!(vectors.len(), 2); assert!(vectors[0][0].abs() < f32::EPSILON); assert!((vectors[1][0] - 1.0).abs() < f32::EPSILON); @@ -655,21 +710,21 @@ mod tests { #[test] fn parse_tei_embeddings_rejects_wrong_count() { let payload = json!([vec![0.0; EMBED_DIM]]); - let error = parse_tei_embeddings(&payload, 2).expect_err("count mismatch"); + let error = parse_tei_embeddings(&payload, EMBED_DIM, 2).expect_err("count mismatch"); assert!(error.contains("count mismatch")); } #[test] fn parse_tei_embeddings_rejects_wrong_dimension() { let payload = json!([vec![0.0; EMBED_DIM - 1]]); - let error = parse_tei_embeddings(&payload, 1).expect_err("dimension mismatch"); + let error = parse_tei_embeddings(&payload, EMBED_DIM, 1).expect_err("dimension mismatch"); assert!(error.contains("dim 1023")); } #[test] fn parse_tei_embeddings_rejects_non_numeric_components() { let payload = json!([vec![json!("not-a-number"); EMBED_DIM]]); - let error = parse_tei_embeddings(&payload, 1).expect_err("numeric component"); + let error = parse_tei_embeddings(&payload, EMBED_DIM, 1).expect_err("numeric component"); assert!(error.contains("non-float component")); } @@ -721,7 +776,7 @@ mod tests { {"index": 1, "embedding": vec![1.0; EMBED_DIM]}, ] }); - let error = parse_openai_embeddings(&payload, 2).expect_err("count mismatch"); + let error = parse_openai_embeddings(&payload, EMBED_DIM, 2).expect_err("count mismatch"); assert!(error.contains("count mismatch")); } @@ -751,7 +806,7 @@ mod tests { {"values": vec![1.0; EMBED_DIM]}, ] }); - let vectors = parse_gemini_embeddings(&payload, 2).expect("parse"); + let vectors = parse_gemini_embeddings(&payload, EMBED_DIM, 2).expect("parse"); assert_eq!(vectors.len(), 2); assert!(vectors[0][0].abs() < f32::EPSILON); assert!((vectors[1][0] - 1.0).abs() < f32::EPSILON); @@ -760,24 +815,101 @@ mod tests { #[test] fn parse_gemini_embeddings_rejects_wrong_count() { let payload = json!({"embeddings": [{"values": vec![0.0; EMBED_DIM]}]}); - let error = parse_gemini_embeddings(&payload, 2).expect_err("count mismatch"); + let error = parse_gemini_embeddings(&payload, EMBED_DIM, 2).expect_err("count mismatch"); assert!(error.contains("count mismatch")); } #[test] fn parse_gemini_embeddings_rejects_wrong_dimension() { let payload = json!({"embeddings": [{"values": vec![0.0; EMBED_DIM - 1]}]}); - let error = parse_gemini_embeddings(&payload, 1).expect_err("dimension mismatch"); + let error = + parse_gemini_embeddings(&payload, EMBED_DIM, 1).expect_err("dimension mismatch"); assert!(error.contains("dim 1023")); } #[test] fn gemini_batch_request_wraps_texts_with_model_and_dimension() { let texts = vec!["alpha".to_string(), "beta".to_string()]; - let body = gemini_batch_request(&texts); + let body = gemini_batch_request("models/gemini-embedding-2", EMBED_DIM, &texts); assert_eq!(body["requests"].as_array().unwrap().len(), 2); assert_eq!(body["requests"][0]["model"], "models/gemini-embedding-2"); assert_eq!(body["requests"][0]["content"]["parts"][0]["text"], "alpha"); assert_eq!(body["requests"][0]["outputDimensionality"], EMBED_DIM); } + + #[test] + fn parsers_accept_a_non_default_configured_width() { + // `google/embeddinggemma-2` is 768-d: the width the parsers were + // hardcoded to reject before it became configurable. + const GEMMA_DIM: usize = 768; + assert_ne!(GEMMA_DIM, DEFAULT_EMBED_DIM); + + let openai = json!({"data": [{"index": 0, "embedding": vec![1.0; GEMMA_DIM]}]}); + let vectors = parse_openai_embeddings(&openai, GEMMA_DIM, 1).expect("openai parse"); + assert_eq!(vectors[0].len(), GEMMA_DIM); + + let tei = json!([vec![1.0; GEMMA_DIM]]); + let vectors = parse_tei_embeddings(&tei, GEMMA_DIM, 1).expect("tei parse"); + assert_eq!(vectors[0].len(), GEMMA_DIM); + + let gemini = json!({"embeddings": [{"values": vec![1.0; GEMMA_DIM]}]}); + let vectors = parse_gemini_embeddings(&gemini, GEMMA_DIM, 1).expect("gemini parse"); + assert_eq!(vectors[0].len(), GEMMA_DIM); + } + + #[test] + fn parsers_reject_a_width_other_than_the_configured_one() { + // A 1024-d response against a 768-d store is the silent-corruption + // case: both widths are individually valid, only the pairing is wrong. + const GEMMA_DIM: usize = 768; + + let openai = json!({"data": [{"index": 0, "embedding": vec![1.0; DEFAULT_EMBED_DIM]}]}); + let error = parse_openai_embeddings(&openai, GEMMA_DIM, 1).expect_err("openai mismatch"); + assert!(error.contains("dim 1024"), "{error}"); + assert!(error.contains("expected 768"), "{error}"); + + let tei = json!([vec![1.0; DEFAULT_EMBED_DIM]]); + let error = parse_tei_embeddings(&tei, GEMMA_DIM, 1).expect_err("tei mismatch"); + assert!(error.contains("expected 768"), "{error}"); + + let gemini = json!({"embeddings": [{"values": vec![1.0; DEFAULT_EMBED_DIM]}]}); + let error = parse_gemini_embeddings(&gemini, GEMMA_DIM, 1).expect_err("gemini mismatch"); + assert!(error.contains("expected 768"), "{error}"); + } + + #[test] + fn gemini_batch_request_carries_the_configured_width_and_model() { + let texts = vec!["alpha".to_string()]; + + let body = gemini_batch_request("models/gemini-embedding-2", 768, &texts); + assert_eq!(body["requests"][0]["outputDimensionality"], 768); + assert_eq!(body["requests"][0]["model"], "models/gemini-embedding-2"); + } + + #[test] + fn stamped_model_is_the_model_actually_sent() { + const GEMINI: &str = "https://generativelanguage.googleapis.com/v1beta/models/x"; + + // A non-Gemini id cannot be served by this endpoint, so it is replaced + // *before* the stamp is taken — otherwise a row embedded by Gemini + // would claim to be bge-m3, the false provenance this guards against. + let sent = resolved_model(GEMINI, DEFAULT_EMBED_MODEL); + assert_eq!(sent, GEMINI_EMBED_MODEL); + let body = gemini_batch_request(&sent, DEFAULT_EMBED_DIM, &["alpha".to_string()]); + assert_eq!(body["requests"][0]["model"], sent); + + // A Gemini id the endpoint accepts travels through untouched. + let sent = resolved_model(GEMINI, "models/gemini-embedding-2"); + assert_eq!(sent, "models/gemini-embedding-2"); + + // Other endpoints never rewrite the configured id. + assert_eq!( + resolved_model("http://localhost:8080/embed", "google/embeddinggemma-2"), + "google/embeddinggemma-2" + ); + assert_eq!( + resolved_model("http://localhost:4000", "google/embeddinggemma-2"), + "google/embeddinggemma-2" + ); + } } diff --git a/src/export.rs b/src/export.rs index 8c57fda..fb74853 100644 --- a/src/export.rs +++ b/src/export.rs @@ -690,11 +690,11 @@ pub fn export_arrow_nodes( "transport": "arrow_ipc", "fca": fca_stats, "embed": { - "model": crate::embed::EMBED_MODEL, - "dim": crate::embed::EMBED_DIM, + "model": crate::config::embed_model(root), + "dim": crate::config::embed_dim(root), "embedded_rows": embedded_rows, "total_rows": row_count, - "encoder": "remote_bge_m3", + "encoder": "remote_dense", "warnings": embed_warnings, }, })), diff --git a/src/node_rows.rs b/src/node_rows.rs index 793d7de..413ba1d 100644 --- a/src/node_rows.rs +++ b/src/node_rows.rs @@ -22,7 +22,7 @@ use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as Arr use serde_json::{Value, json}; use sha2::{Digest, Sha256}; -use crate::embed::{EMBED_DIM, EMBED_MODEL}; +use crate::config::DEFAULT_EMBED_MODEL; use crate::model::RepoIndex; use crate::query::{ AggregatedApiRoute, DockerServiceCandidate, collect_api_routes, collect_docker_services, @@ -229,7 +229,9 @@ fn build_execution_text(kind: &str, name: &str, path: &str, lang: &str) -> Strin /// Build a node entity JSON with zero-placeholder dense vector fields. /// /// The three float vectors (`code_vec`/`semantic_vec`/`ontology_vec`) start as -/// zero placeholders at the BGE-M3 contract dimension (`EMBED_DIM`). An export +/// zero placeholders `vector_dim` wide, the width of the configured encoder +/// ([`crate::config::embed_dim`]), so a later embedding pass overwrites them +/// in place rather than changing the row's shape. An export /// with an embed URL may replace them in one batched pass via /// [`crate::embed::embed_node_entities`], using three distinct text views; /// without a usable encoder, the zero placeholders are written and the row @@ -240,6 +242,7 @@ fn build_execution_text(kind: &str, name: &str, path: &str, lang: &str) -> Strin #[allow(clippy::too_many_arguments)] fn build_node_entity( repo_root: &Path, + vector_dim: usize, logical_node_id: &str, kind: &str, symbol: &str, @@ -274,9 +277,9 @@ fn build_node_entity( "text_snippet": &text_snippet[..text_snippet.len().min(2000)], "embed_model": "none", "embed_dim": 0, - "code_vec": vec![0.0f64; EMBED_DIM], - "semantic_vec": vec![0.0f64; EMBED_DIM], - "ontology_vec": vec![0.0f64; EMBED_DIM], + "code_vec": vec![0.0f64; vector_dim], + "semantic_vec": vec![0.0f64; vector_dim], + "ontology_vec": vec![0.0f64; vector_dim], "execution_vec_bin": exec_bin, }) } @@ -286,6 +289,9 @@ pub(crate) fn visit_node_entities(index: &RepoIndex, repo_root: &Path, mut on where F: FnMut(Value), { + // Resolved once per walk: every row must share one width, and the lookup + // reads `.leio-code/config.toml` from disk. + let vector_dim = crate::config::embed_dim(repo_root); let profile_cartridges = profile_active_cartridges(index); let cartridge_profiles = cartridge_activation_profiles(index); let routes = collect_api_routes(index); @@ -302,6 +308,7 @@ where ); let mut entity = build_node_entity( repo_root, + vector_dim, &file_id, "file", "", @@ -347,6 +354,7 @@ where ); let mut entity = build_node_entity( repo_root, + vector_dim, &logical_id, "symbol", &sym.name, @@ -385,6 +393,7 @@ where ); let mut entity = build_node_entity( repo_root, + vector_dim, &logical_id, "env_var", &ev.name, @@ -423,6 +432,7 @@ where ); let mut entity = build_node_entity( repo_root, + vector_dim, &logical_id, "redis_key", &rk.key, @@ -465,6 +475,7 @@ where ); let mut entity = build_node_entity( repo_root, + vector_dim, &logical_id, "deploy_target", &dt.name, @@ -539,6 +550,7 @@ where ); let mut entity = build_node_entity( repo_root, + vector_dim, &logical_id, "cartridge", &cartridge, @@ -603,6 +615,7 @@ where ); let mut entity = build_node_entity( repo_root, + vector_dim, &logical_id, "profile", &profile.name, @@ -648,6 +661,7 @@ where ); let mut entity = build_node_entity( repo_root, + vector_dim, &logical_id, "secret_set", &secret_set.name, @@ -694,6 +708,7 @@ where ); let mut entity = build_node_entity( repo_root, + vector_dim, &logical_id, "api_route", &route.full_path, @@ -757,6 +772,7 @@ where ); let mut entity = build_node_entity( repo_root, + vector_dim, &logical_id, "docker_service", &service.name, @@ -922,7 +938,10 @@ pub(crate) fn build_leio_row_batch(entities: &[Value]) -> Result Result, String> { #[cfg(test)] mod tests { use super::*; + use crate::config::DEFAULT_EMBED_DIM; /// Env mutation in these tests is unsound under parallel test threads, so /// every scoping test serializes on this lock and restores prior values. @@ -1115,6 +1135,7 @@ mod tests { let restore = EnvRestore::capture(); let entity = build_node_entity( Path::new("/tmp/demo-repo"), + DEFAULT_EMBED_DIM, "file:src/main.rs", "file", "", @@ -1152,6 +1173,7 @@ mod tests { } build_node_entity( Path::new("/tmp/shared-repo"), + DEFAULT_EMBED_DIM, "file:src/main.rs", "file", "",