From 3add6c3d3174a9c8e4abe36f145a9785c4c9fd0b Mon Sep 17 00:00:00 2001 From: Riz Ashraf Date: Wed, 30 Sep 2026 14:37:41 +0100 Subject: [PATCH] Refactor vector search to utilize Qdrant VectorDB --- server/Cargo.toml | 2 +- server/src/handlers/workspaces.rs | 63 +++++++++++++++++++------------ server/src/main.rs | 13 +++++++ server/src/vector_db.rs | 52 ++++++++++++++++++++++--- 4 files changed, 99 insertions(+), 31 deletions(-) diff --git a/server/Cargo.toml b/server/Cargo.toml index 2c66ff3..b558824 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -18,7 +18,7 @@ tantivy = "0.26.1" tokio = { version = "1.53.1", features = ["full"] } tracing = "0.1.44" tracing-subscriber = "0.3.23" -uuid = { version = "1.26.0", features = ["v4"] } +uuid = { version = "1.26.0", features = ["v4", "v5"] } tracing-appender = "0.2.5" rmcp = { version = "3.4.0", features = ["server"] } thiserror = "2.0.20" diff --git a/server/src/handlers/workspaces.rs b/server/src/handlers/workspaces.rs index d3205ad..0168f01 100644 --- a/server/src/handlers/workspaces.rs +++ b/server/src/handlers/workspaces.rs @@ -618,41 +618,54 @@ impl McpTool for SemanticCodeSearchHandler { let query_emb = generate_embedding_async(tool_args.query.clone()).await?; - // 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); - texts_to_embed.push(combined); - metadata.push((snippet.name, snippet.description)); - } - - let sticky = state.sticky.read_with(|s| s.clone()); - for note in sticky { - 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, meta.0, meta.1)); + // Search using VectorDB if available + let mut vdb_search = false; + if let Some(vdb) = &*state.vector_db.read().await { + vdb_search = true; + if let Ok(search_results) = vdb.search(query_emb.clone(), 5).await { + for res in search_results { + results.push((res.score, res.id, res.text)); + } } } - results.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); + // Fallback to manual loop if VectorDB is not initialized + if !vdb_search { + 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); + texts_to_embed.push(combined); + metadata.push((snippet.name, snippet.description)); + } + + let sticky = state.sticky.read_with(|s| s.clone()); + for note in sticky { + 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, meta.0, meta.1)); + } + } + + results.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); + results.truncate(5); + } - let top_results: Vec<_> = results.into_iter().take(5).collect(); - - if top_results.is_empty() { + if results.is_empty() { return Ok(format!("No semantic matches found for query: {}", tool_args.query)); } let mut out = format!("Semantic Search Results for '{}':\n", tool_args.query); - for (score, title, desc) in top_results { + for (score, title, desc) in results { out.push_str(&format!("- [{:.2}] {}: {}\n", score, title, desc)); } diff --git a/server/src/main.rs b/server/src/main.rs index bfc240e..077bd51 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -13,6 +13,7 @@ mod router; mod search; pub mod embedding; pub mod indexer; +pub mod vector_db; mod state; mod store; mod tools; @@ -545,6 +546,18 @@ fn main() -> Result<(), Box> { rt.block_on(async { let state = Arc::new(MemoryState::new(&base.to_string_lossy())); + + // Initialize Qdrant VectorDB (default local URL) + match crate::vector_db::VectorDB::new("http://localhost:6334", "mcp_memory").await { + Ok(vdb) => { + tracing::info!("Successfully connected to Qdrant vector database"); + *state.vector_db.write().await = Some(vdb); + } + Err(e) => { + tracing::warn!("Failed to initialize Qdrant vector database: {}. Vector search will fallback to manual embedding loop. (Is Qdrant running on localhost:6334?)", e); + } + } + if let Err(e) = run_server(state).await { tracing::error!("Server error: {}", e); } diff --git a/server/src/vector_db.rs b/server/src/vector_db.rs index b950a81..eee39b2 100644 --- a/server/src/vector_db.rs +++ b/server/src/vector_db.rs @@ -3,7 +3,7 @@ use qdrant_client::Qdrant; use std::sync::Arc; use std::error::Error; use uuid::Uuid; -use log::{info, error}; +use tracing::{info, error}; use serde::{Deserialize, Serialize}; #[derive(Clone)] @@ -69,10 +69,10 @@ impl VectorDB { } }; - let mut payload = std::collections::HashMap::new(); - payload.insert("doc_type".to_string(), doc_type.into()); - payload.insert("text".to_string(), text.into()); - payload.insert("original_id".to_string(), id.into()); + let mut payload: std::collections::HashMap = std::collections::HashMap::new(); + payload.insert("doc_type".to_string(), serde_json::Value::String(doc_type.to_string())); + payload.insert("text".to_string(), serde_json::Value::String(text.to_string())); + payload.insert("original_id".to_string(), serde_json::Value::String(id.to_string())); let point = PointStruct::new(point_id, vector, payload); @@ -82,4 +82,46 @@ impl VectorDB { Ok(()) } + + pub async fn search( + &self, + query_vector: Vec, + limit: u64, + ) -> Result, Box> { + use qdrant_client::qdrant::SearchPointsBuilder; + + let search_result = self.client + .search_points( + SearchPointsBuilder::new(&self.collection_name, query_vector, limit) + .with_payload(true) + ) + .await?; + + let mut results = Vec::new(); + for point in search_result.result { + let id = point.payload.get("original_id") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .unwrap_or_default(); + + let doc_type = point.payload.get("doc_type") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .unwrap_or_default(); + + let text = point.payload.get("text") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .unwrap_or_default(); + + results.push(VectorSearchResult { + id, + doc_type, + text, + score: point.score, + }); + } + + Ok(results) + } }