Implement global codebase vector indexing
This commit is contained in:
1 parent
ec0472155f
commit
181fc096a1
4 files changed
+208
No files matched your search
Generated
+68
@@ -409,6 +409,16 @@ dependencies = [
|
|||||||
"syn 3.0.6",
|
"syn 3.0.6",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "bstr"
|
||||||
|
version = "1.13.1"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "6bb31b46c14244e20ee9984b11bf5c992b91fb6939fea616e3512c8baecdbe5f"
|
||||||
|
dependencies = [
|
||||||
|
"memchr",
|
||||||
|
"serde_core",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "built"
|
name = "built"
|
||||||
version = "0.8.1"
|
version = "0.8.1"
|
||||||
@@ -1403,6 +1413,19 @@ dependencies = [
|
|||||||
"url",
|
"url",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "globset"
|
||||||
|
version = "0.4.20"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "07c34a9410465b45bd9787443bc7370f37735bad04b0f0cd57ff1a3186c98988"
|
||||||
|
dependencies = [
|
||||||
|
"aho-corasick",
|
||||||
|
"bstr",
|
||||||
|
"log",
|
||||||
|
"regex-automata",
|
||||||
|
"regex-syntax",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "h2"
|
name = "h2"
|
||||||
version = "0.4.19"
|
version = "0.4.19"
|
||||||
@@ -1763,6 +1786,22 @@ dependencies = [
|
|||||||
"icu_properties",
|
"icu_properties",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "ignore"
|
||||||
|
version = "0.4.33"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "00b69833ed729dc5aa7d19541d96d6cf8e9137194207a04916d658e43168402f"
|
||||||
|
dependencies = [
|
||||||
|
"crossbeam-deque",
|
||||||
|
"globset",
|
||||||
|
"log",
|
||||||
|
"memchr",
|
||||||
|
"regex-automata",
|
||||||
|
"same-file",
|
||||||
|
"walkdir",
|
||||||
|
"winapi-util",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "image"
|
name = "image"
|
||||||
version = "0.25.10"
|
version = "0.25.10"
|
||||||
@@ -2114,6 +2153,7 @@ dependencies = [
|
|||||||
"fastembed",
|
"fastembed",
|
||||||
"futures-util",
|
"futures-util",
|
||||||
"git2",
|
"git2",
|
||||||
|
"ignore",
|
||||||
"image",
|
"image",
|
||||||
"redb",
|
"redb",
|
||||||
"reqwest 0.13.5",
|
"reqwest 0.13.5",
|
||||||
@@ -3356,6 +3396,15 @@ dependencies = [
|
|||||||
"tempfile",
|
"tempfile",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "same-file"
|
||||||
|
version = "1.0.6"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
|
||||||
|
dependencies = [
|
||||||
|
"winapi-util",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "schannel"
|
name = "schannel"
|
||||||
version = "0.1.29"
|
version = "0.1.29"
|
||||||
@@ -4521,6 +4570,16 @@ version = "0.9.5"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
|
checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "walkdir"
|
||||||
|
version = "2.5.0"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b"
|
||||||
|
dependencies = [
|
||||||
|
"same-file",
|
||||||
|
"winapi-util",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "want"
|
name = "want"
|
||||||
version = "0.3.1"
|
version = "0.3.1"
|
||||||
@@ -4686,6 +4745,15 @@ version = "0.4.0"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
|
checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "winapi-util"
|
||||||
|
version = "0.1.11"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
|
||||||
|
dependencies = [
|
||||||
|
"windows-sys 0.61.2",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "winapi-x86_64-pc-windows-gnu"
|
name = "winapi-x86_64-pc-windows-gnu"
|
||||||
version = "0.4.0"
|
version = "0.4.0"
|
||||||
|
|||||||
@@ -37,6 +37,7 @@ tree-sitter-java = "0.23.5"
|
|||||||
tree-sitter-c = "0.24.2"
|
tree-sitter-c = "0.24.2"
|
||||||
tree-sitter-cpp = "0.23.4"
|
tree-sitter-cpp = "0.23.4"
|
||||||
tree-sitter-go = "0.25.0"
|
tree-sitter-go = "0.25.0"
|
||||||
|
ignore = "0.4.33"
|
||||||
|
|
||||||
[build-dependencies]
|
[build-dependencies]
|
||||||
chrono = "0.4.45"
|
chrono = "0.4.45"
|
||||||
|
|||||||
@@ -0,0 +1,134 @@
|
|||||||
|
use crate::state::MemoryState;
|
||||||
|
use crate::models::Snippet;
|
||||||
|
use crate::embedding::generate_embeddings_async;
|
||||||
|
use ignore::WalkBuilder;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tree_sitter::{Node, Parser};
|
||||||
|
|
||||||
|
pub async fn start_background_indexer(state: Arc<MemoryState>) {
|
||||||
|
// Determine workspace root. Since mcp-memory is usually run from its project root,
|
||||||
|
// current_dir is a good default for the git repository.
|
||||||
|
let workspace_root = std::env::current_dir().unwrap_or_else(|_| state.base_dir.clone());
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
tracing::info!("Starting background indexer in {:?}", workspace_root);
|
||||||
|
|
||||||
|
let walker = WalkBuilder::new(&workspace_root)
|
||||||
|
.hidden(true)
|
||||||
|
.git_ignore(true)
|
||||||
|
.build();
|
||||||
|
|
||||||
|
let mut files_to_process = Vec::new();
|
||||||
|
for result in walker {
|
||||||
|
match result {
|
||||||
|
Ok(entry) => {
|
||||||
|
if entry.file_type().map_or(false, |ft| ft.is_file()) {
|
||||||
|
let path = entry.path().to_path_buf();
|
||||||
|
let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
|
||||||
|
if ["rs", "ts", "js", "jsx", "tsx", "py", "java", "c", "cpp", "go"].contains(&ext) {
|
||||||
|
files_to_process.push(path);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!("Error walking directory: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let idx = state.get_search_index();
|
||||||
|
|
||||||
|
for file_path in files_to_process {
|
||||||
|
if let Ok(content) = std::fs::read_to_string(&file_path) {
|
||||||
|
let ext = file_path.extension().and_then(|e| e.to_str()).unwrap_or("");
|
||||||
|
let language = match ext {
|
||||||
|
"rs" => tree_sitter_rust::LANGUAGE,
|
||||||
|
"ts" | "tsx" | "js" | "jsx" => tree_sitter_typescript::LANGUAGE_TYPESCRIPT,
|
||||||
|
"py" => tree_sitter_python::LANGUAGE,
|
||||||
|
"java" => tree_sitter_java::LANGUAGE,
|
||||||
|
"c" | "h" => tree_sitter_c::LANGUAGE,
|
||||||
|
"cpp" | "cc" | "cxx" | "hpp" | "hxx" => tree_sitter_cpp::LANGUAGE,
|
||||||
|
"go" => tree_sitter_go::LANGUAGE,
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut parser = Parser::new();
|
||||||
|
if parser.set_language(&language.into()).is_err() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(tree) = parser.parse(&content, None) {
|
||||||
|
let mut chunks = Vec::new();
|
||||||
|
extract_chunks(tree.root_node(), &content, &mut chunks, ext);
|
||||||
|
|
||||||
|
for (name, code, desc) in chunks {
|
||||||
|
// Generate embedding
|
||||||
|
if let Ok(mut emb) = generate_embeddings_async(vec![code.clone()]).await {
|
||||||
|
let embedding = emb.pop();
|
||||||
|
|
||||||
|
let file_name = file_path.file_name().unwrap_or_default().to_string_lossy();
|
||||||
|
let snippet_name = format!("{}:{}", file_name, name);
|
||||||
|
|
||||||
|
let snippet = Snippet {
|
||||||
|
name: snippet_name.to_string(),
|
||||||
|
language: ext.to_string(),
|
||||||
|
code: code.clone(),
|
||||||
|
description: format!("{} in {}", desc, file_path.display()),
|
||||||
|
updated_at: std::time::SystemTime::now()
|
||||||
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs(),
|
||||||
|
embedding,
|
||||||
|
};
|
||||||
|
|
||||||
|
state.snippets.modify(|snippets| {
|
||||||
|
// Prevent duplicates if already indexed
|
||||||
|
if !snippets.iter().any(|s| s.name == snippet.name) {
|
||||||
|
snippets.push(snippet.clone());
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let _ = idx.index_snippet(&snippet).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
tracing::info!("Background indexing completed.");
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
fn extract_chunks(node: Node, code: &str, chunks: &mut Vec<(String, String, String)>, ext: &str) {
|
||||||
|
let kind = node.kind();
|
||||||
|
let is_structural = matches!(
|
||||||
|
kind,
|
||||||
|
"function_item"
|
||||||
|
| "function_declaration"
|
||||||
|
| "function_definition"
|
||||||
|
| "method_definition"
|
||||||
|
| "struct_item"
|
||||||
|
| "class_declaration"
|
||||||
|
);
|
||||||
|
|
||||||
|
if is_structural {
|
||||||
|
let code_text = node.utf8_text(code.as_bytes()).unwrap_or("").to_string();
|
||||||
|
|
||||||
|
let mut name = "unknown".to_string();
|
||||||
|
let mut cursor = node.walk();
|
||||||
|
for child in node.children(&mut cursor) {
|
||||||
|
let child_kind = child.kind();
|
||||||
|
if child_kind == "identifier" || child_kind == "name" || child_kind == "type_identifier" {
|
||||||
|
name = child.utf8_text(code.as_bytes()).unwrap_or("unknown").to_string();
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let desc = format!("{} AST node", kind);
|
||||||
|
chunks.push((name, code_text, desc));
|
||||||
|
} else {
|
||||||
|
let mut cursor = node.walk();
|
||||||
|
for child in node.named_children(&mut cursor) {
|
||||||
|
extract_chunks(child, code, chunks, ext);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -12,6 +12,7 @@ mod models;
|
|||||||
mod router;
|
mod router;
|
||||||
mod search;
|
mod search;
|
||||||
pub mod embedding;
|
pub mod embedding;
|
||||||
|
pub mod indexer;
|
||||||
mod state;
|
mod state;
|
||||||
mod store;
|
mod store;
|
||||||
mod tools;
|
mod tools;
|
||||||
@@ -131,6 +132,10 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
|
|||||||
state_for_index.rebuild_index().await;
|
state_for_index.rebuild_index().await;
|
||||||
tracing::info!("Index rebuild complete.");
|
tracing::info!("Index rebuild complete.");
|
||||||
});
|
});
|
||||||
|
|
||||||
|
// Start the global codebase indexer
|
||||||
|
crate::indexer::start_background_indexer(Arc::clone(&state)).await;
|
||||||
|
|
||||||
tokio::spawn(index_committer_worker(Arc::clone(&state)));
|
tokio::spawn(index_committer_worker(Arc::clone(&state)));
|
||||||
tokio::spawn(ttl_sweeper_worker(Arc::clone(&state)));
|
tokio::spawn(ttl_sweeper_worker(Arc::clone(&state)));
|
||||||
crate::clipboard_watcher::spawn_watcher(Arc::clone(&state));
|
crate::clipboard_watcher::spawn_watcher(Arc::clone(&state));
|
||||||
|
|||||||
Reference in new issue
Block a user