perf(mcp): Batch semantic code search embeddings to resolve CPU bottleneck and fix linting warnings
This commit is contained in:
1 parent
185c3c999e
commit
ec0472155f
4 files changed
+25
-9
No files matched your search
@@ -40,3 +40,14 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
||||
dot_product / (norm_a * norm_b)
|
||||
}
|
||||
}
|
||||
pub async fn generate_embeddings_async(texts: Vec<String>) -> Result<Vec<Vec<f32>>, String> {
|
||||
if texts.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let model_mutex = get_embedding_model()?;
|
||||
let mut model = model_mutex.lock().map_err(|e| e.to_string())?;
|
||||
let embeddings = model.embed(texts, None).map_err(|e| e.to_string())?;
|
||||
Ok(embeddings)
|
||||
}).await.map_err(|e| e.to_string())?
|
||||
}
|
||||
@@ -2,7 +2,7 @@ use crate::router::McpTool;
|
||||
use crate::state::MemoryState;
|
||||
use crate::tools::ReadFileSkeletonTool;
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
use tree_sitter::{Parser, Node};
|
||||
|
||||
|
||||
@@ -596,7 +596,7 @@ impl McpTool for ReadDirectoryArchitectureHandler {
|
||||
}
|
||||
}
|
||||
use crate::tools::SemanticCodeSearchTool;
|
||||
use crate::embedding::{generate_embedding_async, cosine_similarity};
|
||||
use crate::embedding::{generate_embedding_async, generate_embeddings_async, cosine_similarity};
|
||||
|
||||
pub struct SemanticCodeSearchHandler;
|
||||
|
||||
@@ -620,21 +620,26 @@ impl McpTool for SemanticCodeSearchHandler {
|
||||
|
||||
// For MVP, we search across snippets dynamically. A true background codebase indexer would be a separate subsystem.
|
||||
let mut results = Vec::new();
|
||||
let mut texts_to_embed = Vec::new();
|
||||
let mut metadata = Vec::new();
|
||||
|
||||
let snippets = state.snippets.read_with(|snips| snips.clone());
|
||||
for snippet in snippets {
|
||||
let combined = format!("{} {} {}", snippet.name, snippet.description, snippet.code);
|
||||
if let Ok(emb) = generate_embedding_async(combined).await {
|
||||
let sim = cosine_similarity(&query_emb, &emb);
|
||||
results.push((sim, snippet.name, snippet.description));
|
||||
}
|
||||
texts_to_embed.push(combined);
|
||||
metadata.push((snippet.name, snippet.description));
|
||||
}
|
||||
|
||||
let sticky = state.sticky.read_with(|s| s.clone());
|
||||
for note in sticky {
|
||||
if let Ok(emb) = generate_embedding_async(note.content.clone()).await {
|
||||
texts_to_embed.push(note.content.clone());
|
||||
metadata.push(("StickyNote".to_string(), note.content.chars().take(200).collect::<String>()));
|
||||
}
|
||||
|
||||
if let Ok(embeddings) = generate_embeddings_async(texts_to_embed).await {
|
||||
for (emb, meta) in embeddings.into_iter().zip(metadata) {
|
||||
let sim = cosine_similarity(&query_emb, &emb);
|
||||
results.push((sim, "StickyNote".to_string(), note.content.chars().take(200).collect::<String>()));
|
||||
results.push((sim, meta.0, meta.1));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use arboard::{Clipboard, ImageData};
|
||||
use image::GenericImageView;
|
||||
|
||||
use std::borrow::Cow;
|
||||
|
||||
fn main() {
|
||||
|
||||
Reference in new issue
Block a user