Implement global codebase vector indexing

This commit is contained in:
Riz Ashraf committed 2026-09-30 09:45:27 +01:00
1 parent ec0472155f
commit 181fc096a1
4 files changed
+208

No files matched your search

+134
View File
@@ -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);
}
}
}
+5
View File
@@ -12,6 +12,7 @@ mod models;
mod router;
mod search;
pub mod embedding;
pub mod indexer;
mod state;
mod store;
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;
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(ttl_sweeper_worker(Arc::clone(&state)));
crate::clipboard_watcher::spawn_watcher(Arc::clone(&state));