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,
41 #[arg(short = 'k', long, aliases = ["limit", "top-k"], default_value = "10", value_parser = crate::parsers::parse_k_range)]
46 pub k: usize,
47 #[arg(long, default_value = "60")]
49 pub rrf_k: u32,
50 #[arg(long, default_value = "1.0")]
52 pub weight_vec: f32,
53 #[arg(long, default_value = "1.0")]
55 pub weight_fts: f32,
56 #[arg(long, value_enum)]
60 pub r#type: Option<MemoryType>,
61 #[arg(long)]
63 pub namespace: Option<String>,
64 #[arg(long)]
66 pub with_graph: bool,
67 #[arg(long, help = "Skip live query embedding; serve FTS5 BM25 only")]
70 pub fallback_fts_only: bool,
71 #[arg(long)]
73 pub max_hops: Option<u32>,
74 #[arg(long)]
76 pub min_weight: Option<f64>,
77 #[arg(long, value_enum, default_value_t = JsonOutputFormat::Json)]
79 pub format: JsonOutputFormat,
80 #[arg(long)]
82 pub db: Option<String>,
83 #[arg(long, hide = true, help = "No-op; JSON is always emitted on stdout")]
85 pub json: bool,
86}
87
88#[derive(serde::Serialize)]
90pub struct HybridSearchItem {
91 pub memory_id: i64,
93 pub name: String,
95 pub namespace: String,
97 #[serde(rename = "type")]
99 pub memory_type: String,
100 pub description: String,
102 pub body: String,
104 pub snippet: String,
106 pub combined_score: f64,
108 pub score: f64,
110 pub source: String,
112 #[serde(skip_serializing_if = "Option::is_none")]
114 pub vec_rank: Option<usize>,
115 #[serde(skip_serializing_if = "Option::is_none")]
117 pub fts_rank: Option<usize>,
118 #[serde(skip_serializing_if = "Option::is_none")]
120 pub rrf_score: Option<f64>,
121 pub normalized_score: f64,
123 #[serde(skip_serializing_if = "Option::is_none")]
128 pub vec_distance: Option<f64>,
129 #[serde(skip_serializing_if = "Option::is_none")]
132 pub fts_bm25: Option<f64>,
133}
134
135#[derive(serde::Serialize)]
137pub struct Weights {
138 pub vec: f32,
140 pub fts: f32,
142}
143
144#[derive(serde::Serialize)]
146pub struct HybridSearchResponse {
147 pub query: String,
149 pub k: usize,
151 pub rrf_k: u32,
153 pub weights: Weights,
155 pub results: Vec<HybridSearchItem>,
157 pub graph_matches: Vec<RecallItem>,
159 #[serde(skip_serializing_if = "std::ops::Not::not")]
163 pub fts_degraded: bool,
164 #[serde(skip_serializing_if = "Option::is_none")]
168 pub fts_error: Option<String>,
169 #[serde(skip_serializing_if = "std::ops::Not::not")]
173 pub fts_auto_rebuilt: bool,
174 #[serde(skip_serializing_if = "std::ops::Not::not", default)]
178 pub vec_degraded: bool,
179 #[serde(skip_serializing_if = "Option::is_none")]
183 pub vec_error: Option<String>,
184 #[serde(skip_serializing_if = "Option::is_none")]
188 pub warning: Option<String>,
189 #[serde(skip_serializing_if = "Option::is_none")]
193 pub backend_invoked: Option<&'static str>,
194 #[serde(skip_serializing_if = "Option::is_none")]
198 pub vec_degraded_reason: Option<String>,
199 pub elapsed_ms: u64,
201}
202
203#[tracing::instrument(skip_all, level = "debug", name = "hybrid_search")]
205pub fn run(
206 args: HybridSearchArgs,
207 llm_backend: crate::cli::LlmBackendChoice,
208 embedding_backend: crate::cli::EmbeddingBackendChoice,
209) -> Result<(), AppError> {
210 let start = std::time::Instant::now();
211 let _ = args.format;
212 tracing::debug!(target: "hybrid_search", query = %args.query, k = args.k, "fusing results");
213
214 if !args.with_graph {
218 if args.max_hops.is_some() {
219 return Err(AppError::Validation(
220 "--max-hops requires --with-graph to be active".to_string(),
221 ));
222 }
223 if args.min_weight.is_some() {
224 return Err(AppError::Validation(
225 "--min-weight requires --with-graph to be active".to_string(),
226 ));
227 }
228 }
229
230 let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
231 let paths = AppPaths::resolve(args.db.as_deref())?;
232 crate::storage::connection::ensure_db_ready(&paths)?;
233
234 output::emit_progress_i18n(
235 "Computing query embedding...",
236 "Calculando embedding da consulta...",
237 );
238 let conn = open_ro(&paths.db)?;
239 let (embedding, vec_degraded, vec_error, backend_invoked) = if args.fallback_fts_only {
247 (
248 None,
249 true,
250 Some("fallback_fts_only requested".to_string()),
251 None,
252 )
253 } else {
254 match crate::embedder::try_embed_query_with_embedding_choice(
261 &paths.models,
262 &args.query,
263 embedding_backend,
264 llm_backend,
265 ) {
266 Ok((v, backend)) => (Some(v), false, None, Some(backend.as_str())),
267 Err(reason) => {
268 let msg = reason.to_string();
269 tracing::warn!(target: "hybrid_search", fallback_reason = %msg, reason_code = %reason.reason_code(), "live embedding failed; falling back to FTS5");
270 (None, true, Some(msg), None)
271 }
272 }
273 };
274
275 let memory_type_str = args.r#type.map(|t| t.as_str());
276
277 let vec_results: Vec<(i64, f32)> = if let Some(emb) = embedding.as_ref() {
278 memories::knn_search(
279 &conn,
280 emb,
281 std::slice::from_ref(&namespace),
282 memory_type_str,
283 args.k * 2,
284 )?
285 } else {
286 Vec::new()
287 };
288
289 let vec_rank_map: HashMap<i64, usize> = vec_results
291 .iter()
292 .enumerate()
293 .map(|(pos, (id, _))| (*id, pos + 1))
294 .collect();
295
296 let vec_distance_map: HashMap<i64, f64> = vec_results
298 .iter()
299 .map(|(id, dist)| (*id, *dist as f64))
300 .collect();
301
302 let (fts_results, fts_degraded, fts_error, fts_auto_rebuilt) = if args.weight_fts == 0.0 {
303 (vec![], false, None, false)
304 } else {
305 match memories::fts_search(&conn, &args.query, &namespace, memory_type_str, args.k * 2) {
306 Ok(r) => (r, false, None, false),
307 Err(e) => {
308 let err_msg = e.to_string();
309 let is_malformed = err_msg.contains("malformed") || err_msg.contains("corrupt");
310 if is_malformed {
311 tracing::warn!(target: "hybrid_search", "FTS5 index corrupted, attempting auto-rebuild");
312 if conn
313 .execute_batch("INSERT INTO fts_memories(fts_memories) VALUES('rebuild');")
314 .is_ok()
315 {
316 match memories::fts_search(
317 &conn,
318 &args.query,
319 &namespace,
320 memory_type_str,
321 args.k * 2,
322 ) {
323 Ok(r) => (r, false, None, true),
324 Err(e2) => {
325 tracing::error!(target: "hybrid_search", error = %e2, "FTS5 auto-rebuild failed to recover");
326 (vec![], true, Some(e2.to_string()), true)
327 }
328 }
329 } else {
330 (vec![], true, Some(err_msg), false)
331 }
332 } else {
333 tracing::warn!(target: "hybrid_search", error = %e, "FTS5 query failed, falling back to vec-only");
334 (vec![], true, Some(err_msg), false)
335 }
336 }
337 }
338 };
339
340 let fts_rank_map: HashMap<i64, usize> = fts_results
342 .iter()
343 .enumerate()
344 .map(|(pos, row)| (row.id, pos + 1))
345 .collect();
346
347 let rrf_k = args.rrf_k as f64;
348
349 let mut combined_scores: crate::hash::AHashMap<i64, f64> =
351 crate::hash::AHashMap::with_capacity_and_hasher(
352 vec_results.len() + fts_results.len(),
353 Default::default(),
354 );
355
356 for (rank, (memory_id, _)) in vec_results.iter().enumerate() {
357 let score = args.weight_vec as f64 * (1.0 / (rrf_k + rank as f64 + 1.0));
358 *combined_scores.entry(*memory_id).or_insert(0.0) += score;
359 }
360
361 for (rank, row) in fts_results.iter().enumerate() {
362 let score = args.weight_fts as f64 * (1.0 / (rrf_k + rank as f64 + 1.0));
363 *combined_scores.entry(row.id).or_insert(0.0) += score;
364 }
365
366 let mut ranked: Vec<(i64, f64)> = combined_scores.into_iter().collect();
368 ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
369 ranked.truncate(args.k);
370
371 let top_ids: Vec<i64> = ranked.iter().map(|(id, _)| *id).collect();
373
374 let mut memory_data: crate::hash::AHashMap<i64, memories::MemoryRow> =
376 crate::hash::AHashMap::with_capacity_and_hasher(ranked.len(), Default::default());
377 for id in &top_ids {
378 if let Some(row) = memories::read_full(&conn, *id)? {
379 memory_data.insert(*id, row);
380 }
381 }
382
383 let max_possible = args.weight_vec as f64 * (1.0 / (rrf_k + 1.0))
384 + args.weight_fts as f64 * (1.0 / (rrf_k + 1.0));
385
386 let results: Vec<HybridSearchItem> = ranked
388 .into_iter()
389 .filter_map(|(memory_id, combined_score)| {
390 let normalized_score = if max_possible > 0.0 {
391 combined_score / max_possible
392 } else {
393 0.0
394 };
395 memory_data.remove(&memory_id).map(|row| {
396 let snippet: String = row.body.chars().take(300).collect();
397 HybridSearchItem {
398 memory_id: row.id,
399 name: row.name,
400 namespace: row.namespace,
401 memory_type: row.memory_type,
402 description: row.description,
403 body: row.body,
404 snippet,
405 combined_score,
406 score: combined_score,
407 source: "hybrid".to_string(),
408 vec_rank: vec_rank_map.get(&memory_id).copied(),
409 fts_rank: fts_rank_map.get(&memory_id).copied(),
410 rrf_score: Some(combined_score),
411 normalized_score,
412 vec_distance: vec_distance_map.get(&memory_id).copied(),
413 fts_bm25: None,
414 }
415 })
416 })
417 .collect();
418
419 let mut graph_matches: Vec<RecallItem> = Vec::with_capacity(8);
421 if let Some(emb) = args
422 .with_graph
423 .then_some(())
424 .filter(|_| !results.is_empty())
425 .and(embedding.as_ref())
426 {
427 let namespace_for_graph = namespace.clone();
428 let memory_ids: Vec<i64> = results.iter().map(|r| r.memory_id).collect();
429
430 let entity_knn = entities::knn_search(&conn, emb, &namespace_for_graph, 5)?;
431 let entity_ids: Vec<i64> = entity_knn.iter().map(|(id, _)| *id).collect();
432
433 let all_seed_ids: Vec<i64> = memory_ids
434 .iter()
435 .chain(entity_ids.iter())
436 .copied()
437 .collect();
438
439 if !all_seed_ids.is_empty() {
440 let graph_memory_ids = traverse_from_memories_with_hops(
441 &conn,
442 &all_seed_ids,
443 &namespace_for_graph,
444 args.min_weight.unwrap_or(0.3),
445 args.max_hops.unwrap_or(2),
446 )?;
447
448 let already_in_results: std::collections::HashSet<i64> =
449 results.iter().map(|r| r.memory_id).collect();
450
451 for (graph_mem_id, hop) in graph_memory_ids {
452 if already_in_results.contains(&graph_mem_id) {
453 continue;
454 }
455 if let Some(row) = memories::read_full(&conn, graph_mem_id)? {
456 let snippet: String = row.body.chars().take(300).collect();
457 let graph_distance = 1.0 - 1.0 / (hop as f32 + 1.0);
458 graph_matches.push(RecallItem {
459 memory_id: row.id,
460 name: row.name,
461 namespace: row.namespace,
462 memory_type: row.memory_type,
463 description: row.description,
464 snippet,
465 distance: graph_distance,
466 score: RecallItem::score_from_distance(graph_distance),
467 source: "graph".to_string(),
468 graph_depth: Some(hop),
469 });
470 }
471 }
472 }
473 }
474
475 output::emit_json(&HybridSearchResponse {
476 query: args.query,
477 k: args.k,
478 rrf_k: args.rrf_k,
479 weights: Weights {
480 vec: args.weight_vec,
481 fts: args.weight_fts,
482 },
483 results,
484 graph_matches,
485 fts_degraded,
486 fts_error,
487 fts_auto_rebuilt,
488 vec_degraded,
489 vec_error: vec_error.clone(),
490 warning: if vec_degraded {
491 Some(
492 "live query embedding unavailable; results are FTS5 BM25 only (semantic relevance reduced)"
493 .to_string(),
494 )
495 } else {
496 None
497 },
498 backend_invoked,
499 vec_degraded_reason: if vec_degraded { vec_error } else { None },
500 elapsed_ms: start.elapsed().as_millis() as u64,
501 })?;
502
503 Ok(())
504}
505
506#[cfg(test)]
507mod tests {
508 use super::*;
509
510 #[derive(clap::Parser)]
511 struct TestCli {
512 #[command(flatten)]
513 args: HybridSearchArgs,
514 }
515
516 #[test]
517 fn graph_flags_parse_as_none_when_absent() {
518 use clap::Parser;
522 let cli = TestCli::try_parse_from(["hybrid-search", "q"]).expect("bare query parses");
523 assert!(cli.args.max_hops.is_none());
524 assert!(cli.args.min_weight.is_none());
525 let cli = TestCli::try_parse_from(["hybrid-search", "q", "--max-hops", "2"])
526 .expect("explicit flag parses");
527 assert_eq!(cli.args.max_hops, Some(2));
528 }
529
530 fn empty_response(
531 k: usize,
532 rrf_k: u32,
533 weight_vec: f32,
534 weight_fts: f32,
535 ) -> HybridSearchResponse {
536 HybridSearchResponse {
537 query: "test query".to_string(),
538 k,
539 rrf_k,
540 weights: Weights {
541 vec: weight_vec,
542 fts: weight_fts,
543 },
544 results: vec![],
545 graph_matches: vec![],
546 fts_degraded: false,
547 fts_error: None,
548 fts_auto_rebuilt: false,
549 vec_degraded: false,
550 vec_error: None,
551 warning: None,
552 backend_invoked: None,
553 vec_degraded_reason: None,
554 elapsed_ms: 0,
555 }
556 }
557
558 #[test]
559 fn hybrid_search_response_empty_serializes_correct_fields() {
560 let resp = empty_response(10, 60, 1.0, 1.0);
561 let json = serde_json::to_string(&resp).unwrap();
562 assert!(json.contains("\"results\""), "must contain results field");
563 assert!(json.contains("\"query\""), "must contain query field");
564 assert!(json.contains("\"k\""), "must contain k field");
565 assert!(
566 json.contains("\"graph_matches\""),
567 "must contain graph_matches field"
568 );
569 assert!(
570 !json.contains("\"combined_rank\""),
571 "must not contain combined_rank"
572 );
573 assert!(
574 !json.contains("\"vec_rank_list\""),
575 "must not contain vec_rank_list"
576 );
577 assert!(
578 !json.contains("\"fts_rank_list\""),
579 "must not contain fts_rank_list"
580 );
581 }
582
583 #[test]
584 fn hybrid_search_response_serializes_rrf_k_and_weights() {
585 let resp = empty_response(5, 60, 0.7, 0.3);
586 let json = serde_json::to_string(&resp).unwrap();
587 assert!(json.contains("\"rrf_k\""), "must contain rrf_k field");
588 assert!(json.contains("\"weights\""), "must contain weights field");
589 assert!(json.contains("\"vec\""), "must contain weights.vec field");
590 assert!(json.contains("\"fts\""), "must contain weights.fts field");
591 }
592
593 #[test]
594 fn hybrid_search_response_serializes_elapsed_ms() {
595 let mut resp = empty_response(5, 60, 1.0, 1.0);
596 resp.elapsed_ms = 123;
597 let json = serde_json::to_string(&resp).unwrap();
598 assert!(
599 json.contains("\"elapsed_ms\""),
600 "must contain elapsed_ms field"
601 );
602 assert!(json.contains("123"), "deve serializar valor de elapsed_ms");
603 }
604
605 #[test]
606 fn weights_struct_serializes_correctly() {
607 let w = Weights { vec: 0.6, fts: 0.4 };
608 let json = serde_json::to_string(&w).unwrap();
609 assert!(json.contains("\"vec\""));
610 assert!(json.contains("\"fts\""));
611 }
612
613 #[test]
614 fn hybrid_search_item_omits_fts_rank_when_none() {
615 let item = HybridSearchItem {
616 memory_id: 1,
617 name: "mem".to_string(),
618 namespace: "default".to_string(),
619 memory_type: "user".to_string(),
620 description: "desc".to_string(),
621 body: "content".to_string(),
622 snippet: "content".to_string(),
623 combined_score: 0.0328,
624 score: 0.0328,
625 source: "hybrid".to_string(),
626 vec_rank: Some(1),
627 fts_rank: None,
628 rrf_score: Some(0.0328),
629 normalized_score: 1.0,
630 vec_distance: Some(0.12),
631 fts_bm25: None,
632 };
633 let json = serde_json::to_string(&item).unwrap();
634 assert!(
635 json.contains("\"vec_rank\""),
636 "must contain vec_rank when Some"
637 );
638 assert!(
639 !json.contains("\"fts_rank\""),
640 "must not contain fts_rank when None"
641 );
642 }
643
644 #[test]
645 fn hybrid_search_item_omits_vec_rank_when_none() {
646 let item = HybridSearchItem {
647 memory_id: 2,
648 name: "mem2".to_string(),
649 namespace: "default".to_string(),
650 memory_type: "fact".to_string(),
651 description: "desc2".to_string(),
652 body: "corpo2".to_string(),
653 snippet: "corpo2".to_string(),
654 combined_score: 0.016,
655 score: 0.016,
656 source: "hybrid".to_string(),
657 vec_rank: None,
658 fts_rank: Some(2),
659 rrf_score: Some(0.016),
660 normalized_score: 0.5,
661 vec_distance: None,
662 fts_bm25: None,
663 };
664 let json = serde_json::to_string(&item).unwrap();
665 assert!(
666 !json.contains("\"vec_rank\""),
667 "must not contain vec_rank when None"
668 );
669 assert!(
670 json.contains("\"fts_rank\""),
671 "must contain fts_rank when Some"
672 );
673 }
674
675 #[test]
676 fn hybrid_search_item_serializes_both_ranks_when_some() {
677 let item = HybridSearchItem {
678 memory_id: 3,
679 name: "mem3".to_string(),
680 namespace: "ns".to_string(),
681 memory_type: "entity".to_string(),
682 description: "desc3".to_string(),
683 body: "corpo3".to_string(),
684 snippet: "corpo3".to_string(),
685 combined_score: 0.05,
686 score: 0.05,
687 source: "hybrid".to_string(),
688 vec_rank: Some(3),
689 fts_rank: Some(1),
690 rrf_score: Some(0.05),
691 normalized_score: 0.8,
692 vec_distance: Some(0.25),
693 fts_bm25: None,
694 };
695 let json = serde_json::to_string(&item).unwrap();
696 assert!(json.contains("\"vec_rank\""), "must contain vec_rank");
697 assert!(json.contains("\"fts_rank\""), "must contain fts_rank");
698 assert!(json.contains("\"type\""), "deve serializar type renomeado");
699 assert!(!json.contains("memory_type"), "must not expose memory_type");
700 }
701
702 #[test]
703 fn hybrid_search_response_serializes_k_correctly() {
704 let resp = empty_response(5, 60, 1.0, 1.0);
705 let json = serde_json::to_string(&resp).unwrap();
706 assert!(json.contains("\"k\":5"), "deve serializar k=5");
707 }
708
709 #[test]
710 fn hybrid_search_response_with_graph_matches() {
711 use crate::output::RecallItem;
712 let resp = HybridSearchResponse {
713 query: "test".to_string(),
714 k: 5,
715 rrf_k: 60,
716 weights: Weights { vec: 1.0, fts: 1.0 },
717 results: vec![],
718 graph_matches: vec![RecallItem {
719 memory_id: 1,
720 name: "graph-hit".to_string(),
721 namespace: "global".to_string(),
722 memory_type: "document".to_string(),
723 description: "found via graph".to_string(),
724 snippet: "graph content".to_string(),
725 distance: 0.1,
726 score: 0.9,
727 source: "graph".to_string(),
728 graph_depth: Some(1),
729 }],
730 fts_degraded: false,
731 fts_error: None,
732 fts_auto_rebuilt: false,
733 vec_degraded: false,
734 vec_error: None,
735 warning: None,
736 backend_invoked: None,
737 vec_degraded_reason: None,
738 elapsed_ms: 42,
739 };
740 let json = serde_json::to_value(&resp).unwrap();
741 assert_eq!(json["graph_matches"].as_array().unwrap().len(), 1);
742 assert_eq!(json["graph_matches"][0]["source"], "graph");
743 assert_eq!(json["graph_matches"][0]["graph_depth"], 1);
744 }
745
746 #[test]
747 fn fts_degraded_omitted_on_success_present_on_failure() {
748 let ok_resp = empty_response(5, 60, 1.0, 1.0);
750 let ok_json = serde_json::to_string(&ok_resp).unwrap();
751 assert!(
752 !ok_json.contains("\"fts_degraded\""),
753 "fts_degraded must be absent when false"
754 );
755 assert!(
756 !ok_json.contains("\"fts_error\""),
757 "fts_error must be absent when None"
758 );
759
760 let mut degraded_resp = empty_response(5, 60, 1.0, 1.0);
762 degraded_resp.fts_degraded = true;
763 degraded_resp.fts_error = Some("FTS5 table corrupted".to_string());
764 let degraded_json = serde_json::to_string(°raded_resp).unwrap();
765 assert!(
766 degraded_json.contains("\"fts_degraded\":true"),
767 "fts_degraded must be present and true when degraded"
768 );
769 assert!(
770 degraded_json.contains("\"fts_error\""),
771 "fts_error must be present when Some"
772 );
773 assert!(
774 degraded_json.contains("FTS5 table corrupted"),
775 "fts_error must contain the error message"
776 );
777 }
778}