Skip to content

Commit 82653ab

Browse files
committed
Phase 3: Integrate Jina reranking into cloud_search_impl()
Adds intelligent reranking layer after SurrealDB HNSW search: - Overretrieve 3x results for better reranking pool - Create JinaEmbeddingProvider with reranker-v3 model - Prepare documents from node name + content - Call jina_provider.rerank() with query - Return top-N reranked results with relevance scores Graceful degradation on errors: - Skip reranking if JINA_API_KEY not set - Fallback to HNSW scores if provider creation fails - Fallback to HNSW scores if rerank API call fails Performance tracking: - Separate reranking_ms metric - reranking_enabled flag in response - Detailed logging at each step Feature flags: Reranking only active when jina feature enabled. Related: DUAL_MODE_SEARCH_PLAN.md Phase 3
1 parent 2051553 commit 82653ab

File tree

1 file changed

+75
-8
lines changed

1 file changed

+75
-8
lines changed

crates/codegraph-mcp/src/server.rs

Lines changed: 75 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -824,16 +824,82 @@ async fn cloud_search_impl(
824824
let nodes = surrealdb_storage.get_nodes_by_ids(&node_ids).await?;
825825
let load_time = start_load.elapsed().as_millis() as u64;
826826

827-
// 6. TODO: Add Jina reranking here
828-
// For now, use HNSW scores directly
827+
// 6. Jina reranking (if available)
828+
let start_rerank = Instant::now();
829+
let mut rerank_enabled = false;
830+
let mut rerank_time = 0u64;
831+
832+
// Try to create Jina provider for reranking
833+
#[cfg(feature = "jina")]
834+
let reranked_results: Vec<(usize, f32)> = {
835+
match codegraph_vector::JinaConfig::default().api_key.is_empty() {
836+
true => {
837+
tracing::warn!("JINA_API_KEY not set, skipping reranking");
838+
// Return HNSW order
839+
(0..nodes.len()).map(|i| (i, search_results[i].1)).collect()
840+
}
841+
false => {
842+
let jina_config = codegraph_vector::JinaConfig {
843+
enable_reranking: true,
844+
reranking_model: "jina-reranker-v3".to_string(),
845+
reranking_top_n: limit,
846+
..codegraph_vector::JinaConfig::default()
847+
};
848+
849+
match codegraph_vector::JinaEmbeddingProvider::new(jina_config) {
850+
Ok(jina_provider) => {
851+
// Prepare documents for reranking
852+
let documents: Vec<String> = nodes
853+
.iter()
854+
.map(|node| {
855+
format!("{}\n{}", node.name, node.content.as_deref().unwrap_or(""))
856+
})
857+
.collect();
858+
859+
match jina_provider.rerank(&query, documents).await {
860+
Ok(rerank_results) => {
861+
tracing::info!(
862+
"🎯 Jina reranking: {} results",
863+
rerank_results.len()
864+
);
865+
rerank_enabled = true;
866+
// Map rerank results to (index, score)
867+
rerank_results
868+
.into_iter()
869+
.map(|r| (r.index, r.relevance_score))
870+
.collect()
871+
}
872+
Err(e) => {
873+
tracing::warn!("Jina reranking failed: {}, using HNSW scores", e);
874+
// Fallback to HNSW order
875+
(0..nodes.len()).map(|i| (i, search_results[i].1)).collect()
876+
}
877+
}
878+
}
879+
Err(e) => {
880+
tracing::warn!("Failed to create Jina provider: {}, using HNSW scores", e);
881+
(0..nodes.len()).map(|i| (i, search_results[i].1)).collect()
882+
}
883+
}
884+
}
885+
}
886+
};
829887

830-
// 7. Format results
888+
#[cfg(not(feature = "jina"))]
889+
let reranked_results: Vec<(usize, f32)> = {
890+
tracing::info!("Jina feature not enabled, using HNSW scores only");
891+
(0..nodes.len()).map(|i| (i, search_results[i].1)).collect()
892+
};
893+
894+
rerank_time = start_rerank.elapsed().as_millis() as u64;
895+
896+
// 7. Format results using reranked order
831897
let start_format = Instant::now();
832-
let results: Vec<Value> = nodes
898+
let results: Vec<Value> = reranked_results
833899
.iter()
834-
.zip(search_results.iter())
835900
.take(limit)
836-
.map(|(node, (_id, score))| {
901+
.map(|(index, score)| {
902+
let node = &nodes[*index];
837903
json!({
838904
"id": node.id,
839905
"name": node.name,
@@ -842,7 +908,7 @@ async fn cloud_search_impl(
842908
"file_path": node.location.file_path,
843909
"start_line": node.location.line,
844910
"end_line": node.location.end_line,
845-
"score": 1.0 - score, // Convert distance to similarity
911+
"score": *score,
846912
"summary": node.content.as_deref().unwrap_or("").chars().take(160).collect::<String>()
847913
})
848914
})
@@ -858,10 +924,11 @@ async fn cloud_search_impl(
858924
"surrealdb_connection_ms": connect_time,
859925
"hnsw_search_ms": search_time,
860926
"node_loading_ms": load_time,
927+
"reranking_ms": rerank_time,
861928
"formatting_ms": format_time,
862929
"mode": "cloud",
863930
"hnsw_enabled": true,
864-
"reranking_enabled": false // TODO: Phase 2.5
931+
"reranking_enabled": rerank_enabled
865932
}
866933
}))
867934
}

0 commit comments

Comments
 (0)