Feat: Implement SurrealDB HNSW vector search (Phase 2)

jruokola · jruokola · commit 2051553dee67 · 2025-11-07T15:02:47.000+02:00
Adds real HNSW vector search with metadata filtering:

**SurrealDB Storage Layer** (crates/codegraph-graph/src/surrealdb_storage.rs):
- vector_search_knn(): HNSW search using &lt;|K,EF|&gt; operator
- vector_search_with_metadata(): Filtered HNSW search (type/lang/path)
- get_nodes_by_ids(): Batch node loading with caching

**Cloud Search Implementation** (crates/codegraph-mcp/src/server.rs):
- Creates dedicated SurrealDB connection for cloud mode
- Uses vector_search_with_metadata() for filtered HNSW search
- Loads full nodes efficiently with batch query
- Returns results with HNSW similarity scores
- Detailed performance metrics for each step

**Features**:
- ✅ Automatic SurrealDB connection from env vars
- ✅ HNSW search with configurable ef_search parameter
- ✅ Metadata filtering (node_type, language, file_path patterns)
- ✅ OR-pattern support for file paths (src/|lib/)
- ✅ Cache-aware node loading
- ✅ Performance tracking (embedding, connect, search, load, format)

**Status**: Phase 2 complete, compiles successfully
**Next**: Phase 3 - Jina reranking integration (optional enhancement)
**Testing**: Requires SurrealDB instance with HNSW indexes
diff --git a/crates/codegraph-graph/src/surrealdb_storage.rs b/crates/codegraph-graph/src/surrealdb_storage.rs
@@ -288,6 +288,192 @@ impl SurrealDbStorage {
     }
     */
 
+    /// Vector search using SurrealDB HNSW indexes
+    /// Returns node IDs and similarity scores
+    pub async fn vector_search_knn(
+        &self,
+        query_embedding: Vec<f32>,
+        limit: usize,
+        ef_search: usize,
+    ) -> Result<Vec<(String, f32)>> {
+        info!(
+            "Executing HNSW vector search with limit={}, ef_search={}",
+            limit, ef_search
+        );
+
+        // Convert f32 to f64 for SurrealDB
+        let query_vec: Vec<f64> = query_embedding.iter().map(|&f| f as f64).collect();
+
+        // SurrealDB HNSW search using <|K,EF|> operator
+        // vector::distance::knn() reuses pre-computed distance from HNSW
+        let query = r#"
+            SELECT id, vector::distance::knn() AS score
+            FROM nodes
+            WHERE embedding <|$limit,$ef_search|> $query_embedding
+            ORDER BY score ASC
+            LIMIT $limit
+        "#;
+
+        let mut result = self
+            .db
+            .query(query)
+            .bind(("query_embedding", query_vec))
+            .bind(("limit", limit))
+            .bind(("ef_search", ef_search))
+            .await
+            .map_err(|e| CodeGraphError::Database(format!("HNSW search failed: {}", e)))?;
+
+        #[derive(Deserialize)]
+        struct SearchResult {
+            id: String,
+            score: f64,
+        }
+
+        let results: Vec<SearchResult> = result.take(0).map_err(|e| {
+            CodeGraphError::Database(format!("Failed to extract search results: {}", e))
+        })?;
+
+        Ok(results
+            .into_iter()
+            .map(|r| (r.id, r.score as f32))
+            .collect())
+    }
+
+    /// Vector search with metadata filtering
+    pub async fn vector_search_with_metadata(
+        &self,
+        query_embedding: Vec<f32>,
+        limit: usize,
+        ef_search: usize,
+        node_type: Option<String>,
+        language: Option<String>,
+        file_path_pattern: Option<String>,
+    ) -> Result<Vec<(String, f32)>> {
+        info!(
+            "Executing filtered HNSW search: type={:?}, lang={:?}, path={:?}",
+            node_type, language, file_path_pattern
+        );
+
+        let query_vec: Vec<f64> = query_embedding.iter().map(|&f| f as f64).collect();
+
+        // Build dynamic WHERE clause
+        let mut where_clauses =
+            vec!["embedding <|$limit,$ef_search|> $query_embedding".to_string()];
+
+        if let Some(ref nt) = node_type {
+            where_clauses.push(format!("node_type = '{}'", nt));
+        }
+
+        if let Some(ref lang) = language {
+            where_clauses.push(format!("language = '{}'", lang));
+        }
+
+        if let Some(ref path) = file_path_pattern {
+            // Support OR patterns like "src/|lib/"
+            if path.contains('|') {
+                let patterns: Vec<String> = path
+                    .split('|')
+                    .map(|p| format!("file_path CONTAINS '{}'", p))
+                    .collect();
+                where_clauses.push(format!("({})", patterns.join(" OR ")));
+            } else {
+                where_clauses.push(format!("file_path CONTAINS '{}'", path));
+            }
+        }
+
+        let where_clause = where_clauses.join(" AND ");
+
+        let query = format!(
+            r#"
+            SELECT id, vector::distance::knn() AS score
+            FROM nodes
+            WHERE {}
+            ORDER BY score ASC
+            LIMIT $limit
+        "#,
+            where_clause
+        );
+
+        let mut result = self
+            .db
+            .query(&query)
+            .bind(("query_embedding", query_vec))
+            .bind(("limit", limit))
+            .bind(("ef_search", ef_search))
+            .await
+            .map_err(|e| CodeGraphError::Database(format!("Filtered HNSW search failed: {}", e)))?;
+
+        #[derive(Deserialize)]
+        struct SearchResult {
+            id: String,
+            score: f64,
+        }
+
+        let results: Vec<SearchResult> = result.take(0).map_err(|e| {
+            CodeGraphError::Database(format!("Failed to extract filtered results: {}", e))
+        })?;
+
+        Ok(results
+            .into_iter()
+            .map(|r| (r.id, r.score as f32))
+            .collect())
+    }
+
+    /// Get multiple nodes by their IDs in one query
+    pub async fn get_nodes_by_ids(&self, ids: &[String]) -> Result<Vec<CodeNode>> {
+        if ids.is_empty() {
+            return Ok(Vec::new());
+        }
+
+        debug!("Getting {} nodes by IDs", ids.len());
+
+        // Check cache first for all IDs
+        let mut nodes = Vec::new();
+        let mut missing_ids = Vec::new();
+
+        if self.config.cache_enabled {
+            for id_str in ids {
+                if let Ok(id) = NodeId::parse_str(id_str) {
+                    if let Some(cached) = self.node_cache.get(&id) {
+                        nodes.push(cached.clone());
+                    } else {
+                        missing_ids.push(id_str.clone());
+                    }
+                }
+            }
+        } else {
+            missing_ids = ids.to_vec();
+        }
+
+        // Fetch missing nodes from database
+        if !missing_ids.is_empty() {
+            let query = "SELECT * FROM nodes WHERE id IN $ids";
+            let mut result = self
+                .db
+                .query(query)
+                .bind(("ids", missing_ids))
+                .await
+                .map_err(|e| CodeGraphError::Database(format!("Failed to query nodes: {}", e)))?;
+
+            let db_nodes: Vec<HashMap<String, JsonValue>> = result.take(0).map_err(|e| {
+                CodeGraphError::Database(format!("Failed to extract query results: {}", e))
+            })?;
+
+            for data in db_nodes {
+                let node = self.surreal_to_node(data)?;
+
+                // Update cache
+                if self.config.cache_enabled {
+                    self.node_cache.insert(node.id, node.clone());
+                }
+
+                nodes.push(node);
+            }
+        }
+
+        Ok(nodes)
+    }
+
     /// Convert CodeNode to SurrealDB-compatible format
     fn node_to_surreal(&self, node: &CodeNode) -> Result<HashMap<String, JsonValue>> {
         let mut data = HashMap::new();
diff --git a/crates/codegraph-mcp/src/server.rs b/crates/codegraph-mcp/src/server.rs
@@ -736,7 +736,7 @@ async fn cloud_search_impl(
     paths: Option<Vec<String>>,
     langs: Option<Vec<String>>,
     limit: usize,
-    graph: &codegraph_graph::CodeGraph,
+    _graph: &codegraph_graph::CodeGraph,
 ) -> anyhow::Result<Value> {
     use codegraph_core::Language;
 
@@ -745,94 +745,123 @@ async fn cloud_search_impl(
     tracing::info!("🌐 Cloud Mode: SurrealDB HNSW + Jina reranking");
 
     // 1. Generate query embedding using Jina/Cloud provider
+    let start_embedding = Instant::now();
     let embedding_gen = get_embedding_generator().await?;
     let query_embedding = embedding_gen.generate_text_embedding(&query).await?;
+    let embedding_time = start_embedding.elapsed().as_millis() as u64;
 
     // 2. Overretrieve for reranking (3x limit)
     let overretrieve_limit = limit * 3;
 
-    // 3. SurrealDB HNSW search
-    let storage = graph.get_storage();
+    // 3. Create SurrealDB storage connection for cloud mode
+    let start_connect = Instant::now();
+    let surrealdb_config = codegraph_graph::SurrealDbConfig {
+        connection: std::env::var("SURREALDB_URL")
+            .unwrap_or_else(|_| "ws://localhost:3004".to_string()),
+        namespace: std::env::var("SURREALDB_NAMESPACE").unwrap_or_else(|_| "codegraph".to_string()),
+        database: std::env::var("SURREALDB_DATABASE").unwrap_or_else(|_| "main".to_string()),
+        username: std::env::var("SURREALDB_USERNAME").ok(),
+        password: std::env::var("SURREALDB_PASSWORD").ok(),
+        strict_mode: false,
+        auto_migrate: false, // Don't migrate on every search
+        cache_enabled: true,
+    };
 
-    // For now, use get_all_nodes and do in-memory similarity search
-    // TODO: Implement proper HNSW search in SurrealDB storage layer
-    let all_nodes = storage.get_all_nodes().await?;
+    let surrealdb_storage = codegraph_graph::SurrealDbStorage::new(surrealdb_config).await?;
+    let connect_time = start_connect.elapsed().as_millis() as u64;
 
-    if all_nodes.is_empty() {
-        return Ok(json!({
-            "results": [],
-            "message": "No nodes found in database. Please index your codebase first.",
-            "performance": {
-                "total_ms": start_total.elapsed().as_millis(),
-                "mode": "cloud"
-            }
-        }));
-    }
+    // 4. SurrealDB HNSW search with metadata filtering
+    let start_search = Instant::now();
 
-    // Filter by language and path
-    let lang_filter: Option<Vec<Language>> = langs.as_ref().map(|langs| {
+    // Build filter parameters
+    let node_type_filter = langs.as_ref().map(|langs| {
         langs
             .iter()
             .filter_map(|l| match l.to_lowercase().as_str() {
-                "rust" => Some(Language::Rust),
-                "python" => Some(Language::Python),
-                "javascript" | "js" => Some(Language::JavaScript),
-                "typescript" | "ts" => Some(Language::TypeScript),
-                "go" => Some(Language::Go),
-                "java" => Some(Language::Java),
+                "rust" => Some("Rust".to_string()),
+                "python" => Some("Python".to_string()),
+                "javascript" | "js" => Some("JavaScript".to_string()),
+                "typescript" | "ts" => Some("TypeScript".to_string()),
+                "go" => Some("Go".to_string()),
+                "java" => Some("Java".to_string()),
                 _ => None,
             })
-            .collect()
+            .next() // Take first matching language for now
     });
 
-    let filtered_nodes: Vec<_> = all_nodes
-        .into_iter()
-        .filter(|node| {
-            // Filter by language
-            if let Some(ref langs) = lang_filter {
-                if !langs.contains(&node.language) {
-                    return false;
-                }
-            }
+    let file_path_pattern = paths.as_ref().map(|p| p.join("|"));
+
+    let search_results = surrealdb_storage
+        .vector_search_with_metadata(
+            query_embedding.clone(),
+            overretrieve_limit,
+            100, // ef_search parameter for HNSW
+            node_type_filter,
+            None, // language filter (separate from node_type)
+            file_path_pattern,
+        )
+        .await?;
 
-            // Filter by paths
-            if let Some(ref paths) = paths {
-                if !paths.iter().any(|p| node.location.file_path.contains(p)) {
-                    return false;
-                }
+    let search_time = start_search.elapsed().as_millis() as u64;
+
+    if search_results.is_empty() {
+        return Ok(json!({
+            "results": [],
+            "message": "No results found. Ensure codebase is indexed with embeddings.",
+            "performance": {
+                "total_ms": start_total.elapsed().as_millis(),
+                "embedding_ms": embedding_time,
+                "connect_ms": connect_time,
+                "search_ms": search_time,
+                "mode": "cloud"
             }
+        }));
+    }
 
-            true
-        })
-        .collect();
+    // 5. Load full nodes from SurrealDB
+    let start_load = Instant::now();
+    let node_ids: Vec<String> = search_results.iter().map(|(id, _)| id.clone()).collect();
+    let nodes = surrealdb_storage.get_nodes_by_ids(&node_ids).await?;
+    let load_time = start_load.elapsed().as_millis() as u64;
+
+    // 6. TODO: Add Jina reranking here
+    // For now, use HNSW scores directly
 
-    // TODO: For MVP, return basic results without reranking
-    // Full implementation requires Jina reranking integration
-    let results: Vec<Value> = filtered_nodes
+    // 7. Format results
+    let start_format = Instant::now();
+    let results: Vec<Value> = nodes
         .iter()
+        .zip(search_results.iter())
         .take(limit)
-        .map(|node| {
+        .map(|(node, (_id, score))| {
             json!({
                 "id": node.id,
                 "name": node.name,
-                "node_type": format!("{:?}", node.node_type),
-                "language": format!("{:?}", node.language),
+                "node_type": node.node_type.as_ref().map(|nt| format!("{:?}", nt)).unwrap_or_default(),
+                "language": node.language.as_ref().map(|l| format!("{:?}", l)).unwrap_or_default(),
                 "file_path": node.location.file_path,
-                "start_line": node.location.start_line,
+                "start_line": node.location.line,
                 "end_line": node.location.end_line,
-                "score": 0.5, // Placeholder score
+                "score": 1.0 - score,  // Convert distance to similarity
                 "summary": node.content.as_deref().unwrap_or("").chars().take(160).collect::<String>()
             })
         })
         .collect();
+    let format_time = start_format.elapsed().as_millis() as u64;
 
     Ok(json!({
         "results": results,
         "total_results": results.len(),
         "performance": {
             "total_ms": start_total.elapsed().as_millis(),
+            "embedding_generation_ms": embedding_time,
+            "surrealdb_connection_ms": connect_time,
+            "hnsw_search_ms": search_time,
+            "node_loading_ms": load_time,
+            "formatting_ms": format_time,
             "mode": "cloud",
-            "note": "SurrealDB HNSW search - reranking not yet implemented"
+            "hnsw_enabled": true,
+            "reranking_enabled": false  // TODO: Phase 2.5
         }
     }))
 }
