1use crate::cli::MemoryType;
4use crate::errors::AppError;
5use crate::graph::traverse_from_memories_with_hops;
6use crate::output::{self, JsonOutputFormat, RecallItem};
7use crate::paths::AppPaths;
8use crate::storage::connection::open_ro;
9use crate::storage::entities;
10use crate::storage::memories;
11
12use std::collections::HashMap;
13
14#[derive(clap::Args)]
21#[command(after_long_help = "EXAMPLES:\n \
22 # Basic hybrid search combining FTS5 + vector via RRF\n \
23 sqlite-graphrag hybrid-search \"postgres migration deadlock\" --k 10\n\n \
24 # Tune RRF weights to favor keyword matches over semantic similarity\n \
25 sqlite-graphrag hybrid-search \"jwt auth\" --weight-fts 1.5 --weight-vec 0.5 --k 5\n\n \
26 # Add graph traversal matches (entities connected to top results)\n \
27 sqlite-graphrag hybrid-search \"frontend architecture\" --with-graph --k 10\n\n \
28 # Graph traversal with custom depth and minimum edge weight\n \
29 sqlite-graphrag hybrid-search \"auth design\" --with-graph --max-hops 3 --min-weight 0.5 --k 10\n\n \
30NOTES:\n \
31 --with-graph enables entity graph traversal seeded by the top RRF results.\n \
32 Graph matches appear in the `graph_matches` array (separate from `results`).\n \
33 Without --with-graph, `graph_matches` is always empty.")]
34pub struct HybridSearchArgs {
35 #[arg(
36 allow_hyphen_values = true,
37 help = "Hybrid search query (vector KNN + FTS5 BM25 fused via RRF)"
38 )]
39 pub query: String,
40 #[arg(short = 'k', long, aliases = ["limit", "top-k"], default_value = "10", value_parser = crate::parsers::parse_k_range)]
45 pub k: usize,
46 #[arg(long, default_value = "60")]
47 pub rrf_k: u32,
48 #[arg(long, default_value = "1.0")]
49 pub weight_vec: f32,
50 #[arg(long, default_value = "1.0")]
51 pub weight_fts: f32,
52 #[arg(long, value_enum)]
56 pub r#type: Option<MemoryType>,
57 #[arg(long)]
58 pub namespace: Option<String>,
59 #[arg(long)]
60 pub with_graph: bool,
61 #[arg(long, help = "Skip live query embedding; serve FTS5 BM25 only")]
64 pub fallback_fts_only: bool,
65 #[arg(long)]
67 pub max_hops: Option<u32>,
68 #[arg(long)]
70 pub min_weight: Option<f64>,
71 #[arg(long, value_enum, default_value_t = JsonOutputFormat::Json)]
72 pub format: JsonOutputFormat,
73 #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
74 pub db: Option<String>,
75 #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
77 pub json: bool,
78}
79
80#[derive(serde::Serialize)]
81pub struct HybridSearchItem {
82 pub memory_id: i64,
83 pub name: String,
84 pub namespace: String,
85 #[serde(rename = "type")]
86 pub memory_type: String,
87 pub description: String,
88 pub body: String,
89 pub snippet: String,
90 pub combined_score: f64,
91 pub score: f64,
93 pub source: String,
95 #[serde(skip_serializing_if = "Option::is_none")]
96 pub vec_rank: Option<usize>,
97 #[serde(skip_serializing_if = "Option::is_none")]
98 pub fts_rank: Option<usize>,
99 #[serde(skip_serializing_if = "Option::is_none")]
101 pub rrf_score: Option<f64>,
102 pub normalized_score: f64,
104 #[serde(skip_serializing_if = "Option::is_none")]
109 pub vec_distance: Option<f64>,
110 #[serde(skip_serializing_if = "Option::is_none")]
113 pub fts_bm25: Option<f64>,
114}
115
116#[derive(serde::Serialize)]
118pub struct Weights {
119 pub vec: f32,
120 pub fts: f32,
121}
122
123#[derive(serde::Serialize)]
124pub struct HybridSearchResponse {
125 pub query: String,
126 pub k: usize,
127 pub rrf_k: u32,
129 pub weights: Weights,
131 pub results: Vec<HybridSearchItem>,
132 pub graph_matches: Vec<RecallItem>,
133 #[serde(skip_serializing_if = "std::ops::Not::not")]
137 pub fts_degraded: bool,
138 #[serde(skip_serializing_if = "Option::is_none")]
142 pub fts_error: Option<String>,
143 #[serde(skip_serializing_if = "std::ops::Not::not")]
147 pub fts_auto_rebuilt: bool,
148 #[serde(skip_serializing_if = "std::ops::Not::not", default)]
152 pub vec_degraded: bool,
153 #[serde(skip_serializing_if = "Option::is_none")]
157 pub vec_error: Option<String>,
158 #[serde(skip_serializing_if = "Option::is_none")]
162 pub warning: Option<String>,
163 pub elapsed_ms: u64,
165}
166
167#[tracing::instrument(skip_all, level = "debug", name = "hybrid_search")]
168pub fn run(args: HybridSearchArgs) -> Result<(), AppError> {
169 let start = std::time::Instant::now();
170 let _ = args.format;
171 tracing::debug!(target: "hybrid_search", query = %args.query, k = args.k, "fusing results");
172
173 if !args.with_graph {
177 if args.max_hops.is_some() {
178 return Err(AppError::Validation(
179 "--max-hops requires --with-graph to be active".to_string(),
180 ));
181 }
182 if args.min_weight.is_some() {
183 return Err(AppError::Validation(
184 "--min-weight requires --with-graph to be active".to_string(),
185 ));
186 }
187 }
188
189 let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
190 let paths = AppPaths::resolve(args.db.as_deref())?;
191 crate::storage::connection::ensure_db_ready(&paths)?;
192
193 output::emit_progress_i18n(
194 "Computing query embedding...",
195 "Calculando embedding da consulta...",
196 );
197 let conn = open_ro(&paths.db)?;
198 let (embedding, vec_degraded, vec_error) = if args.fallback_fts_only {
203 (None, true, Some("fallback_fts_only requested".to_string()))
204 } else {
205 match crate::embedder::try_embed_query_with_fallback(&paths.models, &args.query) {
206 Ok(v) => (Some(v), false, None),
207 Err(reason) => {
208 let msg = reason.to_string();
209 tracing::warn!(target: "hybrid_search", fallback_reason = %msg, "live embedding failed; falling back to FTS5");
210 (None, true, Some(msg))
211 }
212 }
213 };
214
215 let memory_type_str = args.r#type.map(|t| t.as_str());
216
217 let vec_results: Vec<(i64, f32)> = if let Some(emb) = embedding.as_ref() {
218 memories::knn_search(
219 &conn,
220 emb,
221 std::slice::from_ref(&namespace),
222 memory_type_str,
223 args.k * 2,
224 )?
225 } else {
226 Vec::new()
227 };
228
229 let vec_rank_map: HashMap<i64, usize> = vec_results
231 .iter()
232 .enumerate()
233 .map(|(pos, (id, _))| (*id, pos + 1))
234 .collect();
235
236 let vec_distance_map: HashMap<i64, f64> = vec_results
238 .iter()
239 .map(|(id, dist)| (*id, *dist as f64))
240 .collect();
241
242 let (fts_results, fts_degraded, fts_error, fts_auto_rebuilt) = if args.weight_fts == 0.0 {
243 (vec![], false, None, false)
244 } else {
245 match memories::fts_search(&conn, &args.query, &namespace, memory_type_str, args.k * 2) {
246 Ok(r) => (r, false, None, false),
247 Err(e) => {
248 let err_msg = e.to_string();
249 let is_malformed = err_msg.contains("malformed") || err_msg.contains("corrupt");
250 if is_malformed {
251 tracing::warn!(target: "hybrid_search", "FTS5 index corrupted, attempting auto-rebuild");
252 if conn
253 .execute_batch("INSERT INTO fts_memories(fts_memories) VALUES('rebuild');")
254 .is_ok()
255 {
256 match memories::fts_search(
257 &conn,
258 &args.query,
259 &namespace,
260 memory_type_str,
261 args.k * 2,
262 ) {
263 Ok(r) => (r, false, None, true),
264 Err(e2) => {
265 tracing::error!(target: "hybrid_search", error = %e2, "FTS5 auto-rebuild failed to recover");
266 (vec![], true, Some(e2.to_string()), true)
267 }
268 }
269 } else {
270 (vec![], true, Some(err_msg), false)
271 }
272 } else {
273 tracing::warn!(target: "hybrid_search", error = %e, "FTS5 query failed, falling back to vec-only");
274 (vec![], true, Some(err_msg), false)
275 }
276 }
277 }
278 };
279
280 let fts_rank_map: HashMap<i64, usize> = fts_results
282 .iter()
283 .enumerate()
284 .map(|(pos, row)| (row.id, pos + 1))
285 .collect();
286
287 let rrf_k = args.rrf_k as f64;
288
289 let mut combined_scores: crate::hash::AHashMap<i64, f64> =
291 crate::hash::AHashMap::with_capacity_and_hasher(
292 vec_results.len() + fts_results.len(),
293 Default::default(),
294 );
295
296 for (rank, (memory_id, _)) in vec_results.iter().enumerate() {
297 let score = args.weight_vec as f64 * (1.0 / (rrf_k + rank as f64 + 1.0));
298 *combined_scores.entry(*memory_id).or_insert(0.0) += score;
299 }
300
301 for (rank, row) in fts_results.iter().enumerate() {
302 let score = args.weight_fts as f64 * (1.0 / (rrf_k + rank as f64 + 1.0));
303 *combined_scores.entry(row.id).or_insert(0.0) += score;
304 }
305
306 let mut ranked: Vec<(i64, f64)> = combined_scores.into_iter().collect();
308 ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
309 ranked.truncate(args.k);
310
311 let top_ids: Vec<i64> = ranked.iter().map(|(id, _)| *id).collect();
313
314 let mut memory_data: crate::hash::AHashMap<i64, memories::MemoryRow> =
316 crate::hash::AHashMap::with_capacity_and_hasher(ranked.len(), Default::default());
317 for id in &top_ids {
318 if let Some(row) = memories::read_full(&conn, *id)? {
319 memory_data.insert(*id, row);
320 }
321 }
322
323 let max_possible = args.weight_vec as f64 * (1.0 / (rrf_k + 1.0))
324 + args.weight_fts as f64 * (1.0 / (rrf_k + 1.0));
325
326 let results: Vec<HybridSearchItem> = ranked
328 .into_iter()
329 .filter_map(|(memory_id, combined_score)| {
330 let normalized_score = if max_possible > 0.0 {
331 combined_score / max_possible
332 } else {
333 0.0
334 };
335 memory_data.remove(&memory_id).map(|row| {
336 let snippet: String = row.body.chars().take(300).collect();
337 HybridSearchItem {
338 memory_id: row.id,
339 name: row.name,
340 namespace: row.namespace,
341 memory_type: row.memory_type,
342 description: row.description,
343 body: row.body,
344 snippet,
345 combined_score,
346 score: combined_score,
347 source: "hybrid".to_string(),
348 vec_rank: vec_rank_map.get(&memory_id).copied(),
349 fts_rank: fts_rank_map.get(&memory_id).copied(),
350 rrf_score: Some(combined_score),
351 normalized_score,
352 vec_distance: vec_distance_map.get(&memory_id).copied(),
353 fts_bm25: None,
354 }
355 })
356 })
357 .collect();
358
359 let mut graph_matches: Vec<RecallItem> = Vec::with_capacity(8);
361 if let Some(emb) = args
362 .with_graph
363 .then_some(())
364 .filter(|_| !results.is_empty())
365 .and(embedding.as_ref())
366 {
367 let namespace_for_graph = namespace.clone();
368 let memory_ids: Vec<i64> = results.iter().map(|r| r.memory_id).collect();
369
370 let entity_knn = entities::knn_search(&conn, emb, &namespace_for_graph, 5)?;
371 let entity_ids: Vec<i64> = entity_knn.iter().map(|(id, _)| *id).collect();
372
373 let all_seed_ids: Vec<i64> = memory_ids
374 .iter()
375 .chain(entity_ids.iter())
376 .copied()
377 .collect();
378
379 if !all_seed_ids.is_empty() {
380 let graph_memory_ids = traverse_from_memories_with_hops(
381 &conn,
382 &all_seed_ids,
383 &namespace_for_graph,
384 args.min_weight.unwrap_or(0.3),
385 args.max_hops.unwrap_or(2),
386 )?;
387
388 let already_in_results: std::collections::HashSet<i64> =
389 results.iter().map(|r| r.memory_id).collect();
390
391 for (graph_mem_id, hop) in graph_memory_ids {
392 if already_in_results.contains(&graph_mem_id) {
393 continue;
394 }
395 if let Some(row) = memories::read_full(&conn, graph_mem_id)? {
396 let snippet: String = row.body.chars().take(300).collect();
397 let graph_distance = 1.0 - 1.0 / (hop as f32 + 1.0);
398 graph_matches.push(RecallItem {
399 memory_id: row.id,
400 name: row.name,
401 namespace: row.namespace,
402 memory_type: row.memory_type,
403 description: row.description,
404 snippet,
405 distance: graph_distance,
406 score: RecallItem::score_from_distance(graph_distance),
407 source: "graph".to_string(),
408 graph_depth: Some(hop),
409 });
410 }
411 }
412 }
413 }
414
415 output::emit_json(&HybridSearchResponse {
416 query: args.query,
417 k: args.k,
418 rrf_k: args.rrf_k,
419 weights: Weights {
420 vec: args.weight_vec,
421 fts: args.weight_fts,
422 },
423 results,
424 graph_matches,
425 fts_degraded,
426 fts_error,
427 fts_auto_rebuilt,
428 vec_degraded,
429 vec_error,
430 warning: if vec_degraded {
431 Some(
432 "live query embedding unavailable; results are FTS5 BM25 only (semantic relevance reduced)"
433 .to_string(),
434 )
435 } else {
436 None
437 },
438 elapsed_ms: start.elapsed().as_millis() as u64,
439 })?;
440
441 Ok(())
442}
443
444#[cfg(test)]
445mod tests {
446 use super::*;
447
448 #[derive(clap::Parser)]
449 struct TestCli {
450 #[command(flatten)]
451 args: HybridSearchArgs,
452 }
453
454 #[test]
455 fn graph_flags_parse_as_none_when_absent() {
456 use clap::Parser;
460 let cli = TestCli::try_parse_from(["hybrid-search", "q"]).expect("bare query parses");
461 assert!(cli.args.max_hops.is_none());
462 assert!(cli.args.min_weight.is_none());
463 let cli = TestCli::try_parse_from(["hybrid-search", "q", "--max-hops", "2"])
464 .expect("explicit flag parses");
465 assert_eq!(cli.args.max_hops, Some(2));
466 }
467
468 fn empty_response(
469 k: usize,
470 rrf_k: u32,
471 weight_vec: f32,
472 weight_fts: f32,
473 ) -> HybridSearchResponse {
474 HybridSearchResponse {
475 query: "test query".to_string(),
476 k,
477 rrf_k,
478 weights: Weights {
479 vec: weight_vec,
480 fts: weight_fts,
481 },
482 results: vec![],
483 graph_matches: vec![],
484 fts_degraded: false,
485 fts_error: None,
486 fts_auto_rebuilt: false,
487 vec_degraded: false,
488 vec_error: None,
489 warning: None,
490 elapsed_ms: 0,
491 }
492 }
493
494 #[test]
495 fn hybrid_search_response_empty_serializes_correct_fields() {
496 let resp = empty_response(10, 60, 1.0, 1.0);
497 let json = serde_json::to_string(&resp).unwrap();
498 assert!(json.contains("\"results\""), "must contain results field");
499 assert!(json.contains("\"query\""), "must contain query field");
500 assert!(json.contains("\"k\""), "must contain k field");
501 assert!(
502 json.contains("\"graph_matches\""),
503 "must contain graph_matches field"
504 );
505 assert!(
506 !json.contains("\"combined_rank\""),
507 "must not contain combined_rank"
508 );
509 assert!(
510 !json.contains("\"vec_rank_list\""),
511 "must not contain vec_rank_list"
512 );
513 assert!(
514 !json.contains("\"fts_rank_list\""),
515 "must not contain fts_rank_list"
516 );
517 }
518
519 #[test]
520 fn hybrid_search_response_serializes_rrf_k_and_weights() {
521 let resp = empty_response(5, 60, 0.7, 0.3);
522 let json = serde_json::to_string(&resp).unwrap();
523 assert!(json.contains("\"rrf_k\""), "must contain rrf_k field");
524 assert!(json.contains("\"weights\""), "must contain weights field");
525 assert!(json.contains("\"vec\""), "must contain weights.vec field");
526 assert!(json.contains("\"fts\""), "must contain weights.fts field");
527 }
528
529 #[test]
530 fn hybrid_search_response_serializes_elapsed_ms() {
531 let mut resp = empty_response(5, 60, 1.0, 1.0);
532 resp.elapsed_ms = 123;
533 let json = serde_json::to_string(&resp).unwrap();
534 assert!(
535 json.contains("\"elapsed_ms\""),
536 "must contain elapsed_ms field"
537 );
538 assert!(json.contains("123"), "deve serializar valor de elapsed_ms");
539 }
540
541 #[test]
542 fn weights_struct_serializes_correctly() {
543 let w = Weights { vec: 0.6, fts: 0.4 };
544 let json = serde_json::to_string(&w).unwrap();
545 assert!(json.contains("\"vec\""));
546 assert!(json.contains("\"fts\""));
547 }
548
549 #[test]
550 fn hybrid_search_item_omits_fts_rank_when_none() {
551 let item = HybridSearchItem {
552 memory_id: 1,
553 name: "mem".to_string(),
554 namespace: "default".to_string(),
555 memory_type: "user".to_string(),
556 description: "desc".to_string(),
557 body: "content".to_string(),
558 snippet: "content".to_string(),
559 combined_score: 0.0328,
560 score: 0.0328,
561 source: "hybrid".to_string(),
562 vec_rank: Some(1),
563 fts_rank: None,
564 rrf_score: Some(0.0328),
565 normalized_score: 1.0,
566 vec_distance: Some(0.12),
567 fts_bm25: None,
568 };
569 let json = serde_json::to_string(&item).unwrap();
570 assert!(
571 json.contains("\"vec_rank\""),
572 "must contain vec_rank when Some"
573 );
574 assert!(
575 !json.contains("\"fts_rank\""),
576 "must not contain fts_rank when None"
577 );
578 }
579
580 #[test]
581 fn hybrid_search_item_omits_vec_rank_when_none() {
582 let item = HybridSearchItem {
583 memory_id: 2,
584 name: "mem2".to_string(),
585 namespace: "default".to_string(),
586 memory_type: "fact".to_string(),
587 description: "desc2".to_string(),
588 body: "corpo2".to_string(),
589 snippet: "corpo2".to_string(),
590 combined_score: 0.016,
591 score: 0.016,
592 source: "hybrid".to_string(),
593 vec_rank: None,
594 fts_rank: Some(2),
595 rrf_score: Some(0.016),
596 normalized_score: 0.5,
597 vec_distance: None,
598 fts_bm25: None,
599 };
600 let json = serde_json::to_string(&item).unwrap();
601 assert!(
602 !json.contains("\"vec_rank\""),
603 "must not contain vec_rank when None"
604 );
605 assert!(
606 json.contains("\"fts_rank\""),
607 "must contain fts_rank when Some"
608 );
609 }
610
611 #[test]
612 fn hybrid_search_item_serializes_both_ranks_when_some() {
613 let item = HybridSearchItem {
614 memory_id: 3,
615 name: "mem3".to_string(),
616 namespace: "ns".to_string(),
617 memory_type: "entity".to_string(),
618 description: "desc3".to_string(),
619 body: "corpo3".to_string(),
620 snippet: "corpo3".to_string(),
621 combined_score: 0.05,
622 score: 0.05,
623 source: "hybrid".to_string(),
624 vec_rank: Some(3),
625 fts_rank: Some(1),
626 rrf_score: Some(0.05),
627 normalized_score: 0.8,
628 vec_distance: Some(0.25),
629 fts_bm25: None,
630 };
631 let json = serde_json::to_string(&item).unwrap();
632 assert!(json.contains("\"vec_rank\""), "must contain vec_rank");
633 assert!(json.contains("\"fts_rank\""), "must contain fts_rank");
634 assert!(json.contains("\"type\""), "deve serializar type renomeado");
635 assert!(!json.contains("memory_type"), "must not expose memory_type");
636 }
637
638 #[test]
639 fn hybrid_search_response_serializes_k_correctly() {
640 let resp = empty_response(5, 60, 1.0, 1.0);
641 let json = serde_json::to_string(&resp).unwrap();
642 assert!(json.contains("\"k\":5"), "deve serializar k=5");
643 }
644
645 #[test]
646 fn hybrid_search_response_with_graph_matches() {
647 use crate::output::RecallItem;
648 let resp = HybridSearchResponse {
649 query: "test".to_string(),
650 k: 5,
651 rrf_k: 60,
652 weights: Weights { vec: 1.0, fts: 1.0 },
653 results: vec![],
654 graph_matches: vec![RecallItem {
655 memory_id: 1,
656 name: "graph-hit".to_string(),
657 namespace: "global".to_string(),
658 memory_type: "document".to_string(),
659 description: "found via graph".to_string(),
660 snippet: "graph content".to_string(),
661 distance: 0.1,
662 score: 0.9,
663 source: "graph".to_string(),
664 graph_depth: Some(1),
665 }],
666 fts_degraded: false,
667 fts_error: None,
668 fts_auto_rebuilt: false,
669 vec_degraded: false,
670 vec_error: None,
671 warning: None,
672 elapsed_ms: 42,
673 };
674 let json = serde_json::to_value(&resp).unwrap();
675 assert_eq!(json["graph_matches"].as_array().unwrap().len(), 1);
676 assert_eq!(json["graph_matches"][0]["source"], "graph");
677 assert_eq!(json["graph_matches"][0]["graph_depth"], 1);
678 }
679
680 #[test]
681 fn fts_degraded_omitted_on_success_present_on_failure() {
682 let ok_resp = empty_response(5, 60, 1.0, 1.0);
684 let ok_json = serde_json::to_string(&ok_resp).unwrap();
685 assert!(
686 !ok_json.contains("\"fts_degraded\""),
687 "fts_degraded must be absent when false"
688 );
689 assert!(
690 !ok_json.contains("\"fts_error\""),
691 "fts_error must be absent when None"
692 );
693
694 let mut degraded_resp = empty_response(5, 60, 1.0, 1.0);
696 degraded_resp.fts_degraded = true;
697 degraded_resp.fts_error = Some("FTS5 table corrupted".to_string());
698 let degraded_json = serde_json::to_string(°raded_resp).unwrap();
699 assert!(
700 degraded_json.contains("\"fts_degraded\":true"),
701 "fts_degraded must be present and true when degraded"
702 );
703 assert!(
704 degraded_json.contains("\"fts_error\""),
705 "fts_error must be present when Some"
706 );
707 assert!(
708 degraded_json.contains("FTS5 table corrupted"),
709 "fts_error must contain the error message"
710 );
711 }
712}