diff --git a/Cargo.lock b/Cargo.lock index 6cb1fe9..6fd2335 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -782,7 +782,7 @@ dependencies = [ [[package]] name = "codegraph-harness" -version = "0.18.6" +version = "0.19.0" dependencies = [ "anyhow", "clap", @@ -890,7 +890,7 @@ dependencies = [ [[package]] name = "codegraph-memory" -version = "0.18.6" +version = "0.19.0" dependencies = [ "anyhow", "bincode", @@ -898,16 +898,19 @@ dependencies = [ "criterion", "dashmap", "fastembed", + "half", "hf-hub", "instant-distance", "log", "parking_lot", "rmp-serde", "rocksdb", + "safetensors", "serde", "serde_json", "tempfile", "thiserror 1.0.69", + "tokenizers", "tokio", "tokio-test", "ureq", @@ -1070,7 +1073,7 @@ dependencies = [ [[package]] name = "codegraph-server" -version = "0.18.6" +version = "0.19.0" dependencies = [ "clap", "codegraph", @@ -1612,6 +1615,9 @@ name = "esaxx-rs" version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] [[package]] name = "exr" @@ -3638,6 +3644,16 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "safetensors" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44560c11236a6130a46ce36c836a62936dc81ebf8c36a37947423571be0e55b6" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "same-file" version = "1.0.6" @@ -4046,6 +4062,7 @@ dependencies = [ "derive_builder", "esaxx-rs", "getrandom 0.3.4", + "indicatif", "itertools 0.14.0", "log", "macro_rules_attribute", diff --git a/Cargo.toml b/Cargo.toml index 0060566..922e756 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -60,7 +60,7 @@ members = [ ] [workspace.package] -version = "0.18.6" +version = "0.19.0" edition = "2021" license = "Apache-2.0" repository = "https://github.com/codegraph-ai/codegraph" diff --git a/README.md b/README.md index 2ba9c1a..3e37728 100644 --- a/README.md +++ b/README.md @@ -77,13 +77,29 @@ one tool and exits without the MCP stdio handshake — ideal for scripting. |------|---------|-------------| | `--workspace ` | current dir | Directories to index (repeatable for multi-project) | | `--exclude ` | — | Directories to skip (repeatable) | -| `--embedding-model ` | `bge-small` | `bge-small` (384d, fast), `jina-code-v2` (768d, 6× slower), or `granite-97m` (384d, 32K ctx, ~3× slower) | +| `--embedding-model ` | `bge-small` | `bge-small` (384d, fast), `jina-code-v2` (768d, 6× slower), `granite-97m` (384d, 32K ctx, ~3× slower), or `static` (model2vec, 256d — ~100× faster indexing, no ONNX; needs a local model dir, see below) | | `--full-body-embedding` | `true` | Embed full function body (~50 lines) for better semantic search and duplicate detection | | `--max-files ` | 5000 | Maximum files to index | | `--profile ` | `all` | Filter the exposed MCP tool surface to a named subset (see below) | | `--graph-only` | off | Skip embedding generation — build the graph and serve structural tools only. No ONNX model load, 10-50× faster indexing. Semantic search unavailable. For CI / one-shot graph queries. | | `--run-tool ` | — | One-shot mode: index, run a single tool, print its result, exit. No MCP handshake. Pair with `--tool-args ''`. | +#### `--embedding-model static` — model2vec fast indexing + +Static (model2vec) embeddings replace the ONNX transformer with a token→vector +lookup table: indexing is **~100× faster** (this repo's 5,873 symbols embed in +~1 s vs ~3.4 min with BGE) and there's **no ONNX runtime or 1.5 GB RAM gate**. +Retrieval stays **hybrid (BM25 + semantic)**, so end-to-end quality is **~90% of +BGE**. The VS Code extension ships the model bundled, so `static` works there +with no setup. For the CLI/MCP server it needs a local model directory +(`config.json` + `tokenizer.json` + `model.safetensors`): + +- Point at it with `CODEGRAPH_STATIC_MODEL=/path/to/model` (or the VS Code + `codegraph.staticModelPath` setting to override the bundled model). Default: + `~/.codegraph/static_models/jina-code-static-256`. +- Distill one from any sentence-transformer (Apache-2.0 Jina-Code by default) in + ~30 s on CPU: `python scripts/distill_static_model.py`. + #### `--profile` — narrow the MCP tool surface The full 32-tool surface is convenient but inflates the agent's prompt-context cost. A profile exposes only the slice you need (also settable via the `CODEGRAPH_TOOL_PROFILE` env var): @@ -103,7 +119,8 @@ The full 32-tool surface is convenient but inflates the agent's prompt-context c "codegraph.indexOnStartup": true, "codegraph.indexPaths": ["/path/to/project-a", "/path/to/project-b"], "codegraph.excludePatterns": ["**/cmake-build-debug/**", "**/generated/**"], - "codegraph.embeddingModel": "bge-small", + "codegraph.embeddingModel": "bge-small", // or "static" for ~100× faster indexing + "codegraph.staticModelPath": "", // model2vec model dir when embeddingModel is "static" "codegraph.maxFileSizeKB": 1024, "codegraph.debug": false } @@ -313,6 +330,12 @@ Additional tools available in [CodeGraph Pro](https://codegraph.astudioplus.com/ HTTP handler detection: Python (FastAPI/Flask/Django), TypeScript (NestJS), Java (Spring/JAX-RS), Go (stdlib/Gin/Echo/Fiber), C# (ASP.NET), Ruby (Rails), PHP (Laravel/Symfony). +> **Community vs full builds:** COBOL, Fortran, Perl, Dart, Zig, and R are +> compiled only with `--features extra-languages`. The default community binary +> omits them — they had zero usage in telemetry and their tree-sitter grammars +> add ~25 MB (COBOL's parse tables alone are 30 MB). The other 32 languages are +> always available. + --- ## Architecture diff --git a/crates/codegraph-memory/Cargo.toml b/crates/codegraph-memory/Cargo.toml index 364ea21..379d141 100644 --- a/crates/codegraph-memory/Cargo.toml +++ b/crates/codegraph-memory/Cargo.toml @@ -46,6 +46,12 @@ anyhow = "1.0" # Logging log = "0.4" +# Static (lookup-table) embeddings: HuggingFace tokenizer + a safetensors +# token->vector matrix, mean-pooled. No ONNX — the fast indexing path. +tokenizers = "0.21" +safetensors = "0.4" +half = "2" + # Embeddings - fastembed with BGE-Small-EN-v1.5 # macOS/Linux: static link ONNX Runtime (ort-download-binaries) # Windows: load onnxruntime DLL at runtime (ort-load-dynamic, avoids CRT /MT vs /MD mismatch) diff --git a/crates/codegraph-memory/examples/embed_eval.rs b/crates/codegraph-memory/examples/embed_eval.rs new file mode 100644 index 0000000..2215682 --- /dev/null +++ b/crates/codegraph-memory/examples/embed_eval.rs @@ -0,0 +1,255 @@ +// Copyright 2025-2026 Andrey Vasilevsky +// SPDX-License-Identifier: Apache-2.0 + +//! Real-corpus retrieval eval: static (model2vec) vs ONNX BGE baseline, scored +//! both **pure-semantic** and **hybrid (40% BM25 + 60% semantic)** — the latter +//! mirrors the real `symbol_search` blend, which is what actually ships. +//! +//! Reads a `[{id, signature, doc}]` corpus (from `scripts/extract_eval_corpus.py`): +//! embed `"id: signature"` per symbol, use each symbol's `doc` as the query, +//! target = that symbol. Same task for every embedder → fair comparison. +//! +//! python scripts/extract_eval_corpus.py crates /tmp/codegraph_eval_corpus.json +//! CODEGRAPH_STATIC_MODEL=~/.codegraph/static_models/jina-code-static-256 \ +//! cargo run --release -p codegraph-memory --example embed_eval +//! +//! Env: CODEGRAPH_EVAL_CORPUS (path), CODEGRAPH_SPLIT_IDS=1 (prepend split name). + +use codegraph_memory::embedding::{CodeGraphEmbeddingModel, VectorEngine}; +use serde::Deserialize; +use std::collections::HashMap; + +#[derive(Deserialize)] +struct Sym { + id: String, + signature: String, + doc: String, +} + +struct Scores { + r1: f64, + r5: f64, + r10: f64, + mrr: f64, +} + +/// Split an identifier into camelCase/snake_case words (lowercased). +fn split_id(id: &str) -> String { + let mut out = String::new(); + let mut prev_upper = false; + for ch in id.chars() { + if ch == '_' || ch == '-' { + out.push(' '); + } else if ch.is_uppercase() && !out.is_empty() && !prev_upper { + out.push(' '); + out.extend(ch.to_lowercase()); + } else { + out.extend(ch.to_lowercase()); + } + prev_upper = ch.is_uppercase(); + } + out +} + +/// Tokenize on non-alphanumerics + camelCase, lowercase, drop 1-char tokens. +fn tokenize(s: &str) -> Vec { + let mut toks = Vec::new(); + let mut cur = String::new(); + let mut prev_upper = false; + for ch in s.chars() { + if ch.is_alphanumeric() { + if ch.is_uppercase() && !cur.is_empty() && !prev_upper { + toks.push(std::mem::take(&mut cur)); + } + cur.extend(ch.to_lowercase()); + prev_upper = ch.is_uppercase(); + } else { + if !cur.is_empty() { + toks.push(std::mem::take(&mut cur)); + } + prev_upper = false; + } + } + if !cur.is_empty() { + toks.push(cur); + } + toks.retain(|t| t.len() >= 2); + toks +} + +/// Standard BM25 over the symbol texts, with per-doc term frequencies precomputed. +struct Bm25 { + tf: Vec>, + dl: Vec, + df: HashMap, + avgdl: f64, + n: f64, +} + +impl Bm25 { + fn new(texts: &[String]) -> Self { + let mut df: HashMap = HashMap::new(); + let mut tf = Vec::with_capacity(texts.len()); + let mut dl = Vec::with_capacity(texts.len()); + for t in texts { + let toks = tokenize(t); + let mut m: HashMap = HashMap::new(); + for tok in &toks { + *m.entry(tok.clone()).or_insert(0) += 1; + } + for k in m.keys() { + *df.entry(k.clone()).or_insert(0) += 1; + } + dl.push(toks.len() as f64); + tf.push(m); + } + let n = texts.len() as f64; + let avgdl = dl.iter().sum::() / n.max(1.0); + Self { tf, dl, df, avgdl, n } + } + + fn score(&self, q: &[String], i: usize) -> f64 { + let (k1, b) = (1.2_f64, 0.75_f64); + let mut s = 0.0; + for term in q { + let f = *self.tf[i].get(term).unwrap_or(&0) as f64; + if f == 0.0 { + continue; + } + let df = *self.df.get(term).unwrap_or(&0) as f64; + let idf = ((self.n - df + 0.5) / (df + 0.5) + 1.0).ln(); + s += idf * (f * (k1 + 1.0)) / (f + k1 * (1.0 - b + b * self.dl[i] / self.avgdl)); + } + s + } +} + +fn minmax(v: &mut [f64]) { + let (mn, mx) = v + .iter() + .fold((f64::MAX, f64::MIN), |(a, b), &x| (a.min(x), b.max(x))); + let r = mx - mn; + for x in v.iter_mut() { + *x = if r > 0.0 { (*x - mn) / r } else { 0.0 }; + } +} + +fn rank_of(target: usize, scores: &[f64]) -> usize { + let t = scores[target]; + 1 + scores + .iter() + .enumerate() + .filter(|(j, &s)| *j != target && s > t) + .count() +} + +fn metrics(ranks: &[usize]) -> Scores { + let n = ranks.len() as f64; + let (mut r1, mut r5, mut r10, mut mrr) = (0.0, 0.0, 0.0, 0.0); + for &r in ranks { + if r <= 1 { + r1 += 1.0; + } + if r <= 5 { + r5 += 1.0; + } + if r <= 10 { + r10 += 1.0; + } + mrr += 1.0 / r as f64; + } + Scores { + r1: r1 / n, + r5: r5 / n, + r10: r10 / n, + mrr: mrr / n, + } +} + +/// Returns (semantic-only, hybrid 0.4*BM25 + 0.6*cosine) scores. +fn evaluate(engine: &VectorEngine, syms: &[Sym]) -> (Scores, Scores) { + let split = std::env::var("CODEGRAPH_SPLIT_IDS").is_ok(); + let sym_texts: Vec = syms + .iter() + .map(|s| { + if split { + format!("{} {}: {}", split_id(&s.id), s.id, s.signature) + } else { + format!("{}: {}", s.id, s.signature) + } + }) + .collect(); + let sym_refs: Vec<&str> = sym_texts.iter().map(|s| s.as_str()).collect(); + let sym_vecs = engine.embed_batch(&sym_refs).expect("embed symbols"); + let q_refs: Vec<&str> = syms.iter().map(|s| s.doc.as_str()).collect(); + let q_vecs = engine.embed_batch(&q_refs).expect("embed queries"); + + let bm25 = Bm25::new(&sym_texts); + let mut sem_ranks = Vec::with_capacity(syms.len()); + let mut hyb_ranks = Vec::with_capacity(syms.len()); + + for (i, qv) in q_vecs.iter().enumerate() { + let mut cos: Vec = sym_vecs + .iter() + .map(|sv| engine.similarity(qv, sv) as f64) + .collect(); + sem_ranks.push(rank_of(i, &cos)); + + let qterms = tokenize(&syms[i].doc); + let mut bm: Vec = (0..sym_texts.len()).map(|j| bm25.score(&qterms, j)).collect(); + minmax(&mut bm); + minmax(&mut cos); + let hyb: Vec = bm.iter().zip(&cos).map(|(b, c)| 0.4 * b + 0.6 * c).collect(); + hyb_ranks.push(rank_of(i, &hyb)); + } + + (metrics(&sem_ranks), metrics(&hyb_ranks)) +} + +fn report(label: &str, engine: Result, syms: &[Sym]) { + match engine { + Ok(engine) => { + let (sem, hyb) = evaluate(&engine, syms); + println!( + "{label:<8} {:<28} semantic R@1 {:.3} R@5 {:.3} R@10 {:.3} MRR {:.3}", + engine.model_name(), + sem.r1, + sem.r5, + sem.r10, + sem.mrr + ); + println!( + "{:<8} {:<28} HYBRID R@1 {:.3} R@5 {:.3} R@10 {:.3} MRR {:.3}", + "", + "(0.4 bm25 + 0.6 cos)", + hyb.r1, + hyb.r5, + hyb.r10, + hyb.mrr + ); + } + Err(e) => println!("{label:<8} skipped ({e})"), + } +} + +fn main() { + let home = std::env::var("HOME").unwrap_or_default(); + let corpus_path = std::env::var("CODEGRAPH_EVAL_CORPUS") + .unwrap_or_else(|_| "/tmp/codegraph_eval_corpus.json".to_string()); + let syms: Vec = serde_json::from_slice( + &std::fs::read(&corpus_path).unwrap_or_else(|e| panic!("read {corpus_path}: {e}")), + ) + .expect("parse corpus json"); + println!( + "Retrieval eval: {} symbols (doc -> symbol), corpus {corpus_path}\n", + syms.len() + ); + + let static_dir = std::env::var("CODEGRAPH_STATIC_MODEL") + .map(std::path::PathBuf::from) + .unwrap_or_else(|_| std::path::PathBuf::from(&home).join(".codegraph/static_models/potion-base-8M")); + report("static", VectorEngine::with_static_model(&static_dir), &syms); + + let cache = std::path::PathBuf::from(&home).join(".codegraph/fastembed_cache"); + report("onnx-bge", VectorEngine::with_model(cache, CodeGraphEmbeddingModel::BgeSmall), &syms); +} diff --git a/crates/codegraph-memory/examples/embed_quality.rs b/crates/codegraph-memory/examples/embed_quality.rs new file mode 100644 index 0000000..5c2fa07 --- /dev/null +++ b/crates/codegraph-memory/examples/embed_quality.rs @@ -0,0 +1,125 @@ +// Copyright 2025-2026 Andrey Vasilevsky +// SPDX-License-Identifier: Apache-2.0 + +//! Retrieval-quality micro-eval: static (model2vec) vs ONNX transformer. +//! +//! A small, hand-built `query -> target symbol` set scored by recall@1, recall@3 +//! and MRR — the Phase-0 eval in miniature. It answers, directionally: how far +//! below BGE does the *generic* static floor (`potion-base-8M`) sit? That is the +//! gap a code-distilled Jina static would need to close. NOT a substitute for +//! the full eval on a real indexed repo (12 queries is noisy) — a first signal. +//! +//! cargo run --release -p codegraph-memory --example embed_quality +//! +//! Needs `~/.codegraph/static_models/potion-base-8M`; BGE loads from the +//! fastembed cache. + +use codegraph_memory::embedding::{CodeGraphEmbeddingModel, VectorEngine}; + +/// (symbol id, "name: signature — doc" embed text) +const SYMBOLS: &[(&str, &str)] = &[ + ("authenticate_user", "authenticate_user: fn authenticate_user(token: &Token) -> Result — Verify the bearer token and return the user"), + ("hash_password", "hash_password: fn hash_password(plain: &str) -> String — Bcrypt-hash a plaintext password before storage"), + ("open_database_connection", "open_database_connection: fn open_database_connection(url: &str) -> Pool — Open a pooled postgres connection"), + ("execute_sql_query", "execute_sql_query: fn execute_sql_query(sql: &str, params: &[Value]) -> Rows — Run a parameterized SQL statement"), + ("parse_json_request", "parse_json_request: fn parse_json_request(body: &[u8]) -> Request — Deserialize the HTTP request body as JSON"), + ("serialize_response", "serialize_response: fn serialize_response(resp: &Response) -> String — Encode a response struct to JSON text"), + ("compute_cosine_similarity", "compute_cosine_similarity: fn compute_cosine_similarity(a: &[f32], b: &[f32]) -> f32 — Dot product over norms"), + ("build_hnsw_index", "build_hnsw_index: fn build_hnsw_index(vectors: &[Vec]) -> Hnsw — Construct an approximate nearest-neighbor index"), + ("retry_with_backoff", "retry_with_backoff: fn retry_with_backoff(op: impl Fn() -> Result) -> T — Retry a fallible op with exponential backoff"), + ("rate_limiter", "rate_limiter: struct RateLimiter { tokens: u32 } — Token-bucket throttling of incoming requests"), + ("parse_config_file", "parse_config_file: fn parse_config_file(path: &Path) -> Config — Read and deserialize the TOML configuration file"), + ("watch_file_changes", "watch_file_changes: fn watch_file_changes(dir: &Path) -> Receiver — Notify on filesystem modifications"), + ("spawn_worker_pool", "spawn_worker_pool: fn spawn_worker_pool(n: usize, q: Arc) -> Vec — Start N background consumer threads"), + ("validate_email", "validate_email: fn validate_email(input: &str) -> bool — Lenient RFC-5322 email-address check"), + ("encode_jwt_token", "encode_jwt_token: fn encode_jwt_token(claims: &Claims, secret: &[u8]) -> String — Sign a JSON Web Token with claims"), + ("cache_get_or_insert", "cache_get_or_insert: fn cache_get_or_insert(key: K, f: impl Fn() -> V) -> V — Memoized lookup with fallback compute"), +]; + +/// (natural-language query, target symbol id) +const QUERIES: &[(&str, &str)] = &[ + ("check if a login token is valid", "authenticate_user"), + ("store a user's password securely", "hash_password"), + ("connect to the database", "open_database_connection"), + ("run a SQL statement against the db", "execute_sql_query"), + ("turn the request body into an object", "parse_json_request"), + ("find the nearest vectors quickly", "build_hnsw_index"), + ("measure how similar two embeddings are", "compute_cosine_similarity"), + ("try the operation again if it fails", "retry_with_backoff"), + ("limit how many requests per second", "rate_limiter"), + ("load settings from a config file", "parse_config_file"), + ("sign a json web token", "encode_jwt_token"), + ("check an email address is well formed", "validate_email"), +]; + +struct Scores { + recall_at_1: f64, + recall_at_3: f64, + mrr: f64, +} + +fn evaluate(engine: &VectorEngine) -> Scores { + let sym_texts: Vec<&str> = SYMBOLS.iter().map(|(_, t)| *t).collect(); + let sym_vecs = engine.embed_batch(&sym_texts).expect("embed symbols"); + + let (mut r1, mut r3, mut mrr) = (0.0, 0.0, 0.0); + for (query, target) in QUERIES { + let qv = engine.embed(query).expect("embed query"); + // Rank symbols by cosine similarity to the query. + let mut ranked: Vec<(usize, f32)> = sym_vecs + .iter() + .enumerate() + .map(|(i, sv)| (i, engine.similarity(&qv, sv))) + .collect(); + ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + let rank = 1 + ranked.iter().position(|(i, _)| SYMBOLS[*i].0 == *target).unwrap(); + if rank == 1 { + r1 += 1.0; + } + if rank <= 3 { + r3 += 1.0; + } + mrr += 1.0 / rank as f64; + } + let n = QUERIES.len() as f64; + Scores { + recall_at_1: r1 / n, + recall_at_3: r3 / n, + mrr: mrr / n, + } +} + +fn report(label: &str, engine: Result) { + match engine { + Ok(engine) => { + let s = evaluate(&engine); + println!( + "{label:<8} {:<30} R@1 {:.2} R@3 {:.2} MRR {:.3}", + engine.model_name(), + s.recall_at_1, + s.recall_at_3, + s.mrr + ); + } + Err(e) => println!("{label:<8} skipped ({e})"), + } +} + +fn main() { + let home = std::env::var("HOME").unwrap_or_default(); + println!( + "Retrieval micro-eval: {} queries over {} symbols.\n", + QUERIES.len(), + SYMBOLS.len() + ); + + // Static model dir: override with CODEGRAPH_STATIC_MODEL (e.g. a distilled + // jina-code-static-256), else the potion-base-8M floor. + let static_dir = std::env::var("CODEGRAPH_STATIC_MODEL") + .map(std::path::PathBuf::from) + .unwrap_or_else(|_| std::path::PathBuf::from(&home).join(".codegraph/static_models/potion-base-8M")); + report("static", VectorEngine::with_static_model(&static_dir)); + + let cache = std::path::PathBuf::from(&home).join(".codegraph/fastembed_cache"); + report("onnx", VectorEngine::with_model(cache, CodeGraphEmbeddingModel::BgeSmall)); +} diff --git a/crates/codegraph-memory/examples/embed_throughput.rs b/crates/codegraph-memory/examples/embed_throughput.rs new file mode 100644 index 0000000..b4e5133 --- /dev/null +++ b/crates/codegraph-memory/examples/embed_throughput.rs @@ -0,0 +1,97 @@ +// Copyright 2025-2026 Andrey Vasilevsky +// SPDX-License-Identifier: Apache-2.0 + +//! Embedding throughput: static (model2vec) vs ONNX transformer. +//! +//! The core premise of the static-embeddings work — indexing large codebases is +//! bottlenecked on the per-symbol transformer forward pass; a static lookup +//! table removes it. This prints texts/sec for each backend and the speedup. +//! +//! Run (release is required for a fair number — the static path is pure Rust): +//! cargo run --release -p codegraph-memory --example embed_throughput +//! +//! Needs a static model dir at `~/.codegraph/static_models/potion-base-8M` +//! (config.json + tokenizer.json + model.safetensors). BGE-small loads from the +//! existing fastembed cache (or downloads on first run). + +use codegraph_memory::embedding::{CodeGraphEmbeddingModel, VectorEngine}; +use std::time::Instant; + +/// Synthetic but realistic "name: signature — doc" symbol texts, all unique +/// (the `(variant N)` suffix) so the engine's cache never short-circuits. +fn sample_texts(n: usize) -> Vec { + let templates = [ + "authenticate_user: fn authenticate_user(token: &Token) -> Result — Verify the bearer token and return the user", + "open_database_connection: fn open_database_connection(url: &str) -> Pool — Open a pooled postgres connection", + "parse_config_file: fn parse_config_file(path: &Path) -> Config — Read and deserialize the TOML configuration", + "compute_cosine_similarity: fn compute_cosine_similarity(a: &[f32], b: &[f32]) -> f32 — Dot product over norms", + "RetryPolicy: struct RetryPolicy { max_attempts: u32, backoff: Duration } — Exponential-backoff retry config", + "serialize_to_json: fn serialize_to_json(value: &T) -> String — Encode a value as JSON text", + "spawn_worker_thread: fn spawn_worker_thread(queue: Arc) -> JoinHandle — Start a background consumer", + "validate_email_address: fn validate_email_address(input: &str) -> bool — Lenient RFC-5322 email check", + ]; + (0..n) + .map(|i| format!("{} (variant {i})", templates[i % templates.len()])) + .collect() +} + +/// Returns texts/sec, or None if the engine couldn't be built. +fn measure(label: &str, engine: Result, texts: &[String]) -> Option { + match engine { + Ok(engine) => { + let refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect(); + let t0 = Instant::now(); + let vecs = engine.embed_batch(&refs).expect("embed_batch"); + let secs = t0.elapsed().as_secs_f64(); + let sps = texts.len() as f64 / secs; + let dim = vecs.first().map(|v| v.len()).unwrap_or(0); + println!( + "{label:<8} {:<30} {dim:>4}d {sps:>10.0} texts/sec", + engine.model_name() + ); + Some(sps) + } + Err(e) => { + println!("{label:<8} skipped ({e})"); + None + } + } +} + +fn main() { + // Real symbols from CODEGRAPH_THROUGHPUT_CORPUS ([{id,signature}]), else synthetic. + let texts: Vec = match std::env::var("CODEGRAPH_THROUGHPUT_CORPUS") { + Ok(path) => { + #[derive(serde::Deserialize)] + struct S { + id: String, + signature: String, + } + let v: Vec = + serde_json::from_slice(&std::fs::read(&path).expect("read corpus")).expect("json"); + v.iter().map(|s| format!("{}: {}", s.id, s.signature)).collect() + } + Err(_) => sample_texts(512), + }; + let n = texts.len(); + let home = std::env::var("HOME").unwrap_or_default(); + println!("Embedding {n} unique symbol texts (cold, no cache hits).\n"); + + // Static model dir: override with CODEGRAPH_STATIC_MODEL (e.g. a distilled + // jina-code-static-256), else the potion-base-8M floor. + let static_dir = std::env::var("CODEGRAPH_STATIC_MODEL") + .map(std::path::PathBuf::from) + .unwrap_or_else(|_| std::path::PathBuf::from(&home).join(".codegraph/static_models/potion-base-8M")); + let static_sps = measure("static", VectorEngine::with_static_model(&static_dir), &texts); + + let cache = std::path::PathBuf::from(&home).join(".codegraph/fastembed_cache"); + let onnx_sps = measure( + "onnx", + VectorEngine::with_model(cache, CodeGraphEmbeddingModel::BgeSmall), + &texts, + ); + + if let (Some(s), Some(o)) = (static_sps, onnx_sps) { + println!("\nstatic is {:.1}x faster than ONNX BGE ({s:.0} vs {o:.0} texts/sec)", s / o); + } +} diff --git a/crates/codegraph-memory/src/embedding/engine.rs b/crates/codegraph-memory/src/embedding/engine.rs index 58e2dd5..40d4fce 100644 --- a/crates/codegraph-memory/src/embedding/engine.rs +++ b/crates/codegraph-memory/src/embedding/engine.rs @@ -6,6 +6,7 @@ //! High-level API for generating and caching embeddings. use super::fastembed_embed::{CodeGraphEmbeddingModel, FastembedEmbedding}; +use super::{Embedder, EmbeddingBackend}; use crate::error::Result; use dashmap::DashMap; use std::path::{Path, PathBuf}; @@ -15,7 +16,7 @@ use std::sync::Arc; /// /// Wraps fastembed with a configurable model and DashMap cache for efficient repeated lookups. pub struct VectorEngine { - model: Arc, + model: Arc, cache: DashMap>, dimension: usize, } @@ -44,6 +45,33 @@ impl VectorEngine { }) } + /// Create a VectorEngine backed by a static (lookup-table) model loaded from + /// a model2vec-format directory (`config.json` + `tokenizer.json` + + /// `model.safetensors`). No ONNX — the fast indexing path. + pub fn with_static_model(model_dir: &Path) -> Result { + let model = super::static_embed::StaticEmbedding::from_pretrained(model_dir)?; + let dimension = model.dimension(); + log::info!( + "VectorEngine ready ({}, {}d, static)", + model.model_name(), + dimension + ); + Ok(Self { + model: Arc::new(model), + cache: DashMap::new(), + dimension, + }) + } + + /// Build a VectorEngine from an `EmbeddingBackend` selection — the ONNX + /// fastembed path or the static model2vec path. + pub fn from_backend(cache_dir: PathBuf, backend: &EmbeddingBackend) -> Result { + match backend { + EmbeddingBackend::Fastembed(model) => Self::with_model(cache_dir, *model), + EmbeddingBackend::Static(dir) => Self::with_static_model(dir), + } + } + /// Generate embedding with caching pub fn embed(&self, text: &str) -> Result> { // Check cache first @@ -111,8 +139,8 @@ impl VectorEngine { } /// Get model display name (e.g. "Jina Code V2 (768d)") - pub fn model_name(&self) -> &'static str { - self.model.model_type().display_name() + pub fn model_name(&self) -> &str { + self.model.model_name() } /// Get cache size diff --git a/crates/codegraph-memory/src/embedding/fastembed_embed.rs b/crates/codegraph-memory/src/embedding/fastembed_embed.rs index 344c3e1..5db838e 100644 --- a/crates/codegraph-memory/src/embedding/fastembed_embed.rs +++ b/crates/codegraph-memory/src/embedding/fastembed_embed.rs @@ -234,6 +234,22 @@ impl FastembedEmbedding { } } +impl super::Embedder for FastembedEmbedding { + fn embed(&self, text: &str) -> Result> { + // Explicit path → the inherent method (avoids any trait-method ambiguity). + FastembedEmbedding::embed(self, text) + } + fn embed_batch(&self, texts: &[&str]) -> Result>> { + FastembedEmbedding::embed_batch(self, texts) + } + fn dimension(&self) -> usize { + FastembedEmbedding::dimension(self) + } + fn model_name(&self) -> &str { + self.model_type().display_name() + } +} + /// Bundle of bytes needed to construct a `UserDefinedEmbeddingModel`. struct UserDefinedModelBundle { onnx_bytes: Vec, diff --git a/crates/codegraph-memory/src/embedding/mod.rs b/crates/codegraph-memory/src/embedding/mod.rs index 3331bc4..5d3cea7 100644 --- a/crates/codegraph-memory/src/embedding/mod.rs +++ b/crates/codegraph-memory/src/embedding/mod.rs @@ -7,6 +7,91 @@ mod engine; mod fastembed_embed; +mod static_embed; pub use engine::VectorEngine; pub use fastembed_embed::CodeGraphEmbeddingModel; + +use crate::error::Result; +use std::path::PathBuf; + +/// Pluggable embedding backend. The transformer path (fastembed/ONNX) and a +/// future static (lookup-table) path both implement this; `VectorEngine` holds a +/// `dyn Embedder` so the backend is swappable without touching callers. +pub(crate) trait Embedder: Send + Sync { + fn embed(&self, text: &str) -> Result>; + fn embed_batch(&self, texts: &[&str]) -> Result>>; + fn dimension(&self) -> usize; + fn model_name(&self) -> &str; +} + +/// Selects which backend `VectorEngine` builds: the ONNX transformer path +/// (fastembed) or the static (model2vec) lookup path. +#[derive(Debug, Clone)] +pub enum EmbeddingBackend { + /// ONNX transformer via fastembed (bge-small / jina-code-v2 / granite-97m). + Fastembed(CodeGraphEmbeddingModel), + /// Static model2vec model from a local directory (~100x faster indexing, no + /// ONNX). The directory holds `config.json` + `tokenizer.json` + + /// `model.safetensors`. + Static(PathBuf), +} + +impl Default for EmbeddingBackend { + fn default() -> Self { + Self::Fastembed(CodeGraphEmbeddingModel::default()) + } +} + +impl EmbeddingBackend { + /// Parse a `--embedding-model` value. `static` selects the model2vec path, + /// resolving the model directory from `CODEGRAPH_STATIC_MODEL` or the default + /// `~/.codegraph/static_models/jina-code-static-256`. Unknown values fall + /// back to bge-small. + pub fn parse(s: &str) -> Self { + match s { + "static" | "static-code" | "model2vec" => Self::Static(default_static_model_dir()), + "jina-code-v2" => Self::Fastembed(CodeGraphEmbeddingModel::JinaCodeV2), + "granite-97m" | "granite" | "granite-97m-multilingual-r2" => { + Self::Fastembed(CodeGraphEmbeddingModel::Granite97mMultilingualR2) + } + _ => Self::Fastembed(CodeGraphEmbeddingModel::BgeSmall), + } + } + + /// Human-readable name for logging. + pub fn display_name(&self) -> String { + match self { + Self::Fastembed(m) => m.display_name().to_string(), + Self::Static(dir) => format!( + "static:{} (model2vec)", + dir.file_name().and_then(|s| s.to_str()).unwrap_or("model") + ), + } + } + + /// Short, stable tag for telemetry / adoption grouping. + pub fn telemetry_id(&self) -> &'static str { + match self { + Self::Static(_) => "static", + Self::Fastembed(CodeGraphEmbeddingModel::BgeSmall) => "bge-small", + Self::Fastembed(CodeGraphEmbeddingModel::JinaCodeV2) => "jina-code-v2", + Self::Fastembed(CodeGraphEmbeddingModel::Granite97mMultilingualR2) => "granite-97m", + } + } +} + +/// Resolve the static-model directory: `CODEGRAPH_STATIC_MODEL` env override, +/// else `~/.codegraph/static_models/jina-code-static-256`. +fn default_static_model_dir() -> PathBuf { + if let Ok(p) = std::env::var("CODEGRAPH_STATIC_MODEL") { + return PathBuf::from(p); + } + let home = std::env::var("HOME") + .or_else(|_| std::env::var("USERPROFILE")) + .unwrap_or_else(|_| ".".to_string()); + PathBuf::from(home) + .join(".codegraph") + .join("static_models") + .join("jina-code-static-256") +} diff --git a/crates/codegraph-memory/src/embedding/static_embed.rs b/crates/codegraph-memory/src/embedding/static_embed.rs new file mode 100644 index 0000000..c73b495 --- /dev/null +++ b/crates/codegraph-memory/src/embedding/static_embed.rs @@ -0,0 +1,295 @@ +// Copyright 2025-2026 Andrey Vasilevsky +// SPDX-License-Identifier: Apache-2.0 + +//! Static (lookup-table) embeddings — the fast indexing path. +//! +//! A static embedder runs no neural network at inference. It loads a +//! `token -> vector` matrix (distilled offline from a transformer teacher) plus +//! that teacher's tokenizer, and embeds text as: tokenize -> gather the token +//! rows -> (optional per-token weight) -> mean-pool -> L2-normalize. ~100x +//! faster than an ONNX transformer forward pass, trading away contextualization. +//! +//! Loads the model2vec on-disk format (`config.json` + `tokenizer.json` + +//! `model.safetensors`), so any model2vec / distilled static model — including a +//! Jina-Code distillation — drops in unchanged. Matches model2vec's encode +//! exactly (`model.py`): `emb[ids] * weights[ids]`, then `.mean(axis=0)`, then +//! `/ norm`. The matrix may be F32 (older potions bake the weighting in) or F16 +//! with a separate F64 `weights` vector (model2vec 0.7's SIF weighting). + +use super::Embedder; +use crate::error::{MemoryError, Result}; +use safetensors::tensor::TensorView; +use safetensors::{Dtype, SafeTensors}; +use std::path::Path; +use tokenizers::Tokenizer; + +/// The subset of model2vec's `config.json` we read (unknown fields ignored). +#[derive(serde::Deserialize)] +struct StaticConfig { + #[serde(default = "default_true")] + normalize: bool, + #[serde(default)] + hidden_dim: Option, +} +fn default_true() -> bool { + true +} + +/// A static, lookup-table embedder (no ONNX / no transformer at inference). +pub(crate) struct StaticEmbedding { + tokenizer: Tokenizer, + /// Row-major `vocab * dim` token vectors; token `i` is `matrix[i*dim..(i+1)*dim]`. + matrix: Vec, + /// Optional per-token weight (model2vec SIF). `None` => unweighted mean. + weights: Option>, + vocab: usize, + dim: usize, + normalize: bool, + name: String, +} + +impl StaticEmbedding { + /// Load a model2vec-format directory: `config.json`, `tokenizer.json`, and + /// `model.safetensors` (an `embeddings` [vocab, dim] tensor, F32 or F16, and + /// an optional `weights` [vocab] tensor). + pub(crate) fn from_pretrained(dir: &Path) -> Result { + let cfg: StaticConfig = serde_json::from_slice( + &std::fs::read(dir.join("config.json")) + .map_err(|e| MemoryError::model(format!("read config.json: {e}")))?, + )?; + + let tokenizer = Tokenizer::from_file(dir.join("tokenizer.json")) + .map_err(|e| MemoryError::model(format!("load tokenizer.json: {e}")))?; + + let bytes = std::fs::read(dir.join("model.safetensors")) + .map_err(|e| MemoryError::model(format!("read model.safetensors: {e}")))?; + let st = SafeTensors::deserialize(&bytes) + .map_err(|e| MemoryError::model(format!("parse safetensors: {e}")))?; + + let emb = st + .tensor("embeddings") + .map_err(|e| MemoryError::model(format!("no 'embeddings' tensor: {e}")))?; + let shape = emb.shape(); + if shape.len() != 2 { + return Err(MemoryError::model(format!( + "embeddings must be 2-D [vocab, dim], got {shape:?}" + ))); + } + let (vocab, dim) = (shape[0], shape[1]); + let matrix = tensor_to_f32(&emb)?; + if matrix.len() != vocab * dim { + return Err(MemoryError::model(format!( + "embeddings length {} != vocab*dim {}", + matrix.len(), + vocab * dim + ))); + } + + // Optional per-token SIF weights (model2vec >= 0.4 stores them separately + // instead of baking them into `embeddings`). + let weights = match st.tensor("weights") { + Ok(w) => { + let v = tensor_to_f32(&w)?; + if v.len() != vocab { + return Err(MemoryError::model(format!( + "weights length {} != vocab {vocab}", + v.len() + ))); + } + Some(v) + } + Err(_) => None, + }; + + let name = format!( + "{} ({}d, static)", + dir.file_name().and_then(|s| s.to_str()).unwrap_or("static"), + cfg.hidden_dim.unwrap_or(dim) + ); + + Ok(Self { + tokenizer, + matrix, + weights, + vocab, + dim, + normalize: cfg.normalize, + name, + }) + } + + fn embed_text(&self, text: &str) -> Result> { + // add_special_tokens = false: a bag-of-tokens static model gains nothing + // from [CLS]/[SEP]; they would only dilute the mean. + let enc = self + .tokenizer + .encode(text, false) + .map_err(|e| MemoryError::embedding(format!("tokenize: {e}")))?; + Ok(weighted_mean_l2( + &self.matrix, + self.weights.as_deref(), + enc.get_ids(), + self.dim, + self.vocab, + self.normalize, + )) + } +} + +/// Decode a safetensors float tensor (F32 / F16 / F64) into `Vec`. +fn tensor_to_f32(t: &TensorView) -> Result> { + let raw = t.data(); + let v = match t.dtype() { + Dtype::F32 => raw + .chunks_exact(4) + .map(|b| f32::from_le_bytes(b.try_into().unwrap())) + .collect(), + Dtype::F16 => raw + .chunks_exact(2) + .map(|b| half::f16::from_le_bytes(b.try_into().unwrap()).to_f32()) + .collect(), + Dtype::F64 => raw + .chunks_exact(8) + .map(|b| f64::from_le_bytes(b.try_into().unwrap()) as f32) + .collect(), + other => { + return Err(MemoryError::model(format!( + "unsupported tensor dtype {other:?} (want F32/F16/F64)" + ))) + } + }; + Ok(v) +} + +/// model2vec's pooling: each token row times its weight, mean over tokens, then +/// optional L2-normalize. Out-of-vocab ids are skipped. The whole hot path. +fn weighted_mean_l2( + matrix: &[f32], + weights: Option<&[f32]>, + ids: &[u32], + dim: usize, + vocab: usize, + normalize: bool, +) -> Vec { + let mut acc = vec![0f32; dim]; + let mut n = 0usize; + for &id in ids { + let i = id as usize; + if i >= vocab { + continue; + } + let w = weights.map_or(1.0, |ws| ws[i]); + let row = &matrix[i * dim..(i + 1) * dim]; + for (a, r) in acc.iter_mut().zip(row) { + *a += w * *r; + } + n += 1; + } + if n > 0 { + let inv = 1.0 / n as f32; + for a in acc.iter_mut() { + *a *= inv; + } + } + if normalize { + let norm: f32 = acc.iter().map(|x| x * x).sum::().sqrt(); + if norm > 0.0 { + let inv = 1.0 / norm; + for a in acc.iter_mut() { + *a *= inv; + } + } + } + acc +} + +impl Embedder for StaticEmbedding { + fn embed(&self, text: &str) -> Result> { + self.embed_text(text) + } + fn embed_batch(&self, texts: &[&str]) -> Result>> { + texts.iter().map(|t| self.embed_text(t)).collect() + } + fn dimension(&self) -> usize { + self.dim + } + fn model_name(&self) -> &str { + &self.name + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn weighted_mean_l2_math() { + // vocab 2, dim 2: row0 = [3,0], row1 = [0,4]. + let matrix = vec![3.0, 0.0, 0.0, 4.0]; + // unweighted: mean of the two rows + assert_eq!(weighted_mean_l2(&matrix, None, &[0, 1], 2, 2, false), vec![1.5, 2.0]); + // out-of-vocab id (99) skipped -> only row0 counts + assert_eq!(weighted_mean_l2(&matrix, None, &[0, 99], 2, 2, false), vec![3.0, 0.0]); + // all-OOV -> zero vector, no NaN even with normalize + assert_eq!(weighted_mean_l2(&matrix, None, &[7], 2, 2, true), vec![0.0, 0.0]); + // normalize -> unit length + let v = weighted_mean_l2(&matrix, None, &[0, 1], 2, 2, true); + let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + assert!((norm - 1.0).abs() < 1e-6, "norm={norm}"); + // weighted: row0*2 + row1*0.5, mean over 2 -> [3.0, 1.0] + let w = vec![2.0, 0.5]; + assert_eq!( + weighted_mean_l2(&matrix, Some(&w), &[0, 1], 2, 2, false), + vec![3.0, 1.0] + ); + } + + fn static_dir(name: &str) -> std::path::PathBuf { + std::path::PathBuf::from(std::env::var("HOME").unwrap_or_default()) + .join(".codegraph/static_models") + .join(name) + } + + fn assert_semantically_sane(m: &StaticEmbedding, expected_dim: usize) { + assert_eq!(m.dimension(), expected_dim); + let v = m.embed_text("database connection pool").unwrap(); + assert_eq!(v.len(), expected_dim); + let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + assert!((norm - 1.0).abs() < 1e-3, "should be L2-normalized, norm={norm}"); + + let cos = |a: &[f32], b: &[f32]| a.iter().zip(b).map(|(x, y)| x * y).sum::(); + let a = m.embed_text("open a database connection").unwrap(); + let b = m.embed_text("connect to the postgres database").unwrap(); + let u = m.embed_text("the cat sat on the warm windowsill").unwrap(); + assert!( + cos(&a, &b) > cos(&a, &u), + "related {:.3} should exceed unrelated {:.3}", + cos(&a, &b), + cos(&a, &u) + ); + } + + #[test] + fn loads_potion_f32_no_weights() { + let dir = static_dir("potion-base-8M"); + if !dir.join("model.safetensors").exists() { + eprintln!("skip: potion-base-8M not present"); + return; + } + let m = StaticEmbedding::from_pretrained(&dir).expect("load potion-base-8M"); + assert!(m.weights.is_none(), "potion-base-8M bakes weights into embeddings"); + assert_semantically_sane(&m, 256); + } + + #[test] + fn loads_jina_code_static_f16_weighted() { + let dir = static_dir("jina-code-static-256"); + if !dir.join("model.safetensors").exists() { + eprintln!("skip: jina-code-static-256 not present"); + return; + } + let m = StaticEmbedding::from_pretrained(&dir).expect("load jina-code-static-256"); + assert!(m.weights.is_some(), "model2vec 0.7 stores SIF weights separately"); + assert_semantically_sane(&m, 256); + } +} diff --git a/crates/codegraph-memory/src/lib.rs b/crates/codegraph-memory/src/lib.rs index c8c31d0..c50d276 100644 --- a/crates/codegraph-memory/src/lib.rs +++ b/crates/codegraph-memory/src/lib.rs @@ -43,7 +43,7 @@ pub mod storage; pub mod temporal; // Re-exports for convenience -pub use embedding::{CodeGraphEmbeddingModel, VectorEngine}; +pub use embedding::{CodeGraphEmbeddingModel, EmbeddingBackend, VectorEngine}; pub use error::MemoryError; pub use node::{ CodeLink, IssueSeverity, LinkedNodeType, MemoryId, MemoryKind, MemoryNode, MemoryNodeBuilder, diff --git a/crates/codegraph-server/Cargo.toml b/crates/codegraph-server/Cargo.toml index 8154336..3028865 100644 --- a/crates/codegraph-server/Cargo.toml +++ b/crates/codegraph-server/Cargo.toml @@ -55,13 +55,16 @@ codegraph-ruby.workspace = true codegraph-swift.workspace = true codegraph-tcl.workspace = true codegraph-verilog.workspace = true -codegraph-cobol.workspace = true -codegraph-fortran.workspace = true -codegraph-dart.workspace = true -codegraph-zig.workspace = true +# Heavy, zero-telemetry-usage grammars — gated behind `extra-languages` +# (default-off) to keep the community binary lean. COBOL's parser.c alone is +# 30.7 MB; none of these appear in any indexed workspace (PostHog, 90d). +codegraph-cobol = { workspace = true, optional = true } +codegraph-fortran = { workspace = true, optional = true } +codegraph-dart = { workspace = true, optional = true } +codegraph-zig = { workspace = true, optional = true } codegraph-lua.workspace = true codegraph-groovy.workspace = true -codegraph-r.workspace = true +codegraph-r = { workspace = true, optional = true } codegraph-scala.workspace = true codegraph-bash.workspace = true codegraph-elixir.workspace = true @@ -69,7 +72,7 @@ codegraph-haskell.workspace = true codegraph-hcl.workspace = true codegraph-julia.workspace = true codegraph-ocaml.workspace = true -codegraph-perl.workspace = true +codegraph-perl = { workspace = true, optional = true } codegraph-toml.workspace = true codegraph-yaml.workspace = true codegraph-clojure.workspace = true @@ -86,6 +89,18 @@ codegraph-memory.workspace = true # CLI clap.workspace = true +[features] +# Off by default → the community binary omits these heavy, unused grammars. +# Pro / full builds opt in with `--features extra-languages`. +extra-languages = [ + "dep:codegraph-cobol", + "dep:codegraph-fortran", + "dep:codegraph-dart", + "dep:codegraph-zig", + "dep:codegraph-r", + "dep:codegraph-perl", +] + [dev-dependencies] tempfile = "3" tokio-test = "0.4" diff --git a/crates/codegraph-server/src/ai_query/engine.rs b/crates/codegraph-server/src/ai_query/engine.rs index 4740623..57b3ff1 100644 --- a/crates/codegraph-server/src/ai_query/engine.rs +++ b/crates/codegraph-server/src/ai_query/engine.rs @@ -47,6 +47,9 @@ pub struct QueryEngine { symbol_texts: Arc>>, /// Embed full function body (true) or just name+signature (false, default) full_body_embedding: std::sync::atomic::AtomicBool, + /// Prepend split-identifier words to the embed text (helps static + /// embedders; off by default to leave the transformer path unchanged). + split_identifiers: std::sync::atomic::AtomicBool, } /// Max characters of function body for full-body embedding. @@ -73,6 +76,17 @@ fn available_memory_mb() -> u64 { sys.available_memory() / (1024 * 1024) } +/// Split an identifier into camelCase/snake_case words (deduped, lowercased), +/// reusing the BM25 tokenizer: `getUserById` -> "get user by id". +fn split_identifier_words(name: &str) -> String { + let mut seen = std::collections::HashSet::new(); + super::text_index::tokenize(name) + .into_iter() + .filter(|t| seen.insert(t.clone())) + .collect::>() + .join(" ") +} + impl QueryEngine { /// Create a new query engine with the given graph. pub fn new(graph: Arc>) -> Self { @@ -86,6 +100,7 @@ impl QueryEngine { symbol_vectors: Arc::new(RwLock::new(HashMap::new())), symbol_texts: Arc::new(RwLock::new(HashMap::new())), full_body_embedding: std::sync::atomic::AtomicBool::new(true), + split_identifiers: std::sync::atomic::AtomicBool::new(false), } } @@ -95,6 +110,13 @@ impl QueryEngine { .store(enabled, std::sync::atomic::Ordering::Relaxed); } + /// Enable or disable prepending split-identifier words to the embed text. + /// Helps static (lookup-table) embedders; off by default. + pub fn set_split_identifiers(&self, enabled: bool) { + self.split_identifiers + .store(enabled, std::sync::atomic::Ordering::Relaxed); + } + /// Build the embedding text for a symbol node. /// In signature mode: "name: signature — docstring" /// In full-body mode: "name: signature\n" @@ -103,6 +125,7 @@ impl QueryEngine { node_id: NodeId, name: &str, full_body: bool, + split_identifiers: bool, graph: &CodeGraph, ) -> String { let signature = node.properties.get_string("signature").unwrap_or(""); @@ -119,6 +142,21 @@ impl QueryEngine { name.to_string() }; + // Prepend the camelCase/snake_case-split form of the name when enabled. + // Static (lookup-table) embedders can't subword-recover `authenticateUser` + // from one rare token; the split words ("authenticate user") are their + // strongest signal, front-loaded so they survive truncation. + let base = if split_identifiers { + let words = split_identifier_words(name); + if words.is_empty() || words == name.to_lowercase() { + base + } else { + format!("{words} {base}") + } + } else { + base + }; + if !full_body { return base; } @@ -293,6 +331,8 @@ impl QueryEngine { name, self.full_body_embedding .load(std::sync::atomic::Ordering::Relaxed), + self.split_identifiers + .load(std::sync::atomic::Ordering::Relaxed), &graph, ); @@ -444,6 +484,8 @@ impl QueryEngine { name, self.full_body_embedding .load(std::sync::atomic::Ordering::Relaxed), + self.split_identifiers + .load(std::sync::atomic::Ordering::Relaxed), &graph, ); node_ids.push(node_id); @@ -806,6 +848,8 @@ impl QueryEngine { name, self.full_body_embedding .load(std::sync::atomic::Ordering::Relaxed), + self.split_identifiers + .load(std::sync::atomic::Ordering::Relaxed), &graph, ); @@ -3917,4 +3961,37 @@ mod tests { // Should find main, test, and public API assert!(results.len() >= 3); } + + #[test] + fn split_identifier_words_handles_camel_and_snake() { + assert_eq!(split_identifier_words("authenticate_user"), "authenticate user"); + assert_eq!(split_identifier_words("getUserById"), "get user by id"); + // The existing tokenizer keeps acronym+word runs joined (HTML|Parser is + // NOT split) and drops 1-char tokens — known limitations worth revisiting + // for static-embedding quality (Phase 3 input refinement). + assert_eq!(split_identifier_words("HTMLParser"), "htmlparser"); + // A single lowercase word splits to itself. + assert_eq!(split_identifier_words("foo"), "foo"); + } + + #[test] + fn build_embed_text_prepends_split_name_only_when_enabled() { + let graph = CodeGraph::in_memory().unwrap(); + let node = codegraph::Node::new( + 0, + codegraph::NodeType::Function, + PropertyMap::new().with("signature", "fn getUserById(id: u64) -> User"), + ); + + // Enabled: split name words are front-loaded for the static embedder. + let with_split = + QueryEngine::build_embed_text(&node, 0, "getUserById", false, true, &graph); + assert!(with_split.starts_with("get user by id"), "got: {with_split}"); + + // Disabled (default): original transformer-path text is unchanged. + let without = + QueryEngine::build_embed_text(&node, 0, "getUserById", false, false, &graph); + assert!(without.starts_with("getUserById"), "got: {without}"); + assert!(!without.contains("get user by id")); + } } diff --git a/crates/codegraph-server/src/ai_query/text_index.rs b/crates/codegraph-server/src/ai_query/text_index.rs index c6b8503..bd4b16a 100644 --- a/crates/codegraph-server/src/ai_query/text_index.rs +++ b/crates/codegraph-server/src/ai_query/text_index.rs @@ -330,7 +330,7 @@ impl Default for TextIndexBuilder { /// Tokenize a string into lowercase tokens. /// Splits on non-alphanumeric characters and handles camelCase/snake_case. /// Handles acronyms like "VALIDATE" or "HTMLParser" correctly. -fn tokenize(text: &str) -> Vec { +pub(crate) fn tokenize(text: &str) -> Vec { let mut tokens = Vec::new(); let mut current = String::new(); let mut prev_was_upper = false; diff --git a/crates/codegraph-server/src/backend.rs b/crates/codegraph-server/src/backend.rs index dfdd211..14d4985 100644 --- a/crates/codegraph-server/src/backend.rs +++ b/crates/codegraph-server/src/backend.rs @@ -919,24 +919,9 @@ impl LanguageServer for CodeGraphBackend { let embedding_model = raw_model .and_then(|v| v.as_str()) - .and_then(|s| { + .map(|s| { tracing::info!("[LSP::initialize] Parsing embedding model string: {:?}", s); - match s { - "bge-small" => Some(codegraph_memory::CodeGraphEmbeddingModel::BgeSmall), - "jina-code-v2" => { - Some(codegraph_memory::CodeGraphEmbeddingModel::JinaCodeV2) - } - "granite-97m" | "granite" | "granite-97m-multilingual-r2" => Some( - codegraph_memory::CodeGraphEmbeddingModel::Granite97mMultilingualR2, - ), - _ => { - tracing::warn!( - "[LSP::initialize] Unknown embedding model: {:?}, using default", - s - ); - None - } - } + codegraph_memory::EmbeddingBackend::parse(s) }) .unwrap_or_default(); diff --git a/crates/codegraph-server/src/main.rs b/crates/codegraph-server/src/main.rs index 569e546..6ac0250 100644 --- a/crates/codegraph-server/src/main.rs +++ b/crates/codegraph-server/src/main.rs @@ -368,13 +368,7 @@ async fn run() { } else { args.workspace.clone() }; - let embedding_model = match args.embedding_model.as_str() { - "jina-code-v2" => codegraph_memory::CodeGraphEmbeddingModel::JinaCodeV2, - "granite-97m" | "granite" | "granite-97m-multilingual-r2" => { - codegraph_memory::CodeGraphEmbeddingModel::Granite97mMultilingualR2 - } - _ => codegraph_memory::CodeGraphEmbeddingModel::BgeSmall, - }; + let embedding_model = codegraph_memory::EmbeddingBackend::parse(&args.embedding_model); let tool_args: serde_json::Value = serde_json::from_str(&args.tool_args).unwrap_or_else(|e| { eprintln!("Invalid --tool-args JSON: {e}"); @@ -413,13 +407,7 @@ async fn run() { // Resident socket engine (Model B): one shared model, many thin clients. if args.serve { - let embedding_model = match args.embedding_model.as_str() { - "jina-code-v2" => codegraph_memory::CodeGraphEmbeddingModel::JinaCodeV2, - "granite-97m" | "granite" | "granite-97m-multilingual-r2" => { - codegraph_memory::CodeGraphEmbeddingModel::Granite97mMultilingualR2 - } - _ => codegraph_memory::CodeGraphEmbeddingModel::BgeSmall, - }; + let embedding_model = codegraph_memory::EmbeddingBackend::parse(&args.embedding_model); let cfg = codegraph_server::mcp::engine::EngineConfig { socket_path: args.socket.clone().unwrap_or_else(default_socket_path), embedding_model, @@ -443,13 +431,7 @@ async fn run() { args.workspace }; - let embedding_model = match args.embedding_model.as_str() { - "jina-code-v2" => codegraph_memory::CodeGraphEmbeddingModel::JinaCodeV2, - "granite-97m" | "granite" | "granite-97m-multilingual-r2" => { - codegraph_memory::CodeGraphEmbeddingModel::Granite97mMultilingualR2 - } - _ => codegraph_memory::CodeGraphEmbeddingModel::BgeSmall, - }; + let embedding_model = codegraph_memory::EmbeddingBackend::parse(&args.embedding_model); tracing::info!("Starting CodeGraph MCP server"); tracing::info!("Workspaces: {:?}", workspaces); diff --git a/crates/codegraph-server/src/mcp/engine.rs b/crates/codegraph-server/src/mcp/engine.rs index 1c15d9b..a94ffab 100644 --- a/crates/codegraph-server/src/mcp/engine.rs +++ b/crates/codegraph-server/src/mcp/engine.rs @@ -18,14 +18,14 @@ use std::path::PathBuf; -use codegraph_memory::CodeGraphEmbeddingModel; +use codegraph_memory::EmbeddingBackend; use super::server::McpServer; /// Configuration for a running engine. pub struct EngineConfig { pub socket_path: PathBuf, - pub embedding_model: CodeGraphEmbeddingModel, + pub embedding_model: EmbeddingBackend, pub exclude_dirs: Vec, pub max_files: usize, pub full_body_embedding: bool, @@ -105,7 +105,7 @@ mod imp { vec![ws.clone()], engine.cfg.exclude_dirs.clone(), engine.cfg.max_files, - engine.cfg.embedding_model, + engine.cfg.embedding_model.clone(), engine.cfg.full_body_embedding, ); if let Some(shared) = &engine.shared_engine { @@ -223,7 +223,7 @@ mod imp { tracing::warn!("Engine: <1.5 GB free — running graph-only (no shared model)"); None } else { - match VectorEngine::with_model(model_cache_dir(), cfg.embedding_model) { + match VectorEngine::from_backend(model_cache_dir(), &cfg.embedding_model) { Ok(e) => { tracing::info!("Engine: shared embedding model loaded"); Some(Arc::new(e)) diff --git a/crates/codegraph-server/src/mcp/server.rs b/crates/codegraph-server/src/mcp/server.rs index 4c34168..ce46b85 100644 --- a/crates/codegraph-server/src/mcp/server.rs +++ b/crates/codegraph-server/src/mcp/server.rs @@ -166,7 +166,7 @@ impl McpBackend { workspaces: Vec, exclude_dirs: Vec, max_files: usize, - embedding_model: codegraph_memory::CodeGraphEmbeddingModel, + embedding_model: codegraph_memory::EmbeddingBackend, full_body_embedding: bool, ) -> Self { let primary = workspaces.first().expect("At least one workspace required"); @@ -1228,7 +1228,7 @@ impl McpServer { workspaces: Vec, exclude_dirs: Vec, max_files: usize, - embedding_model: codegraph_memory::CodeGraphEmbeddingModel, + embedding_model: codegraph_memory::EmbeddingBackend, full_body_embedding: bool, ) -> Self { Self::with_pro_provider( @@ -1246,7 +1246,7 @@ impl McpServer { workspaces: Vec, exclude_dirs: Vec, max_files: usize, - embedding_model: codegraph_memory::CodeGraphEmbeddingModel, + embedding_model: codegraph_memory::EmbeddingBackend, full_body_embedding: bool, pro_provider: Arc, ) -> Self { @@ -1395,6 +1395,7 @@ impl McpServer { "os": std::env::consts::OS, "arch": std::env::consts::ARCH, "version": crate::metadata::VERSION, + "embeddingModel": self.backend.memory_manager.embedding_telemetry_id(), })); loop { diff --git a/crates/codegraph-server/src/memory.rs b/crates/codegraph-server/src/memory.rs index 65db599..10ba1fd 100644 --- a/crates/codegraph-server/src/memory.rs +++ b/crates/codegraph-server/src/memory.rs @@ -185,23 +185,20 @@ pub struct MemoryManager { extension_path: Option, /// Cached vector engine (holds model, not DB - safe to keep) engine: Arc>>>, - /// Embedding model selection - embedding_model: codegraph_memory::CodeGraphEmbeddingModel, + /// Embedding backend selection (fastembed model or static model2vec) + embedding_model: codegraph_memory::EmbeddingBackend, } impl MemoryManager { /// Create a new MemoryManager pub fn new(extension_path: Option) -> Self { - Self::with_model( - extension_path, - codegraph_memory::CodeGraphEmbeddingModel::default(), - ) + Self::with_model(extension_path, codegraph_memory::EmbeddingBackend::default()) } - /// Create a new MemoryManager with a specific embedding model + /// Create a new MemoryManager with a specific embedding backend pub fn with_model( extension_path: Option, - embedding_model: codegraph_memory::CodeGraphEmbeddingModel, + embedding_model: codegraph_memory::EmbeddingBackend, ) -> Self { Self { data_dir: Arc::new(RwLock::new(None)), @@ -211,6 +208,12 @@ impl MemoryManager { } } + /// Short telemetry tag for the configured embedding backend + /// (`static` / `bge-small` / `jina-code-v2` / `granite-97m`). + pub fn embedding_telemetry_id(&self) -> &'static str { + self.embedding_model.telemetry_id() + } + /// Initialize the memory manager with workspace path /// /// Resolves the global data directory at `~/.codegraph/projects//memory/`, @@ -302,7 +305,7 @@ impl MemoryManager { // Phase marker: a native crash during the ONNX model load never runs // the panic hook, so stamp the phase for the extension to read post-mortem. crate::crash_phase::mark("onnx_load"); - let engine = VectorEngine::with_model(cache_dir, self.embedding_model).map_err(|e| { + let engine = VectorEngine::from_backend(cache_dir, &self.embedding_model).map_err(|e| { tracing::error!( "[MemoryManager::initialize] VectorEngine initialization failed: {:?}", e diff --git a/crates/codegraph-server/src/parser_registry.rs b/crates/codegraph-server/src/parser_registry.rs index 4dd15ec..af4a0cc 100644 --- a/crates/codegraph-server/src/parser_registry.rs +++ b/crates/codegraph-server/src/parser_registry.rs @@ -7,15 +7,18 @@ use codegraph::CodeGraph; use codegraph_bash::BashParser; use codegraph_c::CParser; use codegraph_clojure::ClojureParser; +#[cfg(feature = "extra-languages")] use codegraph_cobol::CobolParser; use codegraph_cpp::CppParser; use codegraph_csharp::CSharpParser; use codegraph_css::CssParser; +#[cfg(feature = "extra-languages")] use codegraph_dart::DartParser; use codegraph_dockerfile::DockerfileParser; use codegraph_elixir::ElixirParser; use codegraph_elm::ElmParser; use codegraph_erlang::ErlangParser; +#[cfg(feature = "extra-languages")] use codegraph_fortran::FortranParser; use codegraph_go::GoParser; use codegraph_groovy::GroovyParser; @@ -28,9 +31,11 @@ use codegraph_lua::LuaParser; use codegraph_objc::ObjcParser; use codegraph_ocaml::OcamlParser; use codegraph_parser_api::{CodeParser, FileInfo, ParserConfig, ParserError, ParserMetrics}; +#[cfg(feature = "extra-languages")] use codegraph_perl::PerlParser; use codegraph_php::PhpParser; use codegraph_python::PythonParser; +#[cfg(feature = "extra-languages")] use codegraph_r::RParser; use codegraph_ruby::RubyParser; use codegraph_rust::RustParser; @@ -42,6 +47,7 @@ use codegraph_toml::TomlParser; use codegraph_typescript::TypeScriptParser; use codegraph_verilog::VerilogParser; use codegraph_yaml::YamlParser; +#[cfg(feature = "extra-languages")] use codegraph_zig::ZigParser; use std::path::Path; use std::sync::Arc; @@ -51,15 +57,18 @@ pub struct ParserRegistry { bash: Arc, c: Arc, clojure: Arc, + #[cfg(feature = "extra-languages")] cobol: Arc, cpp: Arc, css: Arc, csharp: Arc, + #[cfg(feature = "extra-languages")] dart: Arc, dockerfile: Arc, elixir: Arc, elm: Arc, erlang: Arc, + #[cfg(feature = "extra-languages")] fortran: Arc, go: Arc, groovy: Arc, @@ -71,9 +80,11 @@ pub struct ParserRegistry { lua: Arc, objc: Arc, ocaml: Arc, + #[cfg(feature = "extra-languages")] perl: Arc, php: Arc, python: Arc, + #[cfg(feature = "extra-languages")] r: Arc, ruby: Arc, rust: Arc, @@ -85,6 +96,7 @@ pub struct ParserRegistry { typescript: Arc, verilog: Arc, yaml: Arc, + #[cfg(feature = "extra-languages")] zig: Arc, } @@ -100,15 +112,18 @@ impl ParserRegistry { bash: Arc::new(BashParser::with_config(config.clone())), c: Arc::new(CParser::with_config(config.clone())), clojure: Arc::new(ClojureParser::with_config(config.clone())), + #[cfg(feature = "extra-languages")] cobol: Arc::new(CobolParser::with_config(config.clone())), cpp: Arc::new(CppParser::with_config(config.clone())), css: Arc::new(CssParser::with_config(config.clone())), csharp: Arc::new(CSharpParser::with_config(config.clone())), + #[cfg(feature = "extra-languages")] dart: Arc::new(DartParser::with_config(config.clone())), dockerfile: Arc::new(DockerfileParser::with_config(config.clone())), elixir: Arc::new(ElixirParser::with_config(config.clone())), elm: Arc::new(ElmParser::with_config(config.clone())), erlang: Arc::new(ErlangParser::with_config(config.clone())), + #[cfg(feature = "extra-languages")] fortran: Arc::new(FortranParser::with_config(config.clone())), go: Arc::new(GoParser::with_config(config.clone())), groovy: Arc::new(GroovyParser::with_config(config.clone())), @@ -120,9 +135,11 @@ impl ParserRegistry { lua: Arc::new(LuaParser::with_config(config.clone())), objc: Arc::new(ObjcParser::with_config(config.clone())), ocaml: Arc::new(OcamlParser::with_config(config.clone())), + #[cfg(feature = "extra-languages")] perl: Arc::new(PerlParser::with_config(config.clone())), php: Arc::new(PhpParser::with_config(config.clone())), python: Arc::new(PythonParser::with_config(config.clone())), + #[cfg(feature = "extra-languages")] r: Arc::new(RParser::with_config(config.clone())), ruby: Arc::new(RubyParser::with_config(config.clone())), rust: Arc::new(RustParser::with_config(config.clone())), @@ -134,6 +151,7 @@ impl ParserRegistry { typescript: Arc::new(TypeScriptParser::with_config(config.clone())), verilog: Arc::new(VerilogParser::with_config(config.clone())), yaml: Arc::new(YamlParser::with_config(config.clone())), + #[cfg(feature = "extra-languages")] zig: Arc::new(ZigParser::with_config(config)), } } @@ -144,15 +162,18 @@ impl ParserRegistry { "bash" | "shell" | "sh" => Some(self.bash.clone()), "c" => Some(self.c.clone()), "clojure" => Some(self.clojure.clone()), + #[cfg(feature = "extra-languages")] "cobol" => Some(self.cobol.clone()), "cpp" | "c++" => Some(self.cpp.clone()), "css" => Some(self.css.clone()), "csharp" | "c#" => Some(self.csharp.clone()), + #[cfg(feature = "extra-languages")] "dart" => Some(self.dart.clone()), "dockerfile" | "containerfile" => Some(self.dockerfile.clone()), "elixir" => Some(self.elixir.clone()), "elm" => Some(self.elm.clone()), "erlang" => Some(self.erlang.clone()), + #[cfg(feature = "extra-languages")] "fortran" => Some(self.fortran.clone()), "go" => Some(self.go.clone()), "groovy" => Some(self.groovy.clone()), @@ -164,9 +185,11 @@ impl ParserRegistry { "lua" => Some(self.lua.clone()), "objc" | "objective-c" | "objectivec" => Some(self.objc.clone()), "ocaml" => Some(self.ocaml.clone()), + #[cfg(feature = "extra-languages")] "perl" => Some(self.perl.clone()), "php" => Some(self.php.clone()), "python" => Some(self.python.clone()), + #[cfg(feature = "extra-languages")] "r" => Some(self.r.clone()), "ruby" => Some(self.ruby.clone()), "rust" => Some(self.rust.clone()), @@ -180,6 +203,7 @@ impl ParserRegistry { } "verilog" | "systemverilog" => Some(self.verilog.clone()), "yaml" => Some(self.yaml.clone()), + #[cfg(feature = "extra-languages")] "zig" => Some(self.zig.clone()), _ => None, } @@ -191,20 +215,18 @@ impl ParserRegistry { /// C++-specific extensions (`.hpp`, `.cc`, `.cxx`, `.hh`, `.hxx`) are /// only claimed by the C++ parser and resolve correctly. pub fn parser_for_path(&self, path: &Path) -> Option> { - let parsers: [Arc; 38] = [ + #[cfg_attr(not(feature = "extra-languages"), allow(unused_mut))] + let mut parsers: Vec> = vec![ self.bash.clone(), self.c.clone(), self.clojure.clone(), - self.cobol.clone(), self.cpp.clone(), self.css.clone(), self.csharp.clone(), - self.dart.clone(), self.dockerfile.clone(), self.elixir.clone(), self.elm.clone(), self.erlang.clone(), - self.fortran.clone(), self.go.clone(), self.groovy.clone(), self.haskell.clone(), @@ -215,10 +237,8 @@ impl ParserRegistry { self.lua.clone(), self.objc.clone(), self.ocaml.clone(), - self.perl.clone(), self.php.clone(), self.python.clone(), - self.r.clone(), self.ruby.clone(), self.rust.clone(), self.scala.clone(), @@ -229,8 +249,22 @@ impl ParserRegistry { self.typescript.clone(), self.verilog.clone(), self.yaml.clone(), - self.zig.clone(), ]; + // Gated grammars handle distinct extensions (.cob/.dart/.f90/.pl/.r/.zig), + // so appending them keeps resolution identical. + #[cfg(feature = "extra-languages")] + { + // Explicit element type coerces each Arc to the trait object. + let extra: [Arc; 6] = [ + self.cobol.clone(), + self.dart.clone(), + self.fortran.clone(), + self.perl.clone(), + self.r.clone(), + self.zig.clone(), + ]; + parsers.extend(extra); + } parsers.into_iter().find(|p| p.can_parse(path)) } @@ -241,15 +275,18 @@ impl ParserRegistry { extensions.extend(self.bash.file_extensions().iter().copied()); extensions.extend(self.c.file_extensions().iter().copied()); extensions.extend(self.clojure.file_extensions().iter().copied()); + #[cfg(feature = "extra-languages")] extensions.extend(self.cobol.file_extensions().iter().copied()); extensions.extend(self.cpp.file_extensions().iter().copied()); extensions.extend(self.css.file_extensions().iter().copied()); extensions.extend(self.csharp.file_extensions().iter().copied()); + #[cfg(feature = "extra-languages")] extensions.extend(self.dart.file_extensions().iter().copied()); extensions.extend(self.dockerfile.file_extensions().iter().copied()); extensions.extend(self.elixir.file_extensions().iter().copied()); extensions.extend(self.elm.file_extensions().iter().copied()); extensions.extend(self.erlang.file_extensions().iter().copied()); + #[cfg(feature = "extra-languages")] extensions.extend(self.fortran.file_extensions().iter().copied()); extensions.extend(self.go.file_extensions().iter().copied()); extensions.extend(self.groovy.file_extensions().iter().copied()); @@ -261,9 +298,11 @@ impl ParserRegistry { extensions.extend(self.lua.file_extensions().iter().copied()); extensions.extend(self.objc.file_extensions().iter().copied()); extensions.extend(self.ocaml.file_extensions().iter().copied()); + #[cfg(feature = "extra-languages")] extensions.extend(self.perl.file_extensions().iter().copied()); extensions.extend(self.php.file_extensions().iter().copied()); extensions.extend(self.python.file_extensions().iter().copied()); + #[cfg(feature = "extra-languages")] extensions.extend(self.r.file_extensions().iter().copied()); extensions.extend(self.ruby.file_extensions().iter().copied()); extensions.extend(self.rust.file_extensions().iter().copied()); @@ -275,26 +314,25 @@ impl ParserRegistry { extensions.extend(self.typescript.file_extensions().iter().copied()); extensions.extend(self.verilog.file_extensions().iter().copied()); extensions.extend(self.yaml.file_extensions().iter().copied()); + #[cfg(feature = "extra-languages")] extensions.extend(self.zig.file_extensions().iter().copied()); extensions } /// Get metrics from all parsers. pub fn all_metrics(&self) -> Vec<(&str, ParserMetrics)> { - vec![ + #[cfg_attr(not(feature = "extra-languages"), allow(unused_mut))] + let mut metrics: Vec<(&str, ParserMetrics)> = vec![ ("bash", self.bash.metrics()), ("c", self.c.metrics()), ("clojure", self.clojure.metrics()), - ("cobol", self.cobol.metrics()), ("cpp", self.cpp.metrics()), ("css", self.css.metrics()), ("csharp", self.csharp.metrics()), - ("dart", self.dart.metrics()), ("dockerfile", self.dockerfile.metrics()), ("elixir", self.elixir.metrics()), ("elm", self.elm.metrics()), ("erlang", self.erlang.metrics()), - ("fortran", self.fortran.metrics()), ("go", self.go.metrics()), ("groovy", self.groovy.metrics()), ("haskell", self.haskell.metrics()), @@ -305,10 +343,8 @@ impl ParserRegistry { ("lua", self.lua.metrics()), ("objc", self.objc.metrics()), ("ocaml", self.ocaml.metrics()), - ("perl", self.perl.metrics()), ("php", self.php.metrics()), ("python", self.python.metrics()), - ("r", self.r.metrics()), ("ruby", self.ruby.metrics()), ("rust", self.rust.metrics()), ("scala", self.scala.metrics()), @@ -319,8 +355,17 @@ impl ParserRegistry { ("typescript", self.typescript.metrics()), ("verilog", self.verilog.metrics()), ("yaml", self.yaml.metrics()), + ]; + #[cfg(feature = "extra-languages")] + metrics.extend([ + ("cobol", self.cobol.metrics()), + ("dart", self.dart.metrics()), + ("fortran", self.fortran.metrics()), + ("perl", self.perl.metrics()), + ("r", self.r.metrics()), ("zig", self.zig.metrics()), - ] + ]); + metrics } /// Check if a file path is supported by any parser. @@ -370,16 +415,12 @@ impl ParserRegistry { Some("bash") } else if self.clojure.can_parse(path) { Some("clojure") - } else if self.cobol.can_parse(path) { - Some("cobol") } else if self.cpp.can_parse(path) { Some("cpp") } else if self.csharp.can_parse(path) { Some("csharp") } else if self.css.can_parse(path) { Some("css") - } else if self.dart.can_parse(path) { - Some("dart") } else if self.dockerfile.can_parse(path) { Some("dockerfile") } else if self.elixir.can_parse(path) { @@ -388,8 +429,6 @@ impl ParserRegistry { Some("elm") } else if self.erlang.can_parse(path) { Some("erlang") - } else if self.fortran.can_parse(path) { - Some("fortran") } else if self.go.can_parse(path) { Some("go") } else if self.groovy.can_parse(path) { @@ -410,14 +449,10 @@ impl ParserRegistry { Some("objc") } else if self.ocaml.can_parse(path) { Some("ocaml") - } else if self.perl.can_parse(path) { - Some("perl") } else if self.php.can_parse(path) { Some("php") } else if self.python.can_parse(path) { Some("python") - } else if self.r.can_parse(path) { - Some("r") } else if self.ruby.can_parse(path) { Some("ruby") } else if self.rust.can_parse(path) { @@ -446,12 +481,36 @@ impl ParserRegistry { Some("verilog") } else if self.yaml.can_parse(path) { Some("yaml") + } else { + self.extra_language_for_path(path) + } + } + + /// Resolve a path to one of the `extra-languages` grammars (cobol / dart / + /// fortran / perl / r / zig). `None` when the feature is disabled. + #[cfg(feature = "extra-languages")] + fn extra_language_for_path(&self, path: &Path) -> Option<&'static str> { + if self.cobol.can_parse(path) { + Some("cobol") + } else if self.dart.can_parse(path) { + Some("dart") + } else if self.fortran.can_parse(path) { + Some("fortran") + } else if self.perl.can_parse(path) { + Some("perl") + } else if self.r.can_parse(path) { + Some("r") } else if self.zig.can_parse(path) { Some("zig") } else { None } } + + #[cfg(not(feature = "extra-languages"))] + fn extra_language_for_path(&self, _path: &Path) -> Option<&'static str> { + None + } } impl Default for ParserRegistry { @@ -471,10 +530,13 @@ mod tests { fn test_parser_registry_new() { let registry = ParserRegistry::new(); assert!(registry.get_parser("c").is_some()); + #[cfg(feature = "extra-languages")] assert!(registry.get_parser("cobol").is_some()); assert!(registry.get_parser("cpp").is_some()); assert!(registry.get_parser("csharp").is_some()); + #[cfg(feature = "extra-languages")] assert!(registry.get_parser("dart").is_some()); + #[cfg(feature = "extra-languages")] assert!(registry.get_parser("fortran").is_some()); assert!(registry.get_parser("go").is_some()); assert!(registry.get_parser("java").is_some()); @@ -483,6 +545,7 @@ mod tests { assert!(registry.get_parser("groovy").is_some()); assert!(registry.get_parser("php").is_some()); assert!(registry.get_parser("python").is_some()); + #[cfg(feature = "extra-languages")] assert!(registry.get_parser("r").is_some()); assert!(registry.get_parser("ruby").is_some()); assert!(registry.get_parser("rust").is_some()); @@ -491,6 +554,7 @@ mod tests { assert!(registry.get_parser("tcl").is_some()); assert!(registry.get_parser("typescript").is_some()); assert!(registry.get_parser("verilog").is_some()); + #[cfg(feature = "extra-languages")] assert!(registry.get_parser("zig").is_some()); } @@ -513,9 +577,11 @@ mod tests { assert!(registry.get_parser("C").is_some()); assert!(registry.get_parser("C++").is_some()); assert!(registry.get_parser("C#").is_some()); + #[cfg(feature = "extra-languages")] assert!(registry.get_parser("COBOL").is_some()); assert!(registry.get_parser("Cpp").is_some()); assert!(registry.get_parser("CSharp").is_some()); + #[cfg(feature = "extra-languages")] assert!(registry.get_parser("FORTRAN").is_some()); assert!(registry.get_parser("Go").is_some()); assert!(registry.get_parser("JAVA").is_some()); @@ -530,10 +596,12 @@ mod tests { assert!(registry.get_parser("Swift").is_some()); assert!(registry.get_parser("TCL").is_some()); assert!(registry.get_parser("TypeScript").is_some()); + #[cfg(feature = "extra-languages")] assert!(registry.get_parser("Dart").is_some()); assert!(registry.get_parser("Lua").is_some()); assert!(registry.get_parser("Groovy").is_some()); assert!(registry.get_parser("Scala").is_some()); + #[cfg(feature = "extra-languages")] assert!(registry.get_parser("Zig").is_some()); } @@ -556,6 +624,7 @@ mod tests { fn test_parser_for_path() { let registry = ParserRegistry::new(); assert!(registry.parser_for_path(&PathBuf::from("test.c")).is_some()); + #[cfg(feature = "extra-languages")] assert!(registry .parser_for_path(&PathBuf::from("test.cob")) .is_some()); @@ -565,6 +634,7 @@ mod tests { assert!(registry .parser_for_path(&PathBuf::from("test.cs")) .is_some()); + #[cfg(feature = "extra-languages")] assert!(registry .parser_for_path(&PathBuf::from("test.f90")) .is_some()); @@ -605,9 +675,11 @@ mod tests { assert!(registry .parser_for_path(&PathBuf::from("test.sv")) .is_some()); + #[cfg(feature = "extra-languages")] assert!(registry .parser_for_path(&PathBuf::from("test.dart")) .is_some()); + #[cfg(feature = "extra-languages")] assert!(registry .parser_for_path(&PathBuf::from("test.zig")) .is_some()); @@ -620,6 +692,7 @@ mod tests { assert!(registry .parser_for_path(&PathBuf::from("build.gradle")) .is_some()); + #[cfg(feature = "extra-languages")] assert!(registry.parser_for_path(&PathBuf::from("test.R")).is_some()); assert!(registry .parser_for_path(&PathBuf::from("test.scala")) @@ -661,9 +734,11 @@ mod tests { fn test_can_parse() { let registry = ParserRegistry::new(); assert!(registry.can_parse(Path::new("test.c"))); + #[cfg(feature = "extra-languages")] assert!(registry.can_parse(Path::new("test.cob"))); assert!(registry.can_parse(Path::new("test.cpp"))); assert!(registry.can_parse(Path::new("test.cs"))); + #[cfg(feature = "extra-languages")] assert!(registry.can_parse(Path::new("test.f90"))); assert!(registry.can_parse(Path::new("test.go"))); assert!(registry.can_parse(Path::new("test.h"))); @@ -678,11 +753,14 @@ mod tests { assert!(registry.can_parse(Path::new("test.swift"))); assert!(registry.can_parse(Path::new("test.tcl"))); assert!(registry.can_parse(Path::new("test.ts"))); + #[cfg(feature = "extra-languages")] assert!(registry.can_parse(Path::new("test.dart"))); + #[cfg(feature = "extra-languages")] assert!(registry.can_parse(Path::new("test.zig"))); assert!(registry.can_parse(Path::new("test.lua"))); assert!(registry.can_parse(Path::new("Service.groovy"))); assert!(registry.can_parse(Path::new("build.gradle"))); + #[cfg(feature = "extra-languages")] assert!(registry.can_parse(Path::new("test.R"))); assert!(registry.can_parse(Path::new("test.scala"))); assert!(!registry.can_parse(Path::new("test.txt"))); @@ -693,51 +771,20 @@ mod tests { fn test_all_metrics() { let registry = ParserRegistry::new(); let metrics = registry.all_metrics(); - assert_eq!(metrics.len(), 38); let names: Vec<&str> = metrics.iter().map(|(n, _)| *n).collect(); - assert_eq!( - names, - vec![ - "bash", - "c", - "clojure", - "cobol", - "cpp", - "css", - "csharp", - "dart", - "dockerfile", - "elixir", - "elm", - "erlang", - "fortran", - "go", - "groovy", - "haskell", - "hcl", - "java", - "julia", - "kotlin", - "lua", - "objc", - "ocaml", - "perl", - "php", - "python", - "r", - "ruby", - "rust", - "scala", - "solidity", - "swift", - "tcl", - "toml", - "typescript", - "verilog", - "yaml", - "zig", - ] - ); + #[cfg_attr(not(feature = "extra-languages"), allow(unused_mut))] + let mut expected = vec![ + "bash", "c", "clojure", "cpp", "css", "csharp", "dockerfile", + "elixir", "elm", "erlang", "go", "groovy", "haskell", "hcl", "java", + "julia", "kotlin", "lua", "objc", "ocaml", "php", "python", "ruby", + "rust", "scala", "solidity", "swift", "tcl", "toml", "typescript", + "verilog", "yaml", + ]; + // Gated grammars are appended after the base set (see `all_metrics`). + #[cfg(feature = "extra-languages")] + expected.extend(["cobol", "dart", "fortran", "perl", "r", "zig"]); + assert_eq!(metrics.len(), expected.len()); + assert_eq!(names, expected); } #[test] @@ -751,6 +798,7 @@ mod tests { registry.language_for_path(&PathBuf::from("test.h")), Some("c") ); + #[cfg(feature = "extra-languages")] assert_eq!( registry.language_for_path(&PathBuf::from("test.cob")), Some("cobol") @@ -771,6 +819,7 @@ mod tests { registry.language_for_path(&PathBuf::from("test.cs")), Some("csharp") ); + #[cfg(feature = "extra-languages")] assert_eq!( registry.language_for_path(&PathBuf::from("test.f90")), Some("fortran") @@ -831,10 +880,12 @@ mod tests { registry.language_for_path(&PathBuf::from("test.sv")), Some("verilog") ); + #[cfg(feature = "extra-languages")] assert_eq!( registry.language_for_path(&PathBuf::from("test.dart")), Some("dart") ); + #[cfg(feature = "extra-languages")] assert_eq!( registry.language_for_path(&PathBuf::from("test.zig")), Some("zig") @@ -851,6 +902,7 @@ mod tests { registry.language_for_path(&PathBuf::from("build.gradle")), Some("groovy") ); + #[cfg(feature = "extra-languages")] assert_eq!( registry.language_for_path(&PathBuf::from("test.R")), Some("r") diff --git a/docs/static-embeddings-plan.md b/docs/static-embeddings-plan.md new file mode 100644 index 0000000..83ae503 --- /dev/null +++ b/docs/static-embeddings-plan.md @@ -0,0 +1,277 @@ +# Plan: distill Jina-Code → static embeddings, validated at every step + +## Goal +Replace ONNX transformer embeddings (BGE/Jina via fastembed) with a **static +model distilled from Jina-Code-V2** — ~100× faster indexing at acceptable +quality — and prove the quality at each step rather than hoping. Secondary win: +removing the `ort`/ONNX runtime kills the 1.5 GB RAM gate, the glibc-2.31 shim, +and the `fastembed` diamond conflict. + +## Progress (branch `feat/static-embeddings`) +- **1.1 ✓** `Embedder` trait; `VectorEngine` holds `Arc` (`eb60736`). +- **1.1 ✓** `StaticEmbedding` backend — model2vec format (`tokenizer.json` + + `safetensors`), `tokenize → gather → mean-pool → L2-norm`; validated against + the real `potion-base-8M` (`a30df5e`). +- **1.2 ✓** gated identifier-splitting in `build_embed_text`, default-off + (`de3e45d`). +- **Speed proven** — `examples/embed_throughput.rs` (`1da5309`): static vs ONNX + BGE-small over 512 symbol texts. Debug floor 8.6×; **release 103× — 32,986 vs + 319 texts/sec** (M4, potion-base-8M 256d vs BGE-small 384d). The embedding step + of indexing drops ~100×. +- **Distillation works (Phase 2.2 ✓)** — `scripts/distill_static_model.py` + distilled `jinaai/jina-embeddings-v2-base-code` → 256d in **~32 s on the M4 + CPU** (Apache-2.0 teacher); loads in Rust (F16 + SIF weights) and embeds at + **70× BGE** (23k vs 327 texts/sec — slower than potion's 103× only because + jina's vocab is 2× larger). +- **Quality (micro-eval) — saturated.** On the 12-query NL→symbol set + jina-code-static-256 and the generic potion floor **tie exactly: R@1 0.92 / + R@3 1.00 / MRR 0.958** vs BGE **1.00 / 1.00 / 1.000** (~95% of BGE). Both + always land the answer in the top 3, so the set **cannot distinguish a code + teacher from a generic one** — it proves the static path is real, fast, and + ~95% of BGE, but is too easy to measure the code-teacher delta. +- **Real eval (965-way, pure-semantic) — the micro-eval was misleading.** + `scripts/extract_eval_corpus.py` (965 doc-commented symbols) + + `examples/embed_eval.rs` (doc→symbol retrieval, recall@k): + + | model | R@1 | R@5 | R@10 | MRR | + |---|---|---|---|---| + | BGE-small (ONNX) | 0.591 | 0.824 | 0.873 | 0.691 | + | potion-base-8M (generic) | 0.378 | 0.616 | 0.696 | 0.488 | + | jina-code-static-256 | 0.379 | 0.598 | 0.685 | 0.480 | + + (1) On a hard, unsaturated task, static is **~65% of BGE R@1 / ~70% MRR** — a + real gap, not ~95%. (2) The code teacher **ties the generic potion** at 256d — + plain Jina distillation at 256d is no win. +- **Lever sweep — the gap is fundamental, not tunable cheaply:** + - **512d: no change** (R@1 0.378) — not a compression problem. + - **code teacher vs generic potion: no change** — not a teacher problem. + - **identifier-splitting (`CODEGRAPH_SPLIT_IDS`): +6% relative** (static R@1 + 0.379→0.401, MRR 0.480→0.511; also lifts BGE 0.591→0.608) — real but modest; + validates the Phase-1.2 lever. + + So static's ~65% of BGE is the **contextualization ceiling** (no attention), + not dim/teacher. Only a learned pooling head addresses the root; tokenlearn may + help modestly. +- **Hybrid eval (40% BM25 + 60% semantic) — static is viable.** `embed_eval` now + scores both; BM25 recovers most of static's semantic gap (lexical matching is + strongest exactly where static is weak): + + | config (id-split on) | R@1 sem → hybrid | MRR sem → hybrid | + |---|---|---| + | jina-code-static-256 | 0.401 → **0.547** | 0.511 → **0.656** | + | BGE-small | 0.608 → 0.609 | 0.716 → 0.720 | + + In the **real hybrid system static is ~90% of BGE** (R@1 0.547 vs 0.609 = 90%; + MRR = 91%; R@10 0.854 vs 0.909 = 94%) — at ~70–100× indexing speed. BGE barely + uses BM25 (already-strong semantics); static leans on it (+36% R@1). Caveat: + a faithful *approximation* of `symbol_search`'s blend; the exact server number + needs the server-side eval. +- **Verdict:** static embedding is a **viable default/opt-in** — ~90% of BGE + end-to-end quality, ~100× faster indexing, none of the ONNX / 1.5 GB RAM-gate / + glibc / fastembed baggage. The code teacher ties the generic potion (no edge, + but Apache-clean). Closing the last ~10% would need a learned pooling head; + ~90% at 100× is a strong default trade. +- **Shipped (Phase 4.1 ✓):** `--embedding-model static` (CLI/MCP) + + `codegraph.embeddingModel: static` / `staticModelPath` (VS Code), via a new + `EmbeddingBackend { Fastembed | Static }` threaded through the server; READMEs + + examples updated. +- **Server-side eval — real `symbol_search` hybrid, 300 doc→symbol queries:** + + | model | R@1 | R@5 | R@10 | MRR | + |---|---|---|---|---| + | BGE-small | 0.457 | 0.717 | 0.787 | 0.568 | + | static (jina-code-static-256) | 0.430 | 0.680 | 0.777 | 0.536 | + + **Static is ~94% of BGE** end-to-end (R@1 94%, MRR 94%, R@10 99%) through the + actual hybrid — confirming the embed_eval approximation, at ~100× indexing + speed. The differing numbers prove the semantic side is live (not BM25 + fallback). Verdict holds: static is a viable default/opt-in. + +## Why this is worth doing (grounded in history) +The original static model (`migration.rs` v3) was `potion-base-8M`: a **generic +256d** static model, fed **raw identifiers**, via stock `encode()`. It lost to +BGE because it was the *weakest* static config vs an eventually code-specialized +transformer — not because static is unworkable. A **code-teacher** (Jina-Code), +at **full dim**, with **split identifiers** has never been tried. The distilled +model runs at static speed (~8000 samples/sec) regardless of how slow the Jina +*teacher* is — the teacher is used once, offline, to build the lookup table. + +## Operating principles (the validation philosophy) +1. **Eval-first.** Build the measuring stick (Phase 0) before any embedding work. + It is the single source of truth; every later step is one row on a scoreboard. +2. **ONNX stays a selectable fallback** the whole way. No bridges burned — a + `--embedding-model` value keeps BGE/Jina available for a "max-quality" mode. +3. **One lever at a time.** Never bundle changes; you won't know what helped. +4. **Two validation levels per model:** + - *Distillation-level* (in Python, before codegraph): does the static model + reproduce the teacher's geometry? (correlation of pairwise cosines) + - *Task-level* (codegraph): code-retrieval recall@k / MRR on a real repo. +5. **Stop rule.** Ship when a config reaches the quality bar (e.g. ≥95% of BGE + recall@5) or the levers stop paying; don't gold-plate. + +--- + +## Phase 0 — Evaluation harness + baselines (the measuring stick) +**Objective:** a deterministic code-retrieval eval and a baseline scoreboard. + +- **0.1 Build a labeled eval set** — ~150–300 `(query → expected symbol(s))` + pairs on a real indexed repo (use this repo + one large external repo). + Sources, cheapest first: docstring → its symbol; test name → tested symbol; + commit subject → changed symbol; ~50 hand-curated natural-language queries. + - **Validate:** ≥150 queries; every target symbol exists in the indexed graph + (assert lookups resolve); eyeball 10 pairs for sanity. +- **0.2 Eval runner** — for each query, run `symbol_search` (ai_query/engine.rs) + and compute recall@{1,5,10} and MRR; also record embed-throughput (samples/s) + and full-repo index wall-clock + peak RAM. + - **Validate:** re-running gives identical numbers (deterministic); a trivially + correct query (exact symbol name) scores recall@1 = 1.0. +- **0.3 Baseline table** — run the eval for every current option: + BGE-small-384d (default), Jina-Code-768d, and — if reconstructable — + potion-base-8M-256d as the floor. + - **Validate:** a committed scoreboard (model → recall@k, MRR, throughput, + index-time, RAM). These are the numbers to beat / match. + +**Gate to Phase 1:** baselines captured and reproducible. + +--- + +## Phase 1 — Rust static plumbing, proven with an off-the-shelf model +**Objective:** make the static path real and correct *before* custom distilling. + +- **1.1 `Embedder` trait + `StaticEmbedding` backend** behind `VectorEngine` + (embedding/engine.rs), using the `model2vec` Rust crate (already used + historically). Select via a new `CodeGraphEmbeddingModel::StaticCode` variant; + fastembed remains default. + - **Validate (unit):** loads a model dir; `embed("test")` returns the expected + dim; `embed_batch == map(embed)`; cosine(x, x) = 1.0; cosine(x, ⟂) ≈ 0. +- **1.2 Identifier-splitting into `build_embed_text`** (ai_query/engine.rs) — + reuse the existing splitter used by the BM25 text index + (ai_query/primitives.rs / text_index.rs). Gate it to the static path (and + optionally the transformer path behind a flag). + - **Validate (unit):** `build_embed_text` for `authenticate_user` yields tokens + containing `authenticate` and `user`; snapshot 5 representative symbols. +- **1.3 Eval with off-the-shelf `potion-retrieval-32M`** through the new path + (with and without id-splitting). + - **Validate (task):** scoreboard rows for {potion-retrieval, ±id-split} vs the + BGE baseline and the potion-base-8M floor. This proves the plumbing *and* + quantifies how much *better-general-static + id-splitting* alone recovers. + +**Gate to Phase 2:** static path is correct end-to-end; we know the id-split +delta and whether a better general static helps. (If it already ≈ BGE, great — +distillation is upside, not a requirement.) + +--- + +## Phase 2 — Distill Jina-Code → static (offline, Python, M4 CPU) +**Objective:** a code-teacher static model, validated against the teacher before +it touches codegraph. + +- **2.1 Env + teacher** — `pip install model2vec[distill]`; download + `jinaai/jina-embeddings-v2-base-code`. + - **License:** the teacher is **Apache-2.0** (verified). The distilled model is + a derivative; Apache-2.0 is permissive, so the static model is redistributable + under Apache-2.0 with attribution (§4 — credit the teacher in NOTICE / the + model card). Use the **v2** family only — `jina-embeddings-v3` is + CC-BY-NC-4.0 (non-commercial) and must not be the teacher. + - **Validate:** load the teacher in Python; embed a code snippet → 768d; + confirms it runs locally on the M4 (CPU/MPS). +- **2.2 Distill → 256d** (same dim/speed as the old potion, to isolate one + variable: *code teacher vs generic teacher*). + `distill(model_name="jinaai/jina-embeddings-v2-base-code", pca_dims=256)`; + `save_pretrained("jina-code-static-256")`. + - **Validate (distillation-level):** on a held-out set of ~500 code-snippet + pairs, compute Spearman correlation between *static* pairwise cosines and + *teacher* pairwise cosines. Require ρ ≥ ~0.6 (the static model preserves the + teacher's geometry). Also measure embed throughput → confirm ~static speed + (thousands/sec), independent of Jina's slowness. + - **Caveat:** Jina ships custom modeling (ALiBi, `trust_remote_code=True`); if + distillation fights it, fall back to BGE-small as the teacher (still the + model we trust; re-validate ρ). +- **2.3 Eval `jina-code-static-256` in codegraph** (via Phase-1 plumbing, + id-split on). + - **Validate (task):** scoreboard vs the potion-base-8M-256 floor at **equal + dim and speed**. A meaningful win here proves *the teacher was the problem*. + +**Gate to Phase 3:** code-static-256 beats generic-static-256 on the eval. + +--- + +## Phase 3 — Quality levers, each measured marginally +**Objective:** close the gap to the transformer baseline; keep only levers that +pay. One change → one scoreboard row. + +- **3.1 Dimension → 512d** (re-distill `pca_dims=512`). + - **Validate:** recall@k delta vs 256d; index-size/search-latency delta. Keep + if the quality gain justifies the larger index (embed throughput barely + changes — cost is tokenization, not pooling width). +- **3.2 SIF pooling + drop top principal component** (Rust `StaticEmbedding` + post-process, or model2vec weighting). + - **Validate:** marginal recall delta; keep if positive. +- **3.3 (Optional, GPU-helped) Tokenlearn corpus adaptation** on a code corpus + (the teacher-over-corpus forward is the only GPU-helped step; CPU/overnight ok). + - **Validate:** marginal delta vs the plain distill — is the corpus step worth + the cost? +- **3.4 (Optional, ceiling) tiny learned pooling head** distilled to match + Jina-Code's sentence embedding (one shallow attention/conv layer). + - **Validate:** how much of the remaining gap to the Jina *transformer* it + closes, at what added inference cost (must stay ≫ transformer-fast). + +**Gate to Phase 4:** a chosen config meets the quality bar (define vs BGE, e.g. +≥95% recall@5) or levers plateau. + +--- + +## Phase 4 — Integration + migration + large-repo validation +**Objective:** ship the swap safely; prove the real-world speedup. + +- **4.1 Default/opt-in wiring + DB migration** — new `migration.rs` version + (delete `vec:`, clear embeddings, re-embed on load — mirror v3→v4/v4→v5). + - **Validate:** migration unit test (old-vectors DB migrates, re-embeds, search + still returns results) — mirror the existing `test_migration_with_json_data`. +- **4.2 Index a large real repo on the static path** end-to-end. + - **Validate:** index wall-clock + peak RAM vs the ONNX path (target: large + index-time drop, no 1.5 GB gate, RAM down); eval recall holds vs Phase-3; + confirm the `ort`/fastembed dependency is gone from the build (bonus: this is + what unblocks the in-process Warp linking). +- **4.3 A/B + soak** — run static vs BGE on the same repo; spot-check real + developer queries; check the *other* embedding consumers + (`find_duplicates`, `cluster_symbols`, `find_similar`). + - **Validate:** qualitative parity on real queries; no regression in the other + consumers. + +--- + +## Phase 5 — Decision gate +Final scoreboard: `static-jina-final` vs `BGE` vs `Jina-transformer` across +{recall@k, MRR, index-time, embed-throughput, RAM, index-size, deps}. Decide: +static as default? opt-in? keep ONNX as a `--embedding-model max-quality` +fallback? Record the call. + +## Scoreboard template (the through-line) +| config | recall@1 | recall@5 | recall@10 | MRR | embed/s | index-time | RAM | dim | deps | +|---|---|---|---|---|---|---|---|---|---| +| BGE-small (baseline) | 0.591 | 0.824 | 0.873 | 0.691 | 327 | | | 384 | ort | +| Jina-Code (baseline) | | | | | | | | 768 | ort | +| potion-base-8M (floor) | 0.378 | 0.616 | 0.696 | 0.488 | 32986 | | | 256 | — | +| **jina-code-static-256** | 0.379 | 0.598 | 0.685 | 0.480 | 23037 | | | 256 | — | +| jina-code-static-512 | | | | | | | | 512 | — | +| + tokenlearn / + head | | | | | | | | | — | + +> recall@k from the **965-way pure-semantic eval** (`embed_eval`, doc→symbol). +> The real hybrid (40% BM25 + 60% semantic) system would score higher for every +> row — this isolates *embedding* quality, where static lands ~65% of BGE R@1. + +## Safety / rollback +- fastembed/ONNX stays selectable throughout; nothing is removed until Phase 5. +- Each phase is independently revertable; the eval scoreboard makes any + regression visible immediately. +- Distillation artifacts are reproducible from the one-line `distill` command; + no opaque state. + +## First actions (no Python, no GPU needed) +1. Phase 0.1–0.3: eval harness + baselines. +2. Phase 1.1–1.2: `StaticEmbedding` trait/backend + id-splitting in + `build_embed_text`. +3. Phase 1.3: off-the-shelf `potion-retrieval-32M` through the eval. +These isolate the cheapest wins and stand up the measuring stick before any +distillation. diff --git a/mcp-package/README.md b/mcp-package/README.md index 268aa34..6a71306 100644 --- a/mcp-package/README.md +++ b/mcp-package/README.md @@ -48,7 +48,7 @@ Pass flags after `--`: |------|---------|-------------| | `--workspace ` | current dir | Directories to index (repeatable) | | `--exclude ` | — | Directories to skip (repeatable) | -| `--embedding-model ` | `bge-small` | `bge-small`, `jina-code-v2`, or `granite-97m` (32K context, multilingual) | +| `--embedding-model ` | `bge-small` | `bge-small`, `jina-code-v2`, `granite-97m` (32K, multilingual), or `static` (model2vec, ~100× faster indexing, ~90% of BGE quality; needs a local model dir via `CODEGRAPH_STATIC_MODEL`) | | `--max-files ` | 5000 | Maximum files to index | | `--profile ` | `all` | Scope tool surface: `core` (8), `graph` (16), `memory` (14), `security` (pro), `all` (42) | | `--graph-only` | off | Skip embeddings — graph + structural tools only. No ONNX model load, 10-50× faster indexing. For CI / one-shot graph queries. | diff --git a/mcp-package/bin/postinstall.js b/mcp-package/bin/postinstall.js index 36e2d61..4d2d8d8 100644 --- a/mcp-package/bin/postinstall.js +++ b/mcp-package/bin/postinstall.js @@ -57,6 +57,30 @@ try { console.warn(` ${err.message}`); } +// Fetch the distilled static embedding model (best-effort) from the +// release-independent `model` GitHub release. Only needed for +// `--embedding-model static`; skipped if already present or if +// CODEGRAPH_SKIP_MODEL_FETCH is set. Never fails the install. +if (!process.env.CODEGRAPH_SKIP_MODEL_FETCH) { + const MODEL = "jina-code-static-256"; + const modelDir = path.join(os.homedir(), ".codegraph", "static_models", MODEL); + if (!fs.existsSync(path.join(modelDir, "model.safetensors"))) { + try { + fs.mkdirSync(modelDir, { recursive: true }); + const url = `https://github.com/codegraph-ai/CodeGraph/releases/download/model/${MODEL}.tar.gz`; + const tgz = path.join(modelDir, "_model.tar.gz"); + execFileSync("curl", ["-fsSL", url, "-o", tgz], { timeout: 180000 }); + execFileSync("tar", ["xzf", tgz, "-C", modelDir], { timeout: 60000 }); + fs.unlinkSync(tgz); + console.log(`✓ codegraph-mcp: static embedding model ready (${modelDir})`); + } catch { + console.warn( + `ℹ codegraph-mcp: static model not fetched (optional — only for --embedding-model static)` + ); + } + } +} + // Hint about the optional Claude Code hook. Installation is opt-in to avoid // silently modifying the user's ~/.claude/settings.json. Both Unix // (bash) and Windows (PowerShell) variants are shipped — the installer diff --git a/mcp-package/package.json b/mcp-package/package.json index 4c6c87d..335d570 100644 --- a/mcp-package/package.json +++ b/mcp-package/package.json @@ -1,6 +1,6 @@ { "name": "@astudioplus/codegraph-mcp", - "version": "0.18.6", + "version": "0.19.0", "mcpName": "io.github.codegraph-ai/codegraph", "description": "CodeGraph MCP server — cross-language code intelligence with 42 tools, 38 languages", "author": "Andrey Vasilevsky ", diff --git a/mcp-package/server.json b/mcp-package/server.json index d3a4104..b3a01b3 100644 --- a/mcp-package/server.json +++ b/mcp-package/server.json @@ -6,12 +6,12 @@ "url": "https://github.com/codegraph-ai/CodeGraph", "source": "github" }, - "version": "0.18.6", + "version": "0.19.0", "packages": [ { "registryType": "npm", "identifier": "@astudioplus/codegraph-mcp", - "version": "0.18.6", + "version": "0.19.0", "transport": { "type": "stdio" }, diff --git a/scripts/distill_static_model.py b/scripts/distill_static_model.py new file mode 100644 index 0000000..9aa2854 --- /dev/null +++ b/scripts/distill_static_model.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python3 +"""Distill a transformer embedding model into a static (model2vec) model. + +Produces `config.json` + `tokenizer.json` + `model.safetensors` (an +`embeddings` [vocab, dim] F32 tensor) — exactly the format codegraph-memory's +`StaticEmbedding` (`VectorEngine::with_static_model`) loads. No runtime ONNX. + +Runs on CPU (M4) in minutes — the teacher is used once to build the lookup +table, then discarded. Uses `distill_from_model` with an explicit +`trust_remote_code` load so Jina-Code's custom modeling (ALiBi) doesn't fight +the higher-level `distill()` helper. + +Usage: + python distill_static_model.py [TEACHER] [PCA_DIMS] [OUTPUT_DIR] + +Defaults reproduce the plan's first experiment — the Apache-2.0 code teacher, +256d (same dim/speed as the generic potion floor, isolating "code teacher vs +generic teacher"): + python distill_static_model.py \\ + jinaai/jina-embeddings-v2-base-code 256 \\ + ~/.codegraph/static_models/jina-code-static-256 +""" +import sys +from pathlib import Path + + +def main() -> int: + teacher = sys.argv[1] if len(sys.argv) > 1 else "jinaai/jina-embeddings-v2-base-code" + pca_dims = int(sys.argv[2]) if len(sys.argv) > 2 else 256 + out = ( + Path(sys.argv[3]).expanduser() + if len(sys.argv) > 3 + else Path.home() / ".codegraph/static_models/jina-code-static-256" + ) + + print(f"[distill] teacher={teacher} pca_dims={pca_dims} out={out}", flush=True) + + from transformers import AutoModel, AutoTokenizer + from model2vec.distill import distill_from_model + + # Explicit load with trust_remote_code so Jina's custom code is honored. + model = AutoModel.from_pretrained(teacher, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(teacher) + + static = distill_from_model(model=model, tokenizer=tokenizer, pca_dims=pca_dims) + + out.mkdir(parents=True, exist_ok=True) + static.save_pretrained(str(out)) + + # Sanity: confirm the saved model embeds and has the expected dim. + vec = static.encode(["fn authenticate_user(token: Token) -> User"]) + dim = vec.shape[-1] + print(f"[distill] done — output dim={dim} (expected {pca_dims}); saved to {out}", flush=True) + return 0 if dim == pca_dims else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/extract_eval_corpus.py b/scripts/extract_eval_corpus.py new file mode 100644 index 0000000..fa1dfaa --- /dev/null +++ b/scripts/extract_eval_corpus.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +"""Extract a (name, signature, doc) retrieval-eval corpus from Rust source. + +Scans `/**/*.rs` for doc-commented items (`///` lines, optionally followed +by `#[..]` attributes, then `pub fn|struct|trait|enum NAME`). Emits JSON +`[{id, signature, doc}]` for items whose doc has >= MIN_DOC_WORDS words. + +The eval (examples/embed_eval.rs) embeds `"name: signature"` as the symbol and +the **doc as the query** (the symbol text excludes the doc), so it's a real +natural-language-description -> code-signature retrieval task. The same task is +run for every embedder, so the *comparison* is fair regardless of difficulty. + + python scripts/extract_eval_corpus.py [ROOT=crates] [OUT=/tmp/cg_eval.json] +""" +import json +import re +import sys +from pathlib import Path + +MIN_DOC_WORDS = 6 +ROOT = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("crates") +OUT = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("/tmp/codegraph_eval_corpus.json") + +ITEM_RE = re.compile( + r"^\s*(?:pub(?:\([^)]*\))?\s+)?(?:async\s+)?(?:fn|struct|trait|enum)\s+([A-Za-z_][A-Za-z0-9_]*)" +) + + +def main() -> int: + corpus: dict[str, dict] = {} + for f in ROOT.rglob("*.rs"): + p = str(f) + if any(seg in p for seg in ("/tests/", "/examples/", "/benches/")): + continue + doc: list[str] = [] + for line in f.read_text(errors="ignore").splitlines(): + s = line.strip() + if s.startswith("///"): + doc.append(s[3:].strip()) + elif s.startswith("#["): + pass # attribute between doc and item — keep the doc block + elif s == "": + doc = [] + else: + m = ITEM_RE.match(line) + if m and doc: + name = m.group(1) + sig = s.rstrip().rstrip("{").rstrip() + doctext = " ".join(d for d in doc if d and not d.startswith("#")) + if name not in corpus and len(doctext.split()) >= MIN_DOC_WORDS: + corpus[name] = {"id": name, "signature": sig, "doc": doctext} + doc = [] + items = list(corpus.values()) + OUT.write_text(json.dumps(items)) + print(f"extracted {len(items)} doc-commented symbols -> {OUT}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/extract_fullbody_corpus.py b/scripts/extract_fullbody_corpus.py new file mode 100644 index 0000000..7745789 --- /dev/null +++ b/scripts/extract_fullbody_corpus.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +"""Extract full-body embed texts (decl + body, ~2KB cap) — mirrors codegraph's +`--full-body-embedding` symbol text. Output JSON [{id, signature}] where +`signature` holds the whole text, so examples/embed_throughput.rs can consume it +via CODEGRAPH_THROUGHPUT_CORPUS for a static-vs-BGE throughput test on long texts. + + python scripts/extract_fullbody_corpus.py [ROOT=crates] [OUT=/tmp/cg_fullbody.json] +""" +import json +import re +import sys +from pathlib import Path + +CAP = 2048 # codegraph FULL_BODY_MAX_CHARS +ROOT = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("crates") +OUT = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("/tmp/codegraph_fullbody_corpus.json") +DECL = re.compile( + r"^\s*(?:pub(?:\([^)]*\))?\s+)?(?:async\s+)?(?:fn|struct|trait|enum|impl)\s+([A-Za-z_][A-Za-z0-9_]*)" +) + + +def main() -> int: + items = {} + for f in ROOT.rglob("*.rs"): + p = str(f) + if any(s in p for s in ("/tests/", "/examples/", "/benches/")): + continue + lines = f.read_text(errors="ignore").splitlines() + for i, line in enumerate(lines): + m = DECL.match(line) + if not m: + continue + name = m.group(1) + if name in items: + continue + body = "\n".join(lines[i : i + 45])[:CAP] + if len(body) >= 40: + items[name] = {"id": name, "signature": body} + out = list(items.values()) + OUT.write_text(json.dumps(out)) + print(f"extracted {len(out)} full-body symbol texts -> {OUT}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/fetch-static-model.sh b/scripts/fetch-static-model.sh new file mode 100755 index 0000000..2f886fd --- /dev/null +++ b/scripts/fetch-static-model.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# Fetch the distilled static embedding model from the release-independent +# `model` GitHub release (decoupled from versioned releases — bumping the app +# version never requires re-uploading the model). +# +# Used two ways: +# • Package time — place the model into the VS Code extension bundle: +# scripts/fetch-static-model.sh vscode/bin/jina-code-static-256 +# • Manually — populate the server's default resolve path: +# scripts/fetch-static-model.sh +# +# Usage: scripts/fetch-static-model.sh [DEST_DIR] [MODEL_NAME] +# DEST_DIR default ~/.codegraph/static_models/ +# MODEL_NAME default jina-code-static-256 +# +# Expects a `.tar.gz` asset on the `model` release whose contents +# are the model files at the archive root, created with: +# tar czf jina-code-static-256.tar.gz -C . +set -euo pipefail + +MODEL="${2:-jina-code-static-256}" +DEST="${1:-$HOME/.codegraph/static_models/$MODEL}" +URL="https://github.com/codegraph-ai/CodeGraph/releases/download/model/${MODEL}.tar.gz" + +mkdir -p "$DEST" +echo "fetching ${MODEL}.tar.gz -> $DEST" +curl -fsSL "$URL" | tar xz -C "$DEST" + +if [ ! -f "$DEST/model.safetensors" ]; then + echo "error: model.safetensors missing after extract — check the '$MODEL' release asset" >&2 + exit 1 +fi +echo "static model ready: $(ls "$DEST" | tr '\n' ' ')" diff --git a/scripts/server_eval.py b/scripts/server_eval.py new file mode 100644 index 0000000..e73eddd --- /dev/null +++ b/scripts/server_eval.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +"""Server-side retrieval eval: drive the real codegraph-server MCP and score +`symbol_search` (the actual BM25+semantic hybrid) for static vs BGE. + +For each model: spawn an isolated server (temp HOME, symlinked fastembed cache), +MCP-handshake, force-reindex, wait for embeddings, then run doc->symbol queries +through `codegraph_symbol_search` and score recall@{1,5,10} + MRR. + + python scripts/server_eval.py [N_QUERIES=300] +""" +import json, os, shutil, subprocess, sys, tempfile, time +from pathlib import Path + +REPO = "/Users/anvanster/projects/codegraph" +SERVER = f"{REPO}/target/release/codegraph-server" +CORPUS = "/tmp/codegraph_eval_corpus.json" +CACHE = os.path.expanduser("~/.codegraph/fastembed_cache") +N = int(sys.argv[1]) if len(sys.argv) > 1 else 300 + + +def run_eval(label, model, static_path=None): + home = tempfile.mkdtemp(prefix="cgse.") + os.makedirs(f"{home}/.codegraph", exist_ok=True) + os.symlink(CACHE, f"{home}/.codegraph/fastembed_cache") + env = dict(os.environ, HOME=home, FASTEMBED_CACHE_DIR=CACHE, RUST_LOG="info") + if static_path: + env["CODEGRAPH_STATIC_MODEL"] = static_path + logf = open(f"{home}/server.log", "w+") + proc = subprocess.Popen( + [SERVER, "--mcp", "-w", REPO, "--embedding-model", model, "--full-body-embedding"], + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=logf, env=env, text=True, bufsize=1, + ) + + def send(o): + proc.stdin.write(json.dumps(o) + "\n"); proc.stdin.flush() + + def recv(want, timeout=120): + t0 = time.time() + while time.time() - t0 < timeout: + line = proc.stdout.readline() + if not line: + return None + line = line.strip() + if not line: + continue + try: + m = json.loads(line) + except Exception: + continue + if m.get("id") == want: + return m + return None + + send({"jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "eval", "version": "1"}}}) + recv(1) + send({"jsonrpc": "2.0", "method": "notifications/initialized"}) + send({"jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": {"name": "codegraph_reindex_workspace", "arguments": {"force": True}}}) + recv(2, timeout=120) + + print(f"[{label}] indexing + embedding...", flush=True) + t0 = time.time() + while time.time() - t0 < 1800: + if "embedding generation complete" in open(f"{home}/server.log").read(): + break + time.sleep(5) + embed_secs = time.time() - t0 + print(f"[{label}] embeddings ready in ~{embed_secs:.0f}s", flush=True) + + corpus = json.load(open(CORPUS))[:N] + r1 = r5 = r10 = mrr = 0 + nid = 100 + for item in corpus: + nid += 1 + send({"jsonrpc": "2.0", "id": nid, "method": "tools/call", + "params": {"name": "codegraph_symbol_search", + "arguments": {"query": item["doc"], "limit": 10, "compact": True}}}) + m = recv(nid, timeout=60) + names = [] + if m and "result" in m: + res, data = m["result"], None + try: + if isinstance(res, dict) and "content" in res: + data = json.loads(res["content"][0]["text"]) + elif isinstance(res, dict) and "results" in res: + data = res + elif isinstance(res, str): + data = json.loads(res) + except Exception: + data = None + if data: + names = [r.get("symbol", {}).get("name") for r in data.get("results", [])] + rank = next((i + 1 for i, nm in enumerate(names) if nm == item["id"]), None) + if rank: + r1 += rank <= 1 + r5 += rank <= 5 + r10 += rank <= 10 + mrr += 1.0 / rank + n = len(corpus) + proc.terminate() + try: + proc.wait(timeout=10) + except Exception: + proc.kill() + shutil.rmtree(home, ignore_errors=True) + print(f"RESULT {label}: R@1 {r1/n:.3f} R@5 {r5/n:.3f} R@10 {r10/n:.3f} " + f"MRR {mrr/n:.3f} (n={n}, embed {embed_secs:.0f}s)", flush=True) + + +if __name__ == "__main__": + run_eval("BGE-small", "bge-small") + run_eval("static", "static", static_path=os.path.expanduser("~/.codegraph/static_models/jina-code-static-256")) diff --git a/vscode/README.md b/vscode/README.md index 0e49f6a..a2a53bf 100644 --- a/vscode/README.md +++ b/vscode/README.md @@ -45,7 +45,7 @@ The extension starts the server automatically and registers all tools as Languag |------|---------|-------------| | `--workspace ` | current dir | Directories to index (repeatable for multi-project) | | `--exclude ` | — | Directories to skip (repeatable) | -| `--embedding-model ` | `bge-small` | `bge-small` (384d, fast) or `jina-code-v2` (768d, 6x slower) | +| `--embedding-model ` | `bge-small` | `bge-small` (384d, fast), `jina-code-v2` (768d, 6x slower), or `static` (model2vec, 256d — ~100× faster indexing, no ONNX, ~90% of BGE quality in hybrid search; the model ships bundled in the extension, so no setup is needed) | | `--full-body-embedding` | `true` | Embed full function body (~50 lines) for better semantic search and duplicate detection | | `--max-files ` | 5000 | Maximum files to index | @@ -56,7 +56,8 @@ The extension starts the server automatically and registers all tools as Languag "codegraph.indexOnStartup": true, "codegraph.indexPaths": ["/path/to/project-a", "/path/to/project-b"], "codegraph.excludePatterns": ["**/cmake-build-debug/**", "**/generated/**"], - "codegraph.embeddingModel": "bge-small", + "codegraph.embeddingModel": "bge-small", // or "static" for ~100× faster indexing (model bundled, no path needed) + "codegraph.staticModelPath": "", // optional: override the bundled model2vec dir "codegraph.maxFileSizeKB": 1024, "codegraph.debug": false } diff --git a/vscode/package.json b/vscode/package.json index 23eb5b1..8628768 100644 --- a/vscode/package.json +++ b/vscode/package.json @@ -2,7 +2,7 @@ "name": "codegraph", "displayName": "CodeGraph", "description": "Cross-language code intelligence powered by graph analysis", - "version": "0.18.6", + "version": "0.19.0", "publisher": "aStudioPlus", "author": "Andrey Vasilevsky ", "license": "Apache-2.0", @@ -274,16 +274,24 @@ "type": "string", "enum": [ "bge-small", - "jina-code-v2" + "jina-code-v2", + "static" ], "default": "bge-small", "scope": "resource", "enumDescriptions": [ "BGE-Small-EN-v1.5 (384d) — fast, good quality with full-body embeddings. ~127MB download.", - "Jina Code V2 (768d) — 6x slower indexing, no quality advantage with full-body embeddings. ~642MB download." + "Jina Code V2 (768d) — 6x slower indexing, no quality advantage with full-body embeddings. ~642MB download.", + "Static (model2vec, 256d) — ~100x faster indexing, no ONNX runtime or 1.5GB RAM gate, ~90% of BGE quality in hybrid search. Requires a local model dir (codegraph.staticModelPath or CODEGRAPH_STATIC_MODEL env)." ], "description": "Embedding model for semantic search and code similarity" }, + "codegraph.staticModelPath": { + "type": "string", + "default": "", + "scope": "resource", + "description": "Directory of the model2vec static model (config.json + tokenizer.json + model.safetensors) used when embeddingModel is 'static'. Empty defaults to ~/.codegraph/static_models/jina-code-static-256. Distill one with scripts/distill_static_model.py." + }, "codegraph.fullBodyEmbedding": { "type": "boolean", "default": true, diff --git a/vscode/src/extension.ts b/vscode/src/extension.ts index 522c8cf..8caf67b 100644 --- a/vscode/src/extension.ts +++ b/vscode/src/extension.ts @@ -292,7 +292,20 @@ export async function activate(context: vscode.ExtensionContext): Promise // space the way `shell:true` + cmd.exe did. stdio defaults to pipes, which // vscode-languageclient uses for the LSP transport (stderr → outputChannel). const serverOptions: ServerOptions = () => { - const child = cp.spawn(serverModule, [], { cwd: context.extensionPath }); + // When the static (model2vec) embedding model is selected, point the + // server at the model dir via CODEGRAPH_STATIC_MODEL — the server + // resolves the static path from this env, falling back to + // ~/.codegraph/static_models/jina-code-static-256. + const wsFolder = vscode.workspace.workspaceFolders?.[0]?.uri; + const cfg = vscode.workspace.getConfiguration('codegraph', wsFolder); + const spawnEnv = { ...process.env }; + if (cfg.get('embeddingModel') === 'static') { + // staticModelPath override, else the model bundled next to the binary. + const staticModelPath = cfg.get('staticModelPath') + || path.join(context.extensionPath, 'bin', 'jina-code-static-256'); + spawnEnv.CODEGRAPH_STATIC_MODEL = staticModelPath; + } + const child = cp.spawn(serverModule, [], { cwd: context.extensionPath, env: spawnEnv }); child.once('exit', (code, signal) => { lastExitCode = code; lastExitSignal = signal; @@ -333,6 +346,7 @@ export async function activate(context: vscode.ExtensionContext): Promise indexPaths: latestConfig.get('indexPaths'), maxFileSizeKB: latestConfig.get('maxFileSizeKB'), embeddingModel: latestConfig.get('embeddingModel'), + staticModelPath: latestConfig.get('staticModelPath'), fullBodyEmbedding: latestConfig.get('fullBodyEmbedding'), embedOnOpen: latestConfig.get('embedOnOpen'), };