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(
169 args: HybridSearchArgs,
170 llm_backend: crate::cli::LlmBackendChoice,
171) -> Result<(), AppError> {
172 let start = std::time::Instant::now();
173 let _ = args.format;
174 tracing::debug!(target: "hybrid_search", query = %args.query, k = args.k, "fusing results");
175
176 if !args.with_graph {
180 if args.max_hops.is_some() {
181 return Err(AppError::Validation(
182 "--max-hops requires --with-graph to be active".to_string(),
183 ));
184 }
185 if args.min_weight.is_some() {
186 return Err(AppError::Validation(
187 "--min-weight requires --with-graph to be active".to_string(),
188 ));
189 }
190 }
191
192 let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
193 let paths = AppPaths::resolve(args.db.as_deref())?;
194 crate::storage::connection::ensure_db_ready(&paths)?;
195
196 output::emit_progress_i18n(
197 "Computing query embedding...",
198 "Calculando embedding da consulta...",
199 );
200 let conn = open_ro(&paths.db)?;
201 let (embedding, vec_degraded, vec_error) = if args.fallback_fts_only {
206 (None, true, Some("fallback_fts_only requested".to_string()))
207 } else {
208 match crate::embedder::try_embed_query_with_choice(
210 &paths.models,
211 &args.query,
212 Some(llm_backend),
213 ) {
214 Ok(v) => (Some(v), false, None),
215 Err(reason) => {
216 let msg = reason.to_string();
217 tracing::warn!(target: "hybrid_search", fallback_reason = %msg, "live embedding failed; falling back to FTS5");
218 (None, true, Some(msg))
219 }
220 }
221 };
222
223 let memory_type_str = args.r#type.map(|t| t.as_str());
224
225 let vec_results: Vec<(i64, f32)> = if let Some(emb) = embedding.as_ref() {
226 memories::knn_search(
227 &conn,
228 emb,
229 std::slice::from_ref(&namespace),
230 memory_type_str,
231 args.k * 2,
232 )?
233 } else {
234 Vec::new()
235 };
236
237 let vec_rank_map: HashMap<i64, usize> = vec_results
239 .iter()
240 .enumerate()
241 .map(|(pos, (id, _))| (*id, pos + 1))
242 .collect();
243
244 let vec_distance_map: HashMap<i64, f64> = vec_results
246 .iter()
247 .map(|(id, dist)| (*id, *dist as f64))
248 .collect();
249
250 let (fts_results, fts_degraded, fts_error, fts_auto_rebuilt) = if args.weight_fts == 0.0 {
251 (vec![], false, None, false)
252 } else {
253 match memories::fts_search(&conn, &args.query, &namespace, memory_type_str, args.k * 2) {
254 Ok(r) => (r, false, None, false),
255 Err(e) => {
256 let err_msg = e.to_string();
257 let is_malformed = err_msg.contains("malformed") || err_msg.contains("corrupt");
258 if is_malformed {
259 tracing::warn!(target: "hybrid_search", "FTS5 index corrupted, attempting auto-rebuild");
260 if conn
261 .execute_batch("INSERT INTO fts_memories(fts_memories) VALUES('rebuild');")
262 .is_ok()
263 {
264 match memories::fts_search(
265 &conn,
266 &args.query,
267 &namespace,
268 memory_type_str,
269 args.k * 2,
270 ) {
271 Ok(r) => (r, false, None, true),
272 Err(e2) => {
273 tracing::error!(target: "hybrid_search", error = %e2, "FTS5 auto-rebuild failed to recover");
274 (vec![], true, Some(e2.to_string()), true)
275 }
276 }
277 } else {
278 (vec![], true, Some(err_msg), false)
279 }
280 } else {
281 tracing::warn!(target: "hybrid_search", error = %e, "FTS5 query failed, falling back to vec-only");
282 (vec![], true, Some(err_msg), false)
283 }
284 }
285 }
286 };
287
288 let fts_rank_map: HashMap<i64, usize> = fts_results
290 .iter()
291 .enumerate()
292 .map(|(pos, row)| (row.id, pos + 1))
293 .collect();
294
295 let rrf_k = args.rrf_k as f64;
296
297 let mut combined_scores: crate::hash::AHashMap<i64, f64> =
299 crate::hash::AHashMap::with_capacity_and_hasher(
300 vec_results.len() + fts_results.len(),
301 Default::default(),
302 );
303
304 for (rank, (memory_id, _)) in vec_results.iter().enumerate() {
305 let score = args.weight_vec as f64 * (1.0 / (rrf_k + rank as f64 + 1.0));
306 *combined_scores.entry(*memory_id).or_insert(0.0) += score;
307 }
308
309 for (rank, row) in fts_results.iter().enumerate() {
310 let score = args.weight_fts as f64 * (1.0 / (rrf_k + rank as f64 + 1.0));
311 *combined_scores.entry(row.id).or_insert(0.0) += score;
312 }
313
314 let mut ranked: Vec<(i64, f64)> = combined_scores.into_iter().collect();
316 ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
317 ranked.truncate(args.k);
318
319 let top_ids: Vec<i64> = ranked.iter().map(|(id, _)| *id).collect();
321
322 let mut memory_data: crate::hash::AHashMap<i64, memories::MemoryRow> =
324 crate::hash::AHashMap::with_capacity_and_hasher(ranked.len(), Default::default());
325 for id in &top_ids {
326 if let Some(row) = memories::read_full(&conn, *id)? {
327 memory_data.insert(*id, row);
328 }
329 }
330
331 let max_possible = args.weight_vec as f64 * (1.0 / (rrf_k + 1.0))
332 + args.weight_fts as f64 * (1.0 / (rrf_k + 1.0));
333
334 let results: Vec<HybridSearchItem> = ranked
336 .into_iter()
337 .filter_map(|(memory_id, combined_score)| {
338 let normalized_score = if max_possible > 0.0 {
339 combined_score / max_possible
340 } else {
341 0.0
342 };
343 memory_data.remove(&memory_id).map(|row| {
344 let snippet: String = row.body.chars().take(300).collect();
345 HybridSearchItem {
346 memory_id: row.id,
347 name: row.name,
348 namespace: row.namespace,
349 memory_type: row.memory_type,
350 description: row.description,
351 body: row.body,
352 snippet,
353 combined_score,
354 score: combined_score,
355 source: "hybrid".to_string(),
356 vec_rank: vec_rank_map.get(&memory_id).copied(),
357 fts_rank: fts_rank_map.get(&memory_id).copied(),
358 rrf_score: Some(combined_score),
359 normalized_score,
360 vec_distance: vec_distance_map.get(&memory_id).copied(),
361 fts_bm25: None,
362 }
363 })
364 })
365 .collect();
366
367 let mut graph_matches: Vec<RecallItem> = Vec::with_capacity(8);
369 if let Some(emb) = args
370 .with_graph
371 .then_some(())
372 .filter(|_| !results.is_empty())
373 .and(embedding.as_ref())
374 {
375 let namespace_for_graph = namespace.clone();
376 let memory_ids: Vec<i64> = results.iter().map(|r| r.memory_id).collect();
377
378 let entity_knn = entities::knn_search(&conn, emb, &namespace_for_graph, 5)?;
379 let entity_ids: Vec<i64> = entity_knn.iter().map(|(id, _)| *id).collect();
380
381 let all_seed_ids: Vec<i64> = memory_ids
382 .iter()
383 .chain(entity_ids.iter())
384 .copied()
385 .collect();
386
387 if !all_seed_ids.is_empty() {
388 let graph_memory_ids = traverse_from_memories_with_hops(
389 &conn,
390 &all_seed_ids,
391 &namespace_for_graph,
392 args.min_weight.unwrap_or(0.3),
393 args.max_hops.unwrap_or(2),
394 )?;
395
396 let already_in_results: std::collections::HashSet<i64> =
397 results.iter().map(|r| r.memory_id).collect();
398
399 for (graph_mem_id, hop) in graph_memory_ids {
400 if already_in_results.contains(&graph_mem_id) {
401 continue;
402 }
403 if let Some(row) = memories::read_full(&conn, graph_mem_id)? {
404 let snippet: String = row.body.chars().take(300).collect();
405 let graph_distance = 1.0 - 1.0 / (hop as f32 + 1.0);
406 graph_matches.push(RecallItem {
407 memory_id: row.id,
408 name: row.name,
409 namespace: row.namespace,
410 memory_type: row.memory_type,
411 description: row.description,
412 snippet,
413 distance: graph_distance,
414 score: RecallItem::score_from_distance(graph_distance),
415 source: "graph".to_string(),
416 graph_depth: Some(hop),
417 });
418 }
419 }
420 }
421 }
422
423 output::emit_json(&HybridSearchResponse {
424 query: args.query,
425 k: args.k,
426 rrf_k: args.rrf_k,
427 weights: Weights {
428 vec: args.weight_vec,
429 fts: args.weight_fts,
430 },
431 results,
432 graph_matches,
433 fts_degraded,
434 fts_error,
435 fts_auto_rebuilt,
436 vec_degraded,
437 vec_error,
438 warning: if vec_degraded {
439 Some(
440 "live query embedding unavailable; results are FTS5 BM25 only (semantic relevance reduced)"
441 .to_string(),
442 )
443 } else {
444 None
445 },
446 elapsed_ms: start.elapsed().as_millis() as u64,
447 })?;
448
449 Ok(())
450}
451
452#[cfg(test)]
453mod tests {
454 use super::*;
455
456 #[derive(clap::Parser)]
457 struct TestCli {
458 #[command(flatten)]
459 args: HybridSearchArgs,
460 }
461
462 #[test]
463 fn graph_flags_parse_as_none_when_absent() {
464 use clap::Parser;
468 let cli = TestCli::try_parse_from(["hybrid-search", "q"]).expect("bare query parses");
469 assert!(cli.args.max_hops.is_none());
470 assert!(cli.args.min_weight.is_none());
471 let cli = TestCli::try_parse_from(["hybrid-search", "q", "--max-hops", "2"])
472 .expect("explicit flag parses");
473 assert_eq!(cli.args.max_hops, Some(2));
474 }
475
476 fn empty_response(
477 k: usize,
478 rrf_k: u32,
479 weight_vec: f32,
480 weight_fts: f32,
481 ) -> HybridSearchResponse {
482 HybridSearchResponse {
483 query: "test query".to_string(),
484 k,
485 rrf_k,
486 weights: Weights {
487 vec: weight_vec,
488 fts: weight_fts,
489 },
490 results: vec![],
491 graph_matches: vec![],
492 fts_degraded: false,
493 fts_error: None,
494 fts_auto_rebuilt: false,
495 vec_degraded: false,
496 vec_error: None,
497 warning: None,
498 elapsed_ms: 0,
499 }
500 }
501
502 #[test]
503 fn hybrid_search_response_empty_serializes_correct_fields() {
504 let resp = empty_response(10, 60, 1.0, 1.0);
505 let json = serde_json::to_string(&resp).unwrap();
506 assert!(json.contains("\"results\""), "must contain results field");
507 assert!(json.contains("\"query\""), "must contain query field");
508 assert!(json.contains("\"k\""), "must contain k field");
509 assert!(
510 json.contains("\"graph_matches\""),
511 "must contain graph_matches field"
512 );
513 assert!(
514 !json.contains("\"combined_rank\""),
515 "must not contain combined_rank"
516 );
517 assert!(
518 !json.contains("\"vec_rank_list\""),
519 "must not contain vec_rank_list"
520 );
521 assert!(
522 !json.contains("\"fts_rank_list\""),
523 "must not contain fts_rank_list"
524 );
525 }
526
527 #[test]
528 fn hybrid_search_response_serializes_rrf_k_and_weights() {
529 let resp = empty_response(5, 60, 0.7, 0.3);
530 let json = serde_json::to_string(&resp).unwrap();
531 assert!(json.contains("\"rrf_k\""), "must contain rrf_k field");
532 assert!(json.contains("\"weights\""), "must contain weights field");
533 assert!(json.contains("\"vec\""), "must contain weights.vec field");
534 assert!(json.contains("\"fts\""), "must contain weights.fts field");
535 }
536
537 #[test]
538 fn hybrid_search_response_serializes_elapsed_ms() {
539 let mut resp = empty_response(5, 60, 1.0, 1.0);
540 resp.elapsed_ms = 123;
541 let json = serde_json::to_string(&resp).unwrap();
542 assert!(
543 json.contains("\"elapsed_ms\""),
544 "must contain elapsed_ms field"
545 );
546 assert!(json.contains("123"), "deve serializar valor de elapsed_ms");
547 }
548
549 #[test]
550 fn weights_struct_serializes_correctly() {
551 let w = Weights { vec: 0.6, fts: 0.4 };
552 let json = serde_json::to_string(&w).unwrap();
553 assert!(json.contains("\"vec\""));
554 assert!(json.contains("\"fts\""));
555 }
556
557 #[test]
558 fn hybrid_search_item_omits_fts_rank_when_none() {
559 let item = HybridSearchItem {
560 memory_id: 1,
561 name: "mem".to_string(),
562 namespace: "default".to_string(),
563 memory_type: "user".to_string(),
564 description: "desc".to_string(),
565 body: "content".to_string(),
566 snippet: "content".to_string(),
567 combined_score: 0.0328,
568 score: 0.0328,
569 source: "hybrid".to_string(),
570 vec_rank: Some(1),
571 fts_rank: None,
572 rrf_score: Some(0.0328),
573 normalized_score: 1.0,
574 vec_distance: Some(0.12),
575 fts_bm25: None,
576 };
577 let json = serde_json::to_string(&item).unwrap();
578 assert!(
579 json.contains("\"vec_rank\""),
580 "must contain vec_rank when Some"
581 );
582 assert!(
583 !json.contains("\"fts_rank\""),
584 "must not contain fts_rank when None"
585 );
586 }
587
588 #[test]
589 fn hybrid_search_item_omits_vec_rank_when_none() {
590 let item = HybridSearchItem {
591 memory_id: 2,
592 name: "mem2".to_string(),
593 namespace: "default".to_string(),
594 memory_type: "fact".to_string(),
595 description: "desc2".to_string(),
596 body: "corpo2".to_string(),
597 snippet: "corpo2".to_string(),
598 combined_score: 0.016,
599 score: 0.016,
600 source: "hybrid".to_string(),
601 vec_rank: None,
602 fts_rank: Some(2),
603 rrf_score: Some(0.016),
604 normalized_score: 0.5,
605 vec_distance: None,
606 fts_bm25: None,
607 };
608 let json = serde_json::to_string(&item).unwrap();
609 assert!(
610 !json.contains("\"vec_rank\""),
611 "must not contain vec_rank when None"
612 );
613 assert!(
614 json.contains("\"fts_rank\""),
615 "must contain fts_rank when Some"
616 );
617 }
618
619 #[test]
620 fn hybrid_search_item_serializes_both_ranks_when_some() {
621 let item = HybridSearchItem {
622 memory_id: 3,
623 name: "mem3".to_string(),
624 namespace: "ns".to_string(),
625 memory_type: "entity".to_string(),
626 description: "desc3".to_string(),
627 body: "corpo3".to_string(),
628 snippet: "corpo3".to_string(),
629 combined_score: 0.05,
630 score: 0.05,
631 source: "hybrid".to_string(),
632 vec_rank: Some(3),
633 fts_rank: Some(1),
634 rrf_score: Some(0.05),
635 normalized_score: 0.8,
636 vec_distance: Some(0.25),
637 fts_bm25: None,
638 };
639 let json = serde_json::to_string(&item).unwrap();
640 assert!(json.contains("\"vec_rank\""), "must contain vec_rank");
641 assert!(json.contains("\"fts_rank\""), "must contain fts_rank");
642 assert!(json.contains("\"type\""), "deve serializar type renomeado");
643 assert!(!json.contains("memory_type"), "must not expose memory_type");
644 }
645
646 #[test]
647 fn hybrid_search_response_serializes_k_correctly() {
648 let resp = empty_response(5, 60, 1.0, 1.0);
649 let json = serde_json::to_string(&resp).unwrap();
650 assert!(json.contains("\"k\":5"), "deve serializar k=5");
651 }
652
653 #[test]
654 fn hybrid_search_response_with_graph_matches() {
655 use crate::output::RecallItem;
656 let resp = HybridSearchResponse {
657 query: "test".to_string(),
658 k: 5,
659 rrf_k: 60,
660 weights: Weights { vec: 1.0, fts: 1.0 },
661 results: vec![],
662 graph_matches: vec![RecallItem {
663 memory_id: 1,
664 name: "graph-hit".to_string(),
665 namespace: "global".to_string(),
666 memory_type: "document".to_string(),
667 description: "found via graph".to_string(),
668 snippet: "graph content".to_string(),
669 distance: 0.1,
670 score: 0.9,
671 source: "graph".to_string(),
672 graph_depth: Some(1),
673 }],
674 fts_degraded: false,
675 fts_error: None,
676 fts_auto_rebuilt: false,
677 vec_degraded: false,
678 vec_error: None,
679 warning: None,
680 elapsed_ms: 42,
681 };
682 let json = serde_json::to_value(&resp).unwrap();
683 assert_eq!(json["graph_matches"].as_array().unwrap().len(), 1);
684 assert_eq!(json["graph_matches"][0]["source"], "graph");
685 assert_eq!(json["graph_matches"][0]["graph_depth"], 1);
686 }
687
688 #[test]
689 fn fts_degraded_omitted_on_success_present_on_failure() {
690 let ok_resp = empty_response(5, 60, 1.0, 1.0);
692 let ok_json = serde_json::to_string(&ok_resp).unwrap();
693 assert!(
694 !ok_json.contains("\"fts_degraded\""),
695 "fts_degraded must be absent when false"
696 );
697 assert!(
698 !ok_json.contains("\"fts_error\""),
699 "fts_error must be absent when None"
700 );
701
702 let mut degraded_resp = empty_response(5, 60, 1.0, 1.0);
704 degraded_resp.fts_degraded = true;
705 degraded_resp.fts_error = Some("FTS5 table corrupted".to_string());
706 let degraded_json = serde_json::to_string(°raded_resp).unwrap();
707 assert!(
708 degraded_json.contains("\"fts_degraded\":true"),
709 "fts_degraded must be present and true when degraded"
710 );
711 assert!(
712 degraded_json.contains("\"fts_error\""),
713 "fts_error must be present when Some"
714 );
715 assert!(
716 degraded_json.contains("FTS5 table corrupted"),
717 "fts_error must contain the error message"
718 );
719 }
720}