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.with_graph.then_some(()).filter(|_| !results.is_empty()).and(embedding.as_ref()) {
362 let namespace_for_graph = namespace.clone();
363 let memory_ids: Vec<i64> = results.iter().map(|r| r.memory_id).collect();
364
365 let entity_knn =
366 entities::knn_search(&conn, emb, &namespace_for_graph, 5)?;
367 let entity_ids: Vec<i64> = entity_knn.iter().map(|(id, _)| *id).collect();
368
369 let all_seed_ids: Vec<i64> = memory_ids
370 .iter()
371 .chain(entity_ids.iter())
372 .copied()
373 .collect();
374
375 if !all_seed_ids.is_empty() {
376 let graph_memory_ids = traverse_from_memories_with_hops(
377 &conn,
378 &all_seed_ids,
379 &namespace_for_graph,
380 args.min_weight.unwrap_or(0.3),
381 args.max_hops.unwrap_or(2),
382 )?;
383
384 let already_in_results: std::collections::HashSet<i64> =
385 results.iter().map(|r| r.memory_id).collect();
386
387 for (graph_mem_id, hop) in graph_memory_ids {
388 if already_in_results.contains(&graph_mem_id) {
389 continue;
390 }
391 if let Some(row) = memories::read_full(&conn, graph_mem_id)? {
392 let snippet: String = row.body.chars().take(300).collect();
393 let graph_distance = 1.0 - 1.0 / (hop as f32 + 1.0);
394 graph_matches.push(RecallItem {
395 memory_id: row.id,
396 name: row.name,
397 namespace: row.namespace,
398 memory_type: row.memory_type,
399 description: row.description,
400 snippet,
401 distance: graph_distance,
402 score: RecallItem::score_from_distance(graph_distance),
403 source: "graph".to_string(),
404 graph_depth: Some(hop),
405 });
406 }
407 }
408 }
409 }
410
411 output::emit_json(&HybridSearchResponse {
412 query: args.query,
413 k: args.k,
414 rrf_k: args.rrf_k,
415 weights: Weights {
416 vec: args.weight_vec,
417 fts: args.weight_fts,
418 },
419 results,
420 graph_matches,
421 fts_degraded,
422 fts_error,
423 fts_auto_rebuilt,
424 vec_degraded,
425 vec_error,
426 warning: if vec_degraded {
427 Some(
428 "live query embedding unavailable; results are FTS5 BM25 only (semantic relevance reduced)"
429 .to_string(),
430 )
431 } else {
432 None
433 },
434 elapsed_ms: start.elapsed().as_millis() as u64,
435 })?;
436
437 Ok(())
438}
439
440#[cfg(test)]
441mod tests {
442 use super::*;
443
444 #[derive(clap::Parser)]
445 struct TestCli {
446 #[command(flatten)]
447 args: HybridSearchArgs,
448 }
449
450 #[test]
451 fn graph_flags_parse_as_none_when_absent() {
452 use clap::Parser;
456 let cli = TestCli::try_parse_from(["hybrid-search", "q"]).expect("bare query parses");
457 assert!(cli.args.max_hops.is_none());
458 assert!(cli.args.min_weight.is_none());
459 let cli = TestCli::try_parse_from(["hybrid-search", "q", "--max-hops", "2"])
460 .expect("explicit flag parses");
461 assert_eq!(cli.args.max_hops, Some(2));
462 }
463
464 fn empty_response(
465 k: usize,
466 rrf_k: u32,
467 weight_vec: f32,
468 weight_fts: f32,
469 ) -> HybridSearchResponse {
470 HybridSearchResponse {
471 query: "test query".to_string(),
472 k,
473 rrf_k,
474 weights: Weights {
475 vec: weight_vec,
476 fts: weight_fts,
477 },
478 results: vec![],
479 graph_matches: vec![],
480 fts_degraded: false,
481 fts_error: None,
482 fts_auto_rebuilt: false,
483 vec_degraded: false,
484 vec_error: None,
485 warning: None,
486 elapsed_ms: 0,
487 }
488 }
489
490 #[test]
491 fn hybrid_search_response_empty_serializes_correct_fields() {
492 let resp = empty_response(10, 60, 1.0, 1.0);
493 let json = serde_json::to_string(&resp).unwrap();
494 assert!(json.contains("\"results\""), "must contain results field");
495 assert!(json.contains("\"query\""), "must contain query field");
496 assert!(json.contains("\"k\""), "must contain k field");
497 assert!(
498 json.contains("\"graph_matches\""),
499 "must contain graph_matches field"
500 );
501 assert!(
502 !json.contains("\"combined_rank\""),
503 "must not contain combined_rank"
504 );
505 assert!(
506 !json.contains("\"vec_rank_list\""),
507 "must not contain vec_rank_list"
508 );
509 assert!(
510 !json.contains("\"fts_rank_list\""),
511 "must not contain fts_rank_list"
512 );
513 }
514
515 #[test]
516 fn hybrid_search_response_serializes_rrf_k_and_weights() {
517 let resp = empty_response(5, 60, 0.7, 0.3);
518 let json = serde_json::to_string(&resp).unwrap();
519 assert!(json.contains("\"rrf_k\""), "must contain rrf_k field");
520 assert!(json.contains("\"weights\""), "must contain weights field");
521 assert!(json.contains("\"vec\""), "must contain weights.vec field");
522 assert!(json.contains("\"fts\""), "must contain weights.fts field");
523 }
524
525 #[test]
526 fn hybrid_search_response_serializes_elapsed_ms() {
527 let mut resp = empty_response(5, 60, 1.0, 1.0);
528 resp.elapsed_ms = 123;
529 let json = serde_json::to_string(&resp).unwrap();
530 assert!(
531 json.contains("\"elapsed_ms\""),
532 "must contain elapsed_ms field"
533 );
534 assert!(json.contains("123"), "deve serializar valor de elapsed_ms");
535 }
536
537 #[test]
538 fn weights_struct_serializes_correctly() {
539 let w = Weights { vec: 0.6, fts: 0.4 };
540 let json = serde_json::to_string(&w).unwrap();
541 assert!(json.contains("\"vec\""));
542 assert!(json.contains("\"fts\""));
543 }
544
545 #[test]
546 fn hybrid_search_item_omits_fts_rank_when_none() {
547 let item = HybridSearchItem {
548 memory_id: 1,
549 name: "mem".to_string(),
550 namespace: "default".to_string(),
551 memory_type: "user".to_string(),
552 description: "desc".to_string(),
553 body: "content".to_string(),
554 snippet: "content".to_string(),
555 combined_score: 0.0328,
556 score: 0.0328,
557 source: "hybrid".to_string(),
558 vec_rank: Some(1),
559 fts_rank: None,
560 rrf_score: Some(0.0328),
561 normalized_score: 1.0,
562 vec_distance: Some(0.12),
563 fts_bm25: None,
564 };
565 let json = serde_json::to_string(&item).unwrap();
566 assert!(
567 json.contains("\"vec_rank\""),
568 "must contain vec_rank when Some"
569 );
570 assert!(
571 !json.contains("\"fts_rank\""),
572 "must not contain fts_rank when None"
573 );
574 }
575
576 #[test]
577 fn hybrid_search_item_omits_vec_rank_when_none() {
578 let item = HybridSearchItem {
579 memory_id: 2,
580 name: "mem2".to_string(),
581 namespace: "default".to_string(),
582 memory_type: "fact".to_string(),
583 description: "desc2".to_string(),
584 body: "corpo2".to_string(),
585 snippet: "corpo2".to_string(),
586 combined_score: 0.016,
587 score: 0.016,
588 source: "hybrid".to_string(),
589 vec_rank: None,
590 fts_rank: Some(2),
591 rrf_score: Some(0.016),
592 normalized_score: 0.5,
593 vec_distance: None,
594 fts_bm25: None,
595 };
596 let json = serde_json::to_string(&item).unwrap();
597 assert!(
598 !json.contains("\"vec_rank\""),
599 "must not contain vec_rank when None"
600 );
601 assert!(
602 json.contains("\"fts_rank\""),
603 "must contain fts_rank when Some"
604 );
605 }
606
607 #[test]
608 fn hybrid_search_item_serializes_both_ranks_when_some() {
609 let item = HybridSearchItem {
610 memory_id: 3,
611 name: "mem3".to_string(),
612 namespace: "ns".to_string(),
613 memory_type: "entity".to_string(),
614 description: "desc3".to_string(),
615 body: "corpo3".to_string(),
616 snippet: "corpo3".to_string(),
617 combined_score: 0.05,
618 score: 0.05,
619 source: "hybrid".to_string(),
620 vec_rank: Some(3),
621 fts_rank: Some(1),
622 rrf_score: Some(0.05),
623 normalized_score: 0.8,
624 vec_distance: Some(0.25),
625 fts_bm25: None,
626 };
627 let json = serde_json::to_string(&item).unwrap();
628 assert!(json.contains("\"vec_rank\""), "must contain vec_rank");
629 assert!(json.contains("\"fts_rank\""), "must contain fts_rank");
630 assert!(json.contains("\"type\""), "deve serializar type renomeado");
631 assert!(!json.contains("memory_type"), "must not expose memory_type");
632 }
633
634 #[test]
635 fn hybrid_search_response_serializes_k_correctly() {
636 let resp = empty_response(5, 60, 1.0, 1.0);
637 let json = serde_json::to_string(&resp).unwrap();
638 assert!(json.contains("\"k\":5"), "deve serializar k=5");
639 }
640
641 #[test]
642 fn hybrid_search_response_with_graph_matches() {
643 use crate::output::RecallItem;
644 let resp = HybridSearchResponse {
645 query: "test".to_string(),
646 k: 5,
647 rrf_k: 60,
648 weights: Weights { vec: 1.0, fts: 1.0 },
649 results: vec![],
650 graph_matches: vec![RecallItem {
651 memory_id: 1,
652 name: "graph-hit".to_string(),
653 namespace: "global".to_string(),
654 memory_type: "document".to_string(),
655 description: "found via graph".to_string(),
656 snippet: "graph content".to_string(),
657 distance: 0.1,
658 score: 0.9,
659 source: "graph".to_string(),
660 graph_depth: Some(1),
661 }],
662 fts_degraded: false,
663 fts_error: None,
664 fts_auto_rebuilt: false,
665 vec_degraded: false,
666 vec_error: None,
667 warning: None,
668 elapsed_ms: 42,
669 };
670 let json = serde_json::to_value(&resp).unwrap();
671 assert_eq!(json["graph_matches"].as_array().unwrap().len(), 1);
672 assert_eq!(json["graph_matches"][0]["source"], "graph");
673 assert_eq!(json["graph_matches"][0]["graph_depth"], 1);
674 }
675
676 #[test]
677 fn fts_degraded_omitted_on_success_present_on_failure() {
678 let ok_resp = empty_response(5, 60, 1.0, 1.0);
680 let ok_json = serde_json::to_string(&ok_resp).unwrap();
681 assert!(
682 !ok_json.contains("\"fts_degraded\""),
683 "fts_degraded must be absent when false"
684 );
685 assert!(
686 !ok_json.contains("\"fts_error\""),
687 "fts_error must be absent when None"
688 );
689
690 let mut degraded_resp = empty_response(5, 60, 1.0, 1.0);
692 degraded_resp.fts_degraded = true;
693 degraded_resp.fts_error = Some("FTS5 table corrupted".to_string());
694 let degraded_json = serde_json::to_string(°raded_resp).unwrap();
695 assert!(
696 degraded_json.contains("\"fts_degraded\":true"),
697 "fts_degraded must be present and true when degraded"
698 );
699 assert!(
700 degraded_json.contains("\"fts_error\""),
701 "fts_error must be present when Some"
702 );
703 assert!(
704 degraded_json.contains("FTS5 table corrupted"),
705 "fts_error must contain the error message"
706 );
707 }
708}