diff --git a/server/src/embedding.rs b/server/src/embedding.rs index 414e7c3..8edfd20 100644 --- a/server/src/embedding.rs +++ b/server/src/embedding.rs @@ -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) -> Result>, 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())? +} diff --git a/server/src/handlers/ast.rs b/server/src/handlers/ast.rs index 8a4646d..c87fdce 100644 --- a/server/src/handlers/ast.rs +++ b/server/src/handlers/ast.rs @@ -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}; diff --git a/server/src/handlers/workspaces.rs b/server/src/handlers/workspaces.rs index 1fe7378..d3205ad 100644 --- a/server/src/handlers/workspaces.rs +++ b/server/src/handlers/workspaces.rs @@ -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::())); + } + + 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::())); + results.push((sim, meta.0, meta.1)); } } diff --git a/server/src/test_clip_write.rs b/server/src/test_clip_write.rs index cf90fb1..1498ff6 100644 --- a/server/src/test_clip_write.rs +++ b/server/src/test_clip_write.rs @@ -1,5 +1,5 @@ use arboard::{Clipboard, ImageData}; -use image::GenericImageView; + use std::borrow::Cow; fn main() {