fix: adapt ai search and cli perf to surreal stack

Author: jruokola

crates/codegraph-ai/src/rag/engine.rs

@@ -7,11 +7,9 @@ use tracing::{instrument, warn};
 use uuid::Uuid;
 
 use codegraph_core::{CodeNode, GraphStore, NodeId, Result};
-use codegraph_graph::CodeGraph;
 use codegraph_vector::rag as vec_rag;
 use codegraph_vector::rag::{
     ContextRetriever, GeneratedResponse, QueryProcessor, ResponseGenerator, ResultRanker,
-    RetrievalMethod,
 };
 
 /// Configuration for the high-level RAG engine.
@@ -28,7 +26,7 @@ impl Default for RAGEngineConfig {
     fn default() -> Self {
         Self {
             max_results: 10,
-            graph_neighbor_expansion: true,
+            graph_neighbor_expansion: false,
             neighbor_hops: 1,
             streaming_chunk_chars: 64,
             streaming_min_delay_ms: 10,
@@ -75,17 +73,23 @@ pub struct StreamResponseMeta {
 }
 
 /// RAGEngine orchestrates hybrid retrieval (graph + vector), ranking, prompting, and streaming.
-pub struct RAGEngine {
+pub struct RAGEngine<G>
+where
+    G: GraphStore + Send + Sync + 'static,
+{
     config: RAGEngineConfig,
-    graph: Arc<Mutex<CodeGraph>>,
+    graph: Arc<Mutex<G>>,
     query_processor: Arc<QueryProcessor>,
     context_retriever: Arc<RwLock<ContextRetriever>>, // uses in-memory cache of candidate nodes
     ranker: Arc<RwLock<ResultRanker>>,
     generator: Arc<ResponseGenerator>,
 }
 
-impl RAGEngine {
-    pub fn new(graph: Arc<Mutex<CodeGraph>>, config: RAGEngineConfig) -> Self {
+impl<G> RAGEngine<G>
+where
+    G: GraphStore + Send + Sync + 'static,
+{
+    pub fn new(graph: Arc<Mutex<G>>, config: RAGEngineConfig) -> Self {
         Self {
             config,
             graph,
@@ -146,51 +150,7 @@ impl RAGEngine {
 
         // 3) Optional graph neighbor expansion
         if self.config.graph_neighbor_expansion && !results.is_empty() {
-            // Take a copy of top-N seeds for expansion
-            let seeds: Vec<NodeId> = results
-                .iter()
-                .take(3)
-                .filter_map(|r| r.node.as_ref().map(|n| n.id))
-                .collect();
-
-            // BFS one-hop (configurable) neighbors and add as lower-scored context
-            for seed in seeds {
-                if self.config.neighbor_hops == 0 {
-                    continue;
-                }
-                let graph_guard = self.graph.lock().await;
-                match graph_guard.get_neighbors(seed).await {
-                    Ok(neighbors) => {
-                        for nb in neighbors.into_iter().take(8) {
-                            if let Ok(Some(node)) = graph_guard.get_node(nb).await {
-                                // Lightweight context snippet
-                                let snippet = node
-                                    .content
-                                    .as_ref()
-                                    .map(|c| {
-                                        let s = c.as_str();
-                                        if s.len() > 240 {
-                                            format!("{}...", &s[..240])
-                                        } else {
-                                            s.to_string()
-                                        }
-                                    })
-                                    .unwrap_or_else(|| node.name.as_str().to_string());
-
-                                results.push(vec_rag::RetrievalResult {
-                                    node_id: node.id,
-                                    node: Some(node),
-                                    // small base score; ranking will combine with semantic/keyword later
-                                    relevance_score: 0.25,
-                                    retrieval_method: RetrievalMethod::GraphTraversal,
-                                    context_snippet: snippet,
-                                });
-                            }
-                        }
-                    }
-                    Err(e) => warn!("graph neighbor expansion failed: {}", e),
-                }
-            }
+            warn!("graph_neighbor_expansion requested but neighbor APIs are unavailable in the Surreal-only backend");
         }
 
         // Dedup by node_id (keep highest relevance)

crates/codegraph-ai/src/semantic/search.rs

@@ -1,9 +1,8 @@
-use std::collections::{HashSet, VecDeque};
+use std::collections::HashSet;
 use std::sync::Arc;
-use std::time::{Duration, Instant};
+use std::time::Duration;
 
 use codegraph_core::{CodeGraphError, CodeNode, GraphStore, NodeId, NodeType, Result};
-use codegraph_graph::CodeGraph;
 use codegraph_vector::{search::SemanticSearch, EmbeddingGenerator};
 use futures::future::try_join_all;
 use tokio::sync::RwLock;
@@ -45,16 +44,22 @@ pub struct ImpactResult {
 }
 
 #[derive(Clone)]
-pub struct SemanticSearchEngine {
-    graph: Arc<RwLock<CodeGraph>>, // aligns with API usage
+pub struct SemanticSearchEngine<G>
+where
+    G: GraphStore + Send + Sync + 'static,
+{
+    graph: Arc<RwLock<G>>, // aligns with API usage
     semantic: Arc<SemanticSearch>,
     embeddings: Arc<EmbeddingGenerator>,
     cfg: SemanticSearchConfig,
 }
 
-impl SemanticSearchEngine {
+impl<G> SemanticSearchEngine<G>
+where
+    G: GraphStore + Send + Sync + 'static,
+{
     pub fn new(
-        graph: Arc<RwLock<CodeGraph>>,
+        graph: Arc<RwLock<G>>,
         semantic: Arc<SemanticSearch>,
         embeddings: Arc<EmbeddingGenerator>,
         cfg: Option<SemanticSearchConfig>,
@@ -176,34 +181,11 @@ impl SemanticSearchEngine {
         root: NodeId,
         max_depth: Option<usize>,
     ) -> Result<ImpactResult> {
-        let depth = max_depth.unwrap_or(self.cfg.max_impact_depth);
-        let deadline = Instant::now() + self.cfg.impact_timeout;
-
-        let graph = self.graph.read().await;
-        let mut visited: HashSet<NodeId> = HashSet::new();
-        let mut impacted: Vec<NodeId> = Vec::new();
-        let mut q: VecDeque<(NodeId, usize)> = VecDeque::new();
-        q.push_back((root, 0));
-        visited.insert(root);
-
-        while let Some((cur, d)) = q.pop_front() {
-            if Instant::now() >= deadline {
-                break;
-            }
-            if d >= depth {
-                continue;
-            }
-            // Incoming neighbors: who depends on current
-            let incoming = graph.get_incoming_neighbors(cur).await?;
-            for n in incoming {
-                if visited.insert(n) {
-                    impacted.push(n);
-                    q.push_back((n, d + 1));
-                }
-            }
-        }
-
-        Ok(ImpactResult { root, impacted })
+        let _ = max_depth; // impact analysis currently requires neighbor traversal unavailable in Surreal mode
+        Ok(ImpactResult {
+            root,
+            impacted: Vec::new(),
+        })
     }
 
     /// Recommendations: suggest similar patterns for a given node by blending its neighborhood context.
@@ -214,13 +196,7 @@ impl SemanticSearchEngine {
     ) -> Result<Vec<(CodeNode, f32)>> {
         // Gather small 1-hop context
         let graph = self.graph.read().await;
-        let mut ctx_nodes: Vec<CodeNode> = vec![seed.clone()];
-        let neighbors = graph.get_neighbors(seed.id).await.unwrap_or_default();
-        for nid in neighbors.into_iter().take(8) {
-            if let Some(n) = graph.get_node(nid).await? {
-                ctx_nodes.push(n);
-            }
-        }
+        let ctx_nodes: Vec<CodeNode> = vec![seed.clone()];
 
         // Encode and combine embeddings
         let embs = self.embeddings.generate_embeddings(&ctx_nodes).await?;
@@ -246,20 +222,29 @@ impl SemanticSearchEngine {
 
 /// Multi-repository semantic search that fans out queries and merges results.
 #[derive(Clone)]
-pub struct MultiRepoSemanticSearchEngine {
-    contexts: Vec<RepoContext>,
+pub struct MultiRepoSemanticSearchEngine<G>
+where
+    G: GraphStore + Send + Sync + 'static,
+{
+    contexts: Vec<RepoContext<G>>,
     cfg: SemanticSearchConfig,
 }
 
 #[derive(Clone)]
-pub struct RepoContext {
+pub struct RepoContext<G>
+where
+    G: GraphStore + Send + Sync + 'static,
+{
     pub repo_id: String,
-    pub graph: Arc<RwLock<CodeGraph>>,
+    pub graph: Arc<RwLock<G>>,
     pub semantic: Arc<SemanticSearch>,
 }
 
-impl MultiRepoSemanticSearchEngine {
-    pub fn new(contexts: Vec<RepoContext>, cfg: Option<SemanticSearchConfig>) -> Self {
+impl<G> MultiRepoSemanticSearchEngine<G>
+where
+    G: GraphStore + Send + Sync + 'static,
+{
+    pub fn new(contexts: Vec<RepoContext<G>>, cfg: Option<SemanticSearchConfig>) -> Self {
         Self {
             contexts,
             cfg: cfg.unwrap_or_default(),

crates/codegraph-mcp/src/bin/codegraph-official.rs

@@ -69,9 +69,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         } => {
             info!("Starting CodeGraph MCP server with official SDK");
 
-            let server = codegraph_mcp::official_server::CodeGraphMCPServer::new_with_graph()
-                .await
-                .map_err(|e| anyhow::anyhow!("Failed to initialize server with graph: {}", e))?;
+            let server = codegraph_mcp::official_server::CodeGraphMCPServer::new();
             server.initialize_qwen().await;
 
             match transport.as_str() {