Feat/embedding retrieve rerank traverse

Cargo.toml

@@ -14,6 +14,13 @@ categories = ["command-line-utilities", "data-structures"]
 [lib]
 crate-type = ["rlib"]
 
+[features]
+default = []
+# Enable vector retrieval + cross-encoder rerank + adaptive KG traversal.
+# Adds fastembed (embedding + reranker inference via ONNX) and usearch (ANN index).
+# Off by default to keep the default binary slim.
+embeddings = ["dep:fastembed", "dep:usearch"]
+
 [dependencies]
 clap = { version = "4", features = ["derive"] }
 serde = { version = "1", features = ["derive"] }
@@ -56,6 +63,13 @@ tree-sitter = "0.26"
 tree-sitter-perl = "1.1.2"
 sysinfo = "0.32"
 
+# Embedding retrieval stack (optional, behind `embeddings` feature).
+# fastembed: in-process embedding inference AND cross-encoder reranking
+#   via ONNX Runtime. Covers BGE-small-en-v1.5 (embed) + bge-reranker-v2-m3 (rerank).
+# usearch: HNSW ANN index, SIMD cosine search, file persistence.
+fastembed = { version = "4", optional = true }
+usearch = { version = "2", optional = true }
+
 [target.'cfg(unix)'.dependencies]
 libc = "0.2"
 

docs/mcp-setup.md

@@ -131,3 +131,51 @@ When the MCP server starts with an existing LeanKG project, it checks if the ind
 ## Fallback
 
 If the MCP server reports "LeanKG not initialized", manually run `leankg init` in your project directory, then restart the AI tool.
+
+## Embedding Retrieval (optional, `embeddings` feature)
+
+The `kg_semantic_context` tool — vector retrieve + cross-encoder rerank + adaptive KG traversal — only ships when LeanKG is built with `--features embeddings`. Default builds skip it to keep the binary lean.
+
+### Building with the feature
+
+```bash
+# From a LeanKG checkout:
+cargo build --release --features embeddings
+
+# Or directly install the binary with the feature on:
+cargo install --path . --features embeddings
+```
+
+This pulls in `fastembed` (ONNX-backed embedding + reranker inference) and `usearch` (HNSW ANN index). The first build downloads ONNX Runtime binaries via fastembed's deps.
+
+### One-time setup per machine
+
+```bash
+# 1. Pre-download embedding (BGE-small-en-v1.5, ~130MB) and reranker
+#    (bge-reranker-v2-m3, ~600MB) into ~/.cache/leankg/models:
+leankg embed --init
+
+# 2. Index your project (if not already indexed):
+leankg index ./src
+
+# 3. Build the embedding index (~seconds for incremental, minutes for a
+#    fresh 10k-node repo on CPU):
+leankg embed
+```
+
+### Index lifecycle
+
+`leankg embed` (default) is **incremental**: it reads the `embedding_state` CozoDB table that tracks per-node freshness and only re-embeds nodes that are stale (touched by a recent `index` run), missing (newly added), or whose text blob hash changed. Orphans (state rows whose `qualified_name` is no longer in `code_elements`) are reaped.
+
+`leankg embed --full` ignores state and re-embeds every node. Use after a model swap or suspected index corruption.
+
+The `index` command marks touched elements stale but does **not** trigger `embed` automatically — embedding is a separate explicit step. The MCP tool surfaces a stale-embeddings warning in `diagnostics.embeddings_stale` so callers know when to re-run `embed`.
+
+### Worktree exclusion
+
+By default, `kg_semantic_context` filters out paths under `.worktrees/`, `.claude/worktrees/`, and `.opencode/worktrees/` to avoid duplicate-noise from agent scratch copies. Pass `include_worktrees: true` to include them.
+
+### Reranker fallback
+
+If the reranker fails to load or score, the tool falls back to ANN-order top-N (no cross-encoder). `diagnostics.reranker` will be `"fallback_ann"` instead of `"bge-reranker-v2-m3"`. The most common cause is a partial model download — re-running `leankg embed --init` fixes it.
+

docs/mcp-tools.md

@@ -66,6 +66,40 @@ LeanKG exposes a comprehensive set of MCP tools for AI tools to query the knowle
 
 When the MCP server starts without an existing LeanKG project, it automatically initializes and indexes the current directory. This provides a "plug and play" experience for AI tools.
 
+## Semantic Retrieval (optional, `embeddings` feature)
+
+These tools ship only when LeanKG is built with `--features embeddings`. They add vector retrieval + cross-encoder rerank + adaptive graph traversal on top of the existing keyword/graph search.
+
+| Tool | Description |
+|------|-------------|
+| `kg_semantic_context` | Vector retrieve → rerank → traverse. Best for natural-language questions where keyword search misses (e.g., 'where do we validate access rights'). Returns ranked seed nodes plus 1-2 hop graph context. |
+
+Setup (one-time):
+
+```bash
+cargo run --release --features embeddings -- embed --init        # pre-download models (~700MB)
+cargo run --release --features embeddings -- embed               # build the embedding index
+```
+
+Then call from any MCP client:
+
+```json
+{ "tool": "kg_semantic_context", "arguments": { "query": "where is refund failure handled" } }
+```
+
+Optional arguments: `env` (default `local`), `top_k` (default 50), `rerank_top_n` (default 10), `traverse` (default true), `include_worktrees` (default false), `debug` (default false).
+
+Response shape (debug=false): `{ query, env, seeds[], traversed[] }`. With `debug=true`: adds `diagnostics` with reranker status, candidate counts, per-stage latency, and the edges traversed.
+
+Behavior notes:
+- If the reranker fails to load, the tool silently falls back to ANN-order top-N (Q4 option A). `diagnostics.reranker = "fallback_ann"` surfaces this.
+- If the embedding index is older than the last `index` run, `diagnostics.embeddings_stale = true` (still serves, just warns).
+- Worktree scratch copies (`.worktrees/`, `.claude/worktrees/`, `.opencode/worktrees/`) are filtered out by default to avoid duplicate-noise results.
+
+## Auto-Indexing
+
+When the MCP server starts with an existing LeanKG project, it checks if the index is stale (by comparing git HEAD commit time vs database file modification time). If stale, it automatically runs incremental indexing to ensure AI tools have up-to-date context.
+
 ## Auto-Indexing
 
 When the MCP server starts with an existing LeanKG project, it checks if the index is stale (by comparing git HEAD commit time vs database file modification time). If stale, it automatically runs incremental indexing to ensure AI tools have up-to-date context.