From 72687712d3c5d17c0df408a68a2ad85cd60e1518 Mon Sep 17 00:00:00 2001 From: ajianaz Date: Wed, 22 Jul 2026 23:51:42 +0700 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20Phase=203=20=E2=80=94=20Brain=20Mod?= =?UTF-8?q?e=20hybrid=20search=20(usearch=20+=20RRF)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add CodeVectorIndex (usearch HNSW) with file locking (fs2) - Schema v4: embedding metadata columns on projects table - brain_search(): FTS5 + vector KNN + graph BFS → RRF k=60 fusion - Index-time embedding via static tokens (256d) during cora index - CLI: cora brain [--json] [--limit N] - MCP tool: cora.brain_search Closes #359 --- Cargo.lock | 157 ++++++++++++++++++ Cargo.toml | 4 + src/embed/mod.rs | 4 +- src/index/brain.rs | 262 ++++++++++++++++++++++++++++++ src/index/mod.rs | 17 ++ src/index/schema.rs | 20 ++- src/index/vector.rs | 384 ++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 61 +++++++ src/mcp/tools.rs | 58 ++++++- 9 files changed, 961 insertions(+), 6 deletions(-) create mode 100644 src/index/brain.rs create mode 100644 src/index/vector.rs diff --git a/Cargo.lock b/Cargo.lock index 3a90fcf..6b4f939 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -257,6 +257,17 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" +[[package]] +name = "codespan-reporting" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" +dependencies = [ + "serde", + "termcolor", + "unicode-width", +] + [[package]] name = "colorchoice" version = "1.0.5" @@ -295,6 +306,7 @@ dependencies = [ "clap_complete", "colored", "dirs", + "fs2", "futures-util", "git2", "glob", @@ -320,6 +332,7 @@ dependencies = [ "tree-sitter-python", "tree-sitter-rust", "tree-sitter-typescript", + "usearch", "which", ] @@ -382,6 +395,68 @@ dependencies = [ "typenum", ] +[[package]] +name = "cxx" +version = "1.0.198" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fe442a792c7c736eea18b32a7f8a3b63cf8aafabda6760042dc2fdeda456291" +dependencies = [ + "cc", + "cxx-build", + "cxxbridge-cmd", + "cxxbridge-flags", + "cxxbridge-macro", + "foldhash", + "link-cplusplus", +] + +[[package]] +name = "cxx-build" +version = "1.0.198" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3184a94384c663718698311a78a51ac00c484c10b4eeac06fb0a068c5f64fa2" +dependencies = [ + "cc", + "codespan-reporting", + "indexmap", + "proc-macro2", + "quote", + "scratch", + "syn 3.0.2", +] + +[[package]] +name = "cxxbridge-cmd" +version = "1.0.198" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0148d8fd1199329ddf1d157a5e134e51ceff37c6a7ddd38615c399d81cb05d8d" +dependencies = [ + "clap", + "codespan-reporting", + "indexmap", + "proc-macro2", + "quote", + "syn 3.0.2", +] + +[[package]] +name = "cxxbridge-flags" +version = "1.0.198" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52850339faed2eaadd24e286dc1d8268cc6f8a7bd9524d713adc9099566b4c89" + +[[package]] +name = "cxxbridge-macro" +version = "1.0.198" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c77c856545d886c9bd5215409ebb63b925e262135248b50c79e5a5f194ee47c" +dependencies = [ + "indexmap", + "proc-macro2", + "quote", + "syn 3.0.2", +] + [[package]] name = "difflib" version = "0.4.0" @@ -497,6 +572,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -506,6 +587,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "futures-channel" version = "0.3.33" @@ -1022,6 +1113,15 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "link-cplusplus" +version = "1.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f78c730aaa7d0b9336a299029ea49f9ee53b0ed06e9202e8cb7db9bae7b8c82" +dependencies = [ + "cc", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -1096,6 +1196,15 @@ dependencies = [ "autocfg", ] +[[package]] +name = "numkong" +version = "7.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2694b5274d67e839d7eb977012f531c6b31ea9b9995bd4e8902a2ae68dd9071e" +dependencies = [ + "cc", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -1467,6 +1576,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "scratch" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d68f2ec51b097e4c1a75b681a8bec621909b5e91f15bb7b840c4f2f7b01148b2" + [[package]] name = "serde" version = "1.0.229" @@ -1672,6 +1787,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + [[package]] name = "termtree" version = "0.5.1" @@ -2064,6 +2188,17 @@ dependencies = [ "serde", ] +[[package]] +name = "usearch" +version = "2.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1cd7f672d20412962c457b11c858c6c5aecb949808a5345a95e1d671112bcf72" +dependencies = [ + "cxx", + "cxx-build", + "numkong", +] + [[package]] name = "utf8_iter" version = "1.0.4" @@ -2237,6 +2372,22 @@ dependencies = [ "winsafe", ] +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + [[package]] name = "winapi-util" version = "0.1.11" @@ -2246,6 +2397,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-core" version = "0.62.2" diff --git a/Cargo.toml b/Cargo.toml index a7fdcff..e0f8efd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,6 +59,10 @@ chrono = { version = "0.4.44", features = ["serde"] } # Symbol index (v0.6 — Code Intelligence) rusqlite = { version = "0.31", features = ["bundled"] } +# Vector search (Phase 3 — Brain Mode) +usearch = "2" +fs2 = "0.4" + # Tree-sitter AST parsing (optional — enable with --features tree-sitter) tree-sitter = { version = "0.26", optional = true } tree-sitter-rust = { version = "0.24", optional = true } diff --git a/src/embed/mod.rs b/src/embed/mod.rs index 0059787..d5b61f3 100644 --- a/src/embed/mod.rs +++ b/src/embed/mod.rs @@ -13,9 +13,7 @@ //! Re-exports will be added to this module in Phase 3 when `cora brain` //! commands are wired up. -// Phase 1: public API not yet consumed by any command. -// Clippy dead_code warnings are expected and will resolve in Phase 3. -#![allow(dead_code)] +// Embedding engine — consumed by brain module (Phase 3). pub mod token_vocab; pub mod tokens; diff --git a/src/index/brain.rs b/src/index/brain.rs new file mode 100644 index 0000000..468cebb --- /dev/null +++ b/src/index/brain.rs @@ -0,0 +1,262 @@ +//! Brain Mode — hybrid search combining FTS5 + usearch vectors + graph proximity. +//! +//! RRF fusion (k=60) merges 3 signal sources into ranked results. +//! Pattern adopted from uteke `doc_search_hybrid()`. + +use crate::data_dir::codecora_home; +use crate::embed::tokens::embed_code; +use crate::index::symbols::SymbolQuery; +use crate::index::vector::{cosine_distance_to_similarity, CodeVectorIndex, DEFAULT_DIMS}; +use anyhow::{Context, Result}; +use rusqlite::Connection; +use std::collections::HashMap; + +/// RRF constant (standard value from Cormack et al. 2009). +const RRF_K: f32 = 60.0; + +/// A brain search result with provenance information. +#[derive(Debug, Clone, serde::Serialize)] +pub struct BrainResult { + pub symbol_id: i64, + pub name: String, + pub kind: String, + pub file: String, + pub line: i64, + pub signature: String, + /// Fused RRF score (higher = better). + pub score: f32, + /// Which signals contributed: fts, vector, graph. + pub signals: Vec, +} + +/// Path to the usearch vector index file. +fn vector_index_path() -> std::path::PathBuf { + crate::data_dir::cora_data_dir().join("cora_index.usearch") +} + +/// Embed all symbols for a project into the vector index. +/// +/// Reads symbols from SQLite, embeds via static token method, +/// stores in usearch. Call after `index_project`. +pub fn embed_project(conn: &Connection, project_id: i64) -> Result { + let vi_path = vector_index_path(); + let mut vi = + CodeVectorIndex::load_or_create(&vi_path, DEFAULT_DIMS).context("load vector index")?; + + let mut stmt = conn.prepare( + "SELECT id, name, kind, signature FROM symbols WHERE project_id = ?1", + )?; + let rows: Vec<(i64, String, String, String)> = stmt + .query_map(rusqlite::params![project_id], |row| { + Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)) + })? + .filter_map(|r| r.ok()) + .collect(); + + let mut count = 0; + for (sym_id, name, _kind, signature) in &rows { + let text = if signature.is_empty() || signature == name { + name.clone() + } else { + format!("{name} {signature}") + }; + + let embedding = embed_code(&text); + let vec: Vec = embedding.as_slice().iter().map(|&v| v as f32).collect(); + + vi.insert(*sym_id, &vec).context("insert symbol embedding")?; + count += 1; + } + + if vi.is_dirty() { + vi.save().context("save vector index")?; + } + + conn.execute( + "UPDATE projects SET embedding_tier = 'static', embedding_dims = ?1, \ + last_embedded_at = datetime('now') WHERE id = ?2", + rusqlite::params![DEFAULT_DIMS, project_id], + )?; + + tracing::info!("Embedded {count} symbols for project {project_id}"); + Ok(count) +} + +/// Hybrid brain search: FTS5 + usearch KNN + graph proximity → RRF fusion. +pub fn brain_search( + conn: &Connection, + project_id: i64, + query: &str, + limit: usize, +) -> Result> { + let limit = limit.min(50); + let fetch_limit = limit * 2; + + let fts_hits = fts5_search(conn, project_id, query, fetch_limit); + let vec_hits = vector_search(query, fetch_limit); + let graph_hits = graph_proximity_search(conn, project_id, &fts_hits, fetch_limit); + + // ── RRF Fusion ────────────────────────────────────────────────────── + let mut fused: HashMap)> = HashMap::new(); + + for (rank, (id, _score)) in fts_hits.iter().enumerate() { + let rrf = 1.0 / (RRF_K + (rank as f32 + 1.0)); + let entry = fused.entry(*id).or_insert((0.0, Vec::new())); + entry.0 += rrf; + entry.1.push("fts".into()); + } + + for (rank, (id, _sim)) in vec_hits.iter().enumerate() { + let rrf = 1.0 / (RRF_K + (rank as f32 + 1.0)); + let entry = fused.entry(*id).or_insert((0.0, Vec::new())); + entry.0 += rrf; + entry.1.push("vector".into()); + } + + for (rank, (id, _depth)) in graph_hits.iter().enumerate() { + let rrf = 1.0 / (RRF_K + (rank as f32 + 1.0)); + let entry = fused.entry(*id).or_insert((0.0, Vec::new())); + entry.0 += rrf; + entry.1.push("graph".into()); + } + + let mut ranked: Vec<_> = fused.into_iter().collect(); + ranked.sort_by(|a, b| { + b.1.0 + .partial_cmp(&a.1.0) + .unwrap_or(std::cmp::Ordering::Equal) + }); + ranked.truncate(limit); + + let results: Vec = ranked + .into_iter() + .filter_map(|(id, (score, mut signals))| { + get_symbol_by_id(conn, id).ok().map(|sym| { + signals.sort(); + signals.dedup(); + BrainResult { + symbol_id: id, + name: sym.0, + kind: sym.1, + file: sym.2, + line: sym.3, + signature: sym.4, + score, + signals, + } + }) + }) + .collect(); + + Ok(results) +} + +// ── Signal sources ─────────────────────────────────────────────────────── + +/// FTS5 keyword search → (symbol_id, rank_score) pairs. +fn fts5_search( + conn: &Connection, + project_id: i64, + query: &str, + limit: usize, +) -> Vec<(i64, f64)> { + let sq = SymbolQuery::text(query); + match crate::index::search(conn, project_id, &sq) { + Ok(results) => results + .into_iter() + .take(limit) + .map(|r| (r.symbol.id, r.score)) + .collect(), + Err(e) => { + tracing::warn!("FTS5 search error: {e}"); + Vec::new() + } + } +} + +/// usearch vector search → (symbol_id, cosine_similarity) pairs. +fn vector_search(query: &str, limit: usize) -> Vec<(i64, f32)> { + let vi_path = vector_index_path(); + if !vi_path.exists() { + return Vec::new(); + } + let vi = match CodeVectorIndex::load_or_create(&vi_path, DEFAULT_DIMS) { + Ok(v) => v, + Err(e) => { + tracing::warn!("vector index load error: {e}"); + return Vec::new(); + } + }; + + if vi.is_empty() { + return Vec::new(); + } + + let embedding = embed_code(query); + let vec: Vec = embedding.as_slice().iter().map(|&v| v as f32).collect(); + + vi.search(&vec, limit) + .into_iter() + .map(|(sym_id, dist)| (sym_id, cosine_distance_to_similarity(dist))) + .collect() +} + +/// Graph proximity from top FTS result → (symbol_id, proximity) pairs. +fn graph_proximity_search( + conn: &Connection, + project_id: i64, + fts_results: &[(i64, f64)], + limit: usize, +) -> Vec<(i64, f32)> { + if fts_results.is_empty() { + return Vec::new(); + } + + let top_id = fts_results[0].0; + let top_name: String = match conn.query_row( + "SELECT name FROM symbols WHERE id = ?1", + rusqlite::params![top_id], + |row| row.get(0), + ) { + Ok(n) => n, + Err(_) => return Vec::new(), + }; + + let Ok(mut stmt) = conn.prepare( + "SELECT DISTINCT s.id FROM symbols s + JOIN edges e ON (e.target = s.name OR e.source = s.name) + WHERE (e.source = ?1 OR e.target = ?1) AND e.project_id = ?2 + AND s.project_id = ?2 AND s.id != ?3 + LIMIT ?4", + ) else { + return Vec::new(); + }; + + let ids: Vec<(i64, f32)> = stmt + .query_map(rusqlite::params![top_name, project_id, top_id, limit], |row| { + row.get(0) + }) + .ok() + .map(|rows| { + rows.filter_map(|r| r.ok()) + .enumerate() + .map(|(i, id)| (id, 1.0 / (i as f32 + 2.0))) + .collect() + }) + .unwrap_or_default(); + + ids +} + +/// Fetch symbol row by ID → (name, kind, file, line, signature). +fn get_symbol_by_id( + conn: &Connection, + id: i64, +) -> Result<(String, String, String, i64, String)> { + conn.query_row( + "SELECT name, kind, file, line, signature FROM symbols WHERE id = ?1", + rusqlite::params![id], + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?, row.get(4)?)), + ) + .map_err(Into::into) +} diff --git a/src/index/mod.rs b/src/index/mod.rs index 8d8b1af..20d941b 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -10,6 +10,8 @@ mod extract; pub mod graph; pub mod schema; mod symbols; +pub mod brain; +pub mod vector; use std::collections::HashMap; use std::path::Path; @@ -227,6 +229,20 @@ fn index_project_with_id( stats.files_scanned, stats.files_indexed, stats.symbols_indexed, stats.errors ); + // Embed symbols into vector index for brain search + match brain::embed_project(conn, project_id) { + Ok(n) => { + stats.embedded_symbols = Some(n); + info!("Brain: embedded {n} symbols"); + } + Err(e) => { + if verbose { + eprintln!(" ⚠ Embedding failed (non-fatal): {e}"); + } + tracing::warn!("Embedding failed: {e}"); + } + } + Ok(stats) } @@ -350,6 +366,7 @@ pub struct IndexStats { pub files_skipped: usize, pub symbols_indexed: usize, pub errors: usize, + pub embedded_symbols: Option, } /// Summary of the current index state. diff --git a/src/index/schema.rs b/src/index/schema.rs index 5a69ae9..6c1c394 100644 --- a/src/index/schema.rs +++ b/src/index/schema.rs @@ -4,7 +4,7 @@ use rusqlite::Connection; /// Current schema version. #[allow(dead_code)] -const SCHEMA_VERSION: i32 = 3; +const SCHEMA_VERSION: i32 = 4; /// Run database migrations (creates tables if not exist). pub fn run_migrations(conn: &Connection) -> anyhow::Result<()> { @@ -31,6 +31,9 @@ pub fn run_migrations(conn: &Connection) -> anyhow::Result<()> { if current < 3 { migrate_v3(conn)?; } + if current < 4 { + migrate_v4(conn)?; + } Ok(()) } @@ -194,6 +197,21 @@ fn migrate_v3(conn: &Connection) -> anyhow::Result<()> { Ok(()) } +/// Migration v4: Embedding metadata on projects table. +fn migrate_v4(conn: &Connection) -> anyhow::Result<()> { + conn.execute_batch( + " + ALTER TABLE projects ADD COLUMN embedding_tier TEXT NOT NULL DEFAULT 'static'; + ALTER TABLE projects ADD COLUMN embedding_dims INTEGER NOT NULL DEFAULT 256; + ALTER TABLE projects ADD COLUMN last_embedded_at TEXT; + ", + )?; + + conn.execute("INSERT INTO schema_version (version) VALUES (4)", [])?; + + Ok(()) +} + /// Get or create a project entry by root path. /// /// Returns the project ID. diff --git a/src/index/vector.rs b/src/index/vector.rs new file mode 100644 index 0000000..ad29f95 --- /dev/null +++ b/src/index/vector.rs @@ -0,0 +1,384 @@ +//! Persistent vector index using usearch (HNSW) for symbol embeddings. +//! +//! Pattern copied from uteke-core/src/memory/vector.rs, simplified for cora-code: +//! - Keys are symbol database IDs (i64) instead of UUID strings +//! - Single purpose: code symbol semantic search + +use anyhow::{Context, Result}; +use fs2::FileExt; +use std::collections::HashMap; +use std::fs::File; +use std::path::{Path, PathBuf}; +use usearch::{Index, IndexOptions, MetricKind, ScalarKind}; + +const USEARCH_EXT: &str = "usearch"; + +/// Default embedding dimensions (nomic-embed-code static, 768d). +pub const DEFAULT_DIMS: usize = 256; + +/// Persistent HNSW vector index for code symbol embeddings. +/// +/// - Keys are symbol database IDs (i64) +/// - Cosine distance metric +/// - Disk persistence via buffer-based serialization (same as uteke) +/// - Cross-process safety via exclusive file lock +pub struct CodeVectorIndex { + index: Index, + /// Maps usearch integer key → symbol database ID. + key_to_symbol: HashMap, + /// Maps symbol database ID → usearch integer key. + symbol_to_key: HashMap, + next_key: u64, + path: Option, + dirty: bool, + _lock_file: Option, +} + +impl CodeVectorIndex { + /// Create a new empty in-memory index. + pub fn new(dims: usize) -> Result { + let index = create_usearch_index(dims)?; + Ok(Self { + index, + key_to_symbol: HashMap::new(), + symbol_to_key: HashMap::new(), + next_key: 0, + path: None, + dirty: false, + _lock_file: None, + }) + } + + /// Load from disk, or create empty if not exists. + /// Acquires exclusive file lock for cross-process safety. + pub fn load_or_create(path: &Path, dims: usize) -> Result { + if !path.exists() { + std::fs::write(path, []).context("create usearch file")?; + } + + let mut lock_file = acquire_file_lock(path)?; + + let mut idx = if lock_file.metadata().context("read file metadata")?.len() == 0 { + Self::new(dims)? + } else { + Self::load_from_file(&mut lock_file, path)? + }; + idx.path = Some(path.to_path_buf()); + idx._lock_file = Some(lock_file); + Ok(idx) + } + + fn load_from_file(file: &mut File, path: &Path) -> Result { + use std::io::{Read, Seek, SeekFrom}; + + file.seek(SeekFrom::Start(0)) + .context("seek usearch file")?; + let mut buffer = Vec::new(); + file.read_to_end(&mut buffer) + .context("read usearch file")?; + + let index = + Index::restore_from_buffer(&buffer).context("load usearch index from buffer")?; + + // Rebuild key mappings from sidecar + let mut key_to_symbol = HashMap::new(); + let mut symbol_to_key = HashMap::new(); + let mut next_key = 0u64; + + let mapping_path = path.with_extension("keys"); + if mapping_path.exists() { + let data = std::fs::read_to_string(&mapping_path) + .context("read key mapping file")?; + for line in data.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + if let Some((key_str, sym_id)) = line.split_once('\t') { + if let (Ok(key), Ok(sym)) = (key_str.parse::(), sym_id.parse::()) + { + key_to_symbol.insert(key, sym); + symbol_to_key.insert(sym, key); + next_key = next_key.max(key + 1); + } + } + } + } + + Ok(Self { + index, + key_to_symbol, + symbol_to_key, + next_key, + path: None, + dirty: false, + _lock_file: None, + }) + } + + /// Save index and key mappings to disk. + pub fn save(&mut self) -> Result<()> { + if let Some(ref path) = self.path { + let buf_len = self.index.serialized_length(); + let mut buffer = vec![0u8; buf_len]; + self.index + .save_to_buffer(&mut buffer) + .context("save usearch to buffer")?; + + let tmp_path = path.with_extension(format!("{USEARCH_EXT}.tmp")); + std::fs::write(&tmp_path, &buffer).context("write temp usearch index")?; + std::fs::rename(&tmp_path, path).context("rename temp usearch")?; + + // Save key mapping sidecar + let mapping_path = path.with_extension("keys"); + let mut lines = Vec::new(); + for (&key, &sym_id) in &self.key_to_symbol { + lines.push(format!("{key}\t{sym_id}")); + } + atomic_write(&mapping_path, lines.join("\n").as_bytes())?; + + self.dirty = false; + } + Ok(()) + } + + /// Insert a symbol embedding. If symbol ID exists, replaces it. + pub fn insert(&mut self, symbol_id: i64, embedding: &[f32]) -> Result<()> { + // Remove old entry if exists + if let Some(&old_key) = self.symbol_to_key.get(&symbol_id) { + self.key_to_symbol.remove(&old_key); + self.index + .remove(old_key) + .context("remove old usearch entry")?; + } + + let key = self.next_key; + self.next_key += 1; + self.key_to_symbol.insert(key, symbol_id); + self.symbol_to_key.insert(symbol_id, key); + + // Auto-reserve if at capacity + if self.index.size() >= self.index.capacity() { + let new_cap = (self.index.capacity() + 1024).max(1024); + self.index + .reserve(new_cap) + .context("reserve usearch capacity")?; + } + + self.index + .add(key, embedding) + .context("insert into usearch")?; + + self.dirty = true; + Ok(()) + } + + /// Remove a symbol by database ID. Incremental, no rebuild. + pub fn remove(&mut self, symbol_id: i64) -> bool { + if let Some(&key) = self.symbol_to_key.get(&symbol_id) { + self.key_to_symbol.remove(&key); + self.symbol_to_key.remove(&symbol_id); + if let Err(e) = self.index.remove(key) { + tracing::error!("Failed to remove from usearch: {e}"); + } + self.dirty = true; + true + } else { + false + } + } + + /// Search for k nearest neighbors. Returns (symbol_id, cosine_distance) pairs. + pub fn search(&self, query: &[f32], k: usize) -> Vec<(i64, f32)> { + if self.index.size() == 0 { + return Vec::new(); + } + let count = k.max(1); + let results = match self.index.search(query, count) { + Ok(r) => r, + Err(e) => { + tracing::error!("usearch search failed: {e}"); + return Vec::new(); + } + }; + + results + .keys + .iter() + .zip(results.distances.iter()) + .filter_map(|(key, dist)| { + self.key_to_symbol.get(key).map(|&id| (id, *dist)) + }) + .collect() + } + + /// Number of vectors in the index. + pub fn len(&self) -> usize { + self.index.size() + } + + /// Embedding dimensionality. + pub fn dims(&self) -> usize { + self.index.dimensions() + } + + /// Whether the index has unsaved changes. + pub fn is_dirty(&self) -> bool { + self.dirty + } + + /// Whether the index is empty. + pub fn is_empty(&self) -> bool { + self.index.size() == 0 + } +} + +fn create_usearch_index(dims: usize) -> Result { + let options = IndexOptions { + dimensions: dims, + metric: MetricKind::Cos, + quantization: ScalarKind::F32, + ..Default::default() + }; + Index::new(&options).context("create usearch index") +} + +/// Convert cosine distance (0..2) to cosine similarity (0..1). +pub fn cosine_distance_to_similarity(distance: f32) -> f32 { + (1.0 - distance).clamp(0.0, 1.0) +} + +fn acquire_file_lock(path: &Path) -> Result { + let file = File::options() + .read(true) + .write(true) + .open(path) + .with_context(|| format!("open usearch file for locking: {}", path.display()))?; + + if file.try_lock_exclusive().is_ok() { + tracing::debug!("usearch file lock acquired: {}", path.display()); + } else { + tracing::debug!("usearch file lock busy, waiting..."); + file.lock_exclusive() + .context("acquire exclusive file lock on usearch")?; + } + Ok(file) +} + +fn atomic_write(path: &std::path::Path, data: &[u8]) -> Result<()> { + let tmp_path = path.with_extension("keys.tmp"); + std::fs::write(&tmp_path, data).context("write temp key mapping")?; + std::fs::rename(&tmp_path, path).context("rename temp to final key mapping")?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::f32; + + fn make_unit_vec(dims: usize, idx: usize) -> Vec { + let mut v = vec![0.0f32; dims]; + if idx < dims { + v[idx] = 1.0; + } + v + } + + #[test] + fn test_empty_search() { + let idx = CodeVectorIndex::new(768).unwrap(); + assert!(idx.is_empty()); + let results = idx.search(&[0.0f32; 768], 5); + assert!(results.is_empty()); + } + + #[test] + fn test_insert_and_search() { + let mut idx = CodeVectorIndex::new(768).unwrap(); + + let v1 = make_unit_vec(768, 0); // unit vector along dim 0 + let v2 = make_unit_vec(768, 1); // unit vector along dim 1 + let mut v3 = vec![0.0f32; 768]; + v3[0] = 0.9; + v3[1] = 0.1; + let norm = v3.iter().map(|x| x * x).sum::().sqrt(); + v3.iter_mut().for_each(|x| *x /= norm); + + idx.insert(100, &v1).unwrap(); + idx.insert(200, &v2).unwrap(); + idx.insert(300, &v3).unwrap(); + assert_eq!(idx.len(), 3); + + // Query with v1 — should return v3 closest (similar direction), then v1 (exact) + let results = idx.search(&v1, 3); + assert_eq!(results.len(), 3); + // First result should be symbol 100 (exact match, dist ~0) + assert_eq!(results[0].0, 100); + // v3 is closer to v1 than v2 is + let d_v3 = results.iter().find(|(id, _)| *id == 300).map(|(_, d)| *d); + let d_v2 = results.iter().find(|(id, _)| *id == 200).map(|(_, d)| *d); + assert!(d_v3.unwrap() < d_v2.unwrap()); + } + + #[test] + fn test_replace_on_duplicate_insert() { + let mut idx = CodeVectorIndex::new(64).unwrap(); + + let v1 = make_unit_vec(64, 0); + let v2 = make_unit_vec(64, 1); + + idx.insert(42, &v1).unwrap(); + assert_eq!(idx.len(), 1); + + // Insert same symbol ID with different vector — should replace + idx.insert(42, &v2).unwrap(); + assert_eq!(idx.len(), 1); // still 1, not 2 + + let results = idx.search(&v2, 1); + assert_eq!(results[0].0, 42); + } + + #[test] + fn test_remove() { + let mut idx = CodeVectorIndex::new(64).unwrap(); + + idx.insert(1, &make_unit_vec(64, 0)).unwrap(); + idx.insert(2, &make_unit_vec(64, 1)).unwrap(); + assert_eq!(idx.len(), 2); + + assert!(idx.remove(1)); + assert_eq!(idx.len(), 1); + + let results = idx.search(&make_unit_vec(64, 0), 5); + assert!(results.iter().all(|(id, _)| *id != 1)); + } + + #[test] + fn test_save_and_load() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("test.usearch"); + + { + let mut idx = CodeVectorIndex::new(64).unwrap(); + idx.insert(10, &make_unit_vec(64, 0)).unwrap(); + idx.insert(20, &make_unit_vec(64, 1)).unwrap(); + idx.path = Some(path.clone()); + idx.save().unwrap(); + } + + let idx2 = CodeVectorIndex::load_or_create(&path, DEFAULT_DIMS).unwrap(); + assert_eq!(idx2.len(), 2); + + let results = idx2.search(&make_unit_vec(64, 0), 1); + assert_eq!(results[0].0, 10); + } + + #[test] + fn test_cosine_distance_to_similarity() { + assert!((cosine_distance_to_similarity(0.0) - 1.0).abs() < f32::EPSILON); + assert!((cosine_distance_to_similarity(1.0) - 0.0).abs() < f32::EPSILON); + assert!((cosine_distance_to_similarity(2.0) - 0.0).abs() < f32::EPSILON); + assert!((cosine_distance_to_similarity(0.5) - 0.5).abs() < f32::EPSILON); + } +} diff --git a/src/main.rs b/src/main.rs index 15df50c..d56ce74 100644 --- a/src/main.rs +++ b/src/main.rs @@ -187,6 +187,20 @@ enum Command { json: bool, }, + /// Hybrid search: FTS5 + vector + graph → RRF fusion + Brain { + /// Search query + query: Vec, + + /// Maximum results (default: 20) + #[clap(long, default_value = "20")] + limit: usize, + + /// Output as JSON + #[clap(long)] + json: bool, + }, + /// Find tests affected by changed files Affected { /// Changed files (space-separated). If empty, reads from git diff. @@ -930,6 +944,53 @@ async fn main() -> Result<()> { 0 } + Command::Brain { query, limit, json } => { + let query_str = query.join(" "); + if query_str.trim().is_empty() { + eprintln!("{}", "Usage: cora brain ".yellow()); + std::process::exit(1); + } + + let project_root = std::env::current_dir()?; + let db_path = crate::data_dir::graph_db_path(); + if !db_path.exists() { + eprintln!("{}", "No index found. Run `cora index` first.".yellow()); + std::process::exit(1); + } + let conn = index::open_global_index()?; + let project_id = index::ensure_project(&conn, &project_root)?; + + let results = index::brain::brain_search(&conn, project_id, &query_str, limit)?; + + if json { + println!("{}", serde_json::to_string_pretty(&results)?); + } else { + if results.is_empty() { + println!("{}", "No results found.".dimmed()); + } else { + println!( + "{} {} results for '{}'", + "🧠".magenta(), + results.len().to_string().bold(), + query_str + ); + println!("{}", "─".repeat(60).dimmed()); + for r in &results { + let signals = r.signals.join(","); + let score = format!("{:.4}", r.score); + println!( + " {:<40} {:<15} {} {}", + format!("{} ({})", r.name, r.file).green(), + format!("[{}]", r.kind).dimmed(), + format!("score={}", score).cyan(), + format!("signals={}", signals).dimmed(), + ); + } + } + } + 0 + } + Command::Affected { files, stdin, diff --git a/src/mcp/tools.rs b/src/mcp/tools.rs index 9290e5b..7918fae 100644 --- a/src/mcp/tools.rs +++ b/src/mcp/tools.rs @@ -181,6 +181,19 @@ pub fn list_tools() -> Vec { "required": ["query"] }), }, + // Brain Mode (Phase 3) + Tool { + name: "cora.brain_search".to_string(), + description: "Hybrid code search: FTS5 + vector embeddings + graph proximity → RRF fusion. Better than plain text search for semantic code queries.".to_string(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "query": { "type": "string", "description": "Search query (code concept, symbol name, or description)" }, + "limit": { "type": "integer", "description": "Max results (default: 20)" } + }, + "required": ["query"] + }), + }, ] } @@ -206,6 +219,8 @@ pub fn handle_tool_call(name: &str, params: &serde_json::Value) -> ToolResult { // Context Enrichment (Phase 3) "cora.get_project_info" => handle_get_project_info(), "cora.get_memory" => handle_get_memory(params), + // Brain Mode (Phase 3) + "cora.brain_search" => handle_brain_search(params), _ => ToolResult::error(format!("Unknown tool: {name}")), } } @@ -875,6 +890,45 @@ fn handle_get_memory(params: &serde_json::Value) -> ToolResult { ToolResult::text(serde_json::to_string_pretty(&json).unwrap_or_default()) } +// ─── Brain Mode Handlers (Phase 3) ─── + +fn handle_brain_search(params: &serde_json::Value) -> ToolResult { + let query = match params.get("query").and_then(|v| v.as_str()) { + Some(q) => q, + None => return ToolResult::error("Missing required parameter: query"), + }; + let limit = params.get("limit").and_then(|v| v.as_u64()).unwrap_or(20) as usize; + + let (conn, project_id) = match open_index_db() { + Ok((c, pid)) => (c, pid), + Err(e) => return ToolResult::error(e.to_string()), + }; + + match crate::index::brain::brain_search(&conn, project_id, query, limit) { + Ok(results) => { + if results.is_empty() { + return ToolResult::text(format!("No results found for '{query}'.")); + } + let json: Vec = results + .iter() + .map(|r| { + serde_json::json!({ + "name": r.name, + "kind": r.kind, + "file": r.file, + "line": r.line, + "signature": r.signature, + "score": r.score, + "signals": r.signals, + }) + }) + .collect(); + ToolResult::text(serde_json::to_string_pretty(&json).unwrap_or_default()) + } + Err(e) => ToolResult::error(format!("Brain search failed: {e}")), + } +} + /// Load project config safely (no API keys exposed). fn load_project_config() -> anyhow::Result { let mut config = crate::config::schema::Config::default(); @@ -1055,7 +1109,7 @@ mod tests { #[test] fn total_tool_count() { let tools = list_tools(); - // Phase 1 (5 existing) + Phase 1 code intel (5) + Phase 2 (2) + Phase 3 (2) = 14 - assert_eq!(tools.len(), 14); + // Phase 1 (5) + code intel (5) + Phase 2 (2) + Phase 3 (3) = 15 + assert_eq!(tools.len(), 15); } } From f28845a65349abd64bd05338b3741a66d4da3860 Mon Sep 17 00:00:00 2001 From: ajianaz Date: Thu, 23 Jul 2026 06:03:00 +0700 Subject: [PATCH 2/2] fix: resolve CI dead_code errors for -D warnings clippy - Add #![allow(dead_code)] to vector.rs (Phase 4+ API) and token_vocab.rs (Phase 5) - Add #[allow(dead_code)] on TokenEmbedding::into_vec - Add #[cfg(test)] on cosine_similarity (test-only) - Remove unused codecora_home import from brain.rs --- src/embed/token_vocab.rs | 3 +++ src/embed/tokens.rs | 2 ++ src/index/brain.rs | 40 ++++++++++++++++++++-------------------- src/index/mod.rs | 4 ++-- src/index/vector.rs | 19 ++++++++----------- 5 files changed, 35 insertions(+), 33 deletions(-) diff --git a/src/embed/token_vocab.rs b/src/embed/token_vocab.rs index cd676bb..4080275 100644 --- a/src/embed/token_vocab.rs +++ b/src/embed/token_vocab.rs @@ -1,3 +1,6 @@ +//! Pre-trained token vocabulary — reserved for Phase 5 (Voyage-4-Nano ONNX). +#![allow(dead_code)] + //! Pre-trained token vocabulary from nomic-embed-code. //! //! Provides real 768-dimensional code embeddings distilled from diff --git a/src/embed/tokens.rs b/src/embed/tokens.rs index 9730a15..382ee23 100644 --- a/src/embed/tokens.rs +++ b/src/embed/tokens.rs @@ -51,6 +51,7 @@ impl TokenEmbedding { } /// Consumes self and returns the underlying vector. + #[allow(dead_code)] pub fn into_vec(self) -> Vec { self.vec } @@ -278,6 +279,7 @@ pub fn embed_code(code: &str) -> TokenEmbedding { /// /// Because vectors are pre-normalised by [`embed`], this is just the dot /// product. Returns a value in `[-1, 1]`. +#[cfg(test)] pub fn cosine_similarity(a: &TokenEmbedding, b: &TokenEmbedding) -> f64 { a.vec.iter().zip(b.vec.iter()).map(|(x, y)| x * y).sum() } diff --git a/src/index/brain.rs b/src/index/brain.rs index 468cebb..639e07a 100644 --- a/src/index/brain.rs +++ b/src/index/brain.rs @@ -3,10 +3,9 @@ //! RRF fusion (k=60) merges 3 signal sources into ranked results. //! Pattern adopted from uteke `doc_search_hybrid()`. -use crate::data_dir::codecora_home; use crate::embed::tokens::embed_code; use crate::index::symbols::SymbolQuery; -use crate::index::vector::{cosine_distance_to_similarity, CodeVectorIndex, DEFAULT_DIMS}; +use crate::index::vector::{CodeVectorIndex, DEFAULT_DIMS, cosine_distance_to_similarity}; use anyhow::{Context, Result}; use rusqlite::Connection; use std::collections::HashMap; @@ -43,9 +42,8 @@ pub fn embed_project(conn: &Connection, project_id: i64) -> Result { let mut vi = CodeVectorIndex::load_or_create(&vi_path, DEFAULT_DIMS).context("load vector index")?; - let mut stmt = conn.prepare( - "SELECT id, name, kind, signature FROM symbols WHERE project_id = ?1", - )?; + let mut stmt = + conn.prepare("SELECT id, name, kind, signature FROM symbols WHERE project_id = ?1")?; let rows: Vec<(i64, String, String, String)> = stmt .query_map(rusqlite::params![project_id], |row| { Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)) @@ -64,7 +62,8 @@ pub fn embed_project(conn: &Connection, project_id: i64) -> Result { let embedding = embed_code(&text); let vec: Vec = embedding.as_slice().iter().map(|&v| v as f32).collect(); - vi.insert(*sym_id, &vec).context("insert symbol embedding")?; + vi.insert(*sym_id, &vec) + .context("insert symbol embedding")?; count += 1; } @@ -154,12 +153,7 @@ pub fn brain_search( // ── Signal sources ─────────────────────────────────────────────────────── /// FTS5 keyword search → (symbol_id, rank_score) pairs. -fn fts5_search( - conn: &Connection, - project_id: i64, - query: &str, - limit: usize, -) -> Vec<(i64, f64)> { +fn fts5_search(conn: &Connection, project_id: i64, query: &str, limit: usize) -> Vec<(i64, f64)> { let sq = SymbolQuery::text(query); match crate::index::search(conn, project_id, &sq) { Ok(results) => results @@ -233,9 +227,10 @@ fn graph_proximity_search( }; let ids: Vec<(i64, f32)> = stmt - .query_map(rusqlite::params![top_name, project_id, top_id, limit], |row| { - row.get(0) - }) + .query_map( + rusqlite::params![top_name, project_id, top_id, limit], + |row| row.get(0), + ) .ok() .map(|rows| { rows.filter_map(|r| r.ok()) @@ -249,14 +244,19 @@ fn graph_proximity_search( } /// Fetch symbol row by ID → (name, kind, file, line, signature). -fn get_symbol_by_id( - conn: &Connection, - id: i64, -) -> Result<(String, String, String, i64, String)> { +fn get_symbol_by_id(conn: &Connection, id: i64) -> Result<(String, String, String, i64, String)> { conn.query_row( "SELECT name, kind, file, line, signature FROM symbols WHERE id = ?1", rusqlite::params![id], - |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?, row.get(4)?)), + |row| { + Ok(( + row.get(0)?, + row.get(1)?, + row.get(2)?, + row.get(3)?, + row.get(4)?, + )) + }, ) .map_err(Into::into) } diff --git a/src/index/mod.rs b/src/index/mod.rs index 20d941b..bfedbee 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -6,11 +6,11 @@ #[cfg(feature = "tree-sitter")] mod ast; +pub mod brain; mod extract; pub mod graph; pub mod schema; mod symbols; -pub mod brain; pub mod vector; use std::collections::HashMap; @@ -366,7 +366,7 @@ pub struct IndexStats { pub files_skipped: usize, pub symbols_indexed: usize, pub errors: usize, - pub embedded_symbols: Option, + pub embedded_symbols: Option, } /// Summary of the current index state. diff --git a/src/index/vector.rs b/src/index/vector.rs index ad29f95..122ff67 100644 --- a/src/index/vector.rs +++ b/src/index/vector.rs @@ -4,6 +4,9 @@ //! - Keys are symbol database IDs (i64) instead of UUID strings //! - Single purpose: code symbol semantic search +// Public API reserved for Phase 4+ (remove, dims, etc.). +#![allow(dead_code)] + use anyhow::{Context, Result}; use fs2::FileExt; use std::collections::HashMap; @@ -71,11 +74,9 @@ impl CodeVectorIndex { fn load_from_file(file: &mut File, path: &Path) -> Result { use std::io::{Read, Seek, SeekFrom}; - file.seek(SeekFrom::Start(0)) - .context("seek usearch file")?; + file.seek(SeekFrom::Start(0)).context("seek usearch file")?; let mut buffer = Vec::new(); - file.read_to_end(&mut buffer) - .context("read usearch file")?; + file.read_to_end(&mut buffer).context("read usearch file")?; let index = Index::restore_from_buffer(&buffer).context("load usearch index from buffer")?; @@ -87,16 +88,14 @@ impl CodeVectorIndex { let mapping_path = path.with_extension("keys"); if mapping_path.exists() { - let data = std::fs::read_to_string(&mapping_path) - .context("read key mapping file")?; + let data = std::fs::read_to_string(&mapping_path).context("read key mapping file")?; for line in data.lines() { let line = line.trim(); if line.is_empty() { continue; } if let Some((key_str, sym_id)) = line.split_once('\t') { - if let (Ok(key), Ok(sym)) = (key_str.parse::(), sym_id.parse::()) - { + if let (Ok(key), Ok(sym)) = (key_str.parse::(), sym_id.parse::()) { key_to_symbol.insert(key, sym); symbol_to_key.insert(sym, key); next_key = next_key.max(key + 1); @@ -206,9 +205,7 @@ impl CodeVectorIndex { .keys .iter() .zip(results.distances.iter()) - .filter_map(|(key, dist)| { - self.key_to_symbol.get(key).map(|&id| (id, *dist)) - }) + .filter_map(|(key, dist)| self.key_to_symbol.get(key).map(|&id| (id, *dist))) .collect() }