Refactor vector search to utilize Qdrant VectorDB
This commit is contained in:
1 parent
61c0e88ad8
commit
3add6c3d31
4 files changed
+79
-11
No files matched your search
+1
-1
@@ -18,7 +18,7 @@ tantivy = "0.26.1"
|
|||||||
tokio = { version = "1.53.1", features = ["full"] }
|
tokio = { version = "1.53.1", features = ["full"] }
|
||||||
tracing = "0.1.44"
|
tracing = "0.1.44"
|
||||||
tracing-subscriber = "0.3.23"
|
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"
|
tracing-appender = "0.2.5"
|
||||||
rmcp = { version = "3.4.0", features = ["server"] }
|
rmcp = { version = "3.4.0", features = ["server"] }
|
||||||
thiserror = "2.0.20"
|
thiserror = "2.0.20"
|
||||||
|
|||||||
@@ -618,8 +618,21 @@ impl McpTool for SemanticCodeSearchHandler {
|
|||||||
|
|
||||||
let query_emb = generate_embedding_async(tool_args.query.clone()).await?;
|
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 results = Vec::new();
|
||||||
|
|
||||||
|
// 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));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Fallback to manual loop if VectorDB is not initialized
|
||||||
|
if !vdb_search {
|
||||||
let mut texts_to_embed = Vec::new();
|
let mut texts_to_embed = Vec::new();
|
||||||
let mut metadata = Vec::new();
|
let mut metadata = Vec::new();
|
||||||
|
|
||||||
@@ -644,15 +657,15 @@ impl McpTool for SemanticCodeSearchHandler {
|
|||||||
}
|
}
|
||||||
|
|
||||||
results.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
|
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 results.is_empty() {
|
||||||
|
|
||||||
if top_results.is_empty() {
|
|
||||||
return Ok(format!("No semantic matches found for query: {}", tool_args.query));
|
return Ok(format!("No semantic matches found for query: {}", tool_args.query));
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut out = format!("Semantic Search Results for '{}':\n", 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));
|
out.push_str(&format!("- [{:.2}] {}: {}\n", score, title, desc));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -13,6 +13,7 @@ mod router;
|
|||||||
mod search;
|
mod search;
|
||||||
pub mod embedding;
|
pub mod embedding;
|
||||||
pub mod indexer;
|
pub mod indexer;
|
||||||
|
pub mod vector_db;
|
||||||
mod state;
|
mod state;
|
||||||
mod store;
|
mod store;
|
||||||
mod tools;
|
mod tools;
|
||||||
@@ -545,6 +546,18 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
|
|
||||||
rt.block_on(async {
|
rt.block_on(async {
|
||||||
let state = Arc::new(MemoryState::new(&base.to_string_lossy()));
|
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 {
|
if let Err(e) = run_server(state).await {
|
||||||
tracing::error!("Server error: {}", e);
|
tracing::error!("Server error: {}", e);
|
||||||
}
|
}
|
||||||
|
|||||||
+47
-5
@@ -3,7 +3,7 @@ use qdrant_client::Qdrant;
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::error::Error;
|
use std::error::Error;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
use log::{info, error};
|
use tracing::{info, error};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
@@ -69,10 +69,10 @@ impl VectorDB {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut payload = std::collections::HashMap::new();
|
let mut payload: std::collections::HashMap<String, serde_json::Value> = std::collections::HashMap::new();
|
||||||
payload.insert("doc_type".to_string(), doc_type.into());
|
payload.insert("doc_type".to_string(), serde_json::Value::String(doc_type.to_string()));
|
||||||
payload.insert("text".to_string(), text.into());
|
payload.insert("text".to_string(), serde_json::Value::String(text.to_string()));
|
||||||
payload.insert("original_id".to_string(), id.into());
|
payload.insert("original_id".to_string(), serde_json::Value::String(id.to_string()));
|
||||||
|
|
||||||
let point = PointStruct::new(point_id, vector, payload);
|
let point = PointStruct::new(point_id, vector, payload);
|
||||||
|
|
||||||
@@ -82,4 +82,46 @@ impl VectorDB {
|
|||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn search(
|
||||||
|
&self,
|
||||||
|
query_vector: Vec<f32>,
|
||||||
|
limit: u64,
|
||||||
|
) -> Result<Vec<VectorSearchResult>, Box<dyn Error + Send + Sync>> {
|
||||||
|
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)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
Reference in new issue
Block a user