1use crate::errors::AppError;
8use crate::graph::{
9 bfs_with_predecessors, traverse_from_memories_with_hops_capped, PredecessorMap,
10};
11use crate::output;
12use crate::paths::AppPaths;
13use crate::storage::connection::open_ro;
14use crate::storage::fusion::{rrf_fuse, rrf_max_possible};
15use crate::storage::{entities, memories};
16
17use serde::Serialize;
18use std::collections::HashSet;
19use std::sync::Arc;
20use tokio::sync::Semaphore;
21use tokio::task::JoinSet;
22
23#[derive(clap::Args)]
25#[command(
26 about = "Deep parallel multi-hop GraphRAG research via query decomposition",
27 after_long_help = "EXAMPLES:\n \
28 # Basic deep research\n \
29 sqlite-graphrag deep-research \"auth architecture decisions\"\n\n \
30 # With custom parameters\n \
31 sqlite-graphrag deep-research \"auth\" --k 20 --max-hops 3 --max-sub-queries 7\n\n \
32 # Include full memory bodies in output\n \
33 sqlite-graphrag deep-research \"auth\" --with-bodies\n\n \
34 # Tune RRF and graph scoring\n \
35 sqlite-graphrag deep-research \"auth and deployment\" --rrf-k 60 --graph-decay 0.7"
36)]
37pub struct DeepResearchArgs {
38 #[arg(
40 value_name = "QUERY",
41 allow_hyphen_values = true,
42 help = "Research query to decompose and search"
43 )]
44 pub query: String,
45 #[arg(
47 long,
48 short,
49 aliases = ["limit", "top-k"],
50 default_value_t = 20,
51 help = "Results per sub-query (Recall@20 captures 95%+ relevant hits)"
52 )]
53 pub k: usize,
54 #[arg(
56 long,
57 default_value_t = 7,
58 help = "Maximum sub-queries (covers complex multi-hop queries)"
59 )]
60 pub max_sub_queries: usize,
61 #[arg(
63 long,
64 default_value_t = 3,
65 help = "Multi-hop graph traversal depth (sweet spot: 2-3 hops)"
66 )]
67 pub max_hops: usize,
68 #[arg(
70 long,
71 default_value_t = 0.3,
72 help = "Minimum edge weight for graph traversal"
73 )]
74 pub min_weight: f64,
75 #[arg(long, help = "Maximum concurrent sub-queries (default: min(cpus, 8))")]
77 pub max_concurrency: Option<usize>,
78 #[arg(long, default_value_t = 30, help = "Timeout per sub-query in seconds")]
80 pub timeout: u64,
81 #[arg(
83 long,
84 default_value_t = false,
85 help = "Include full memory bodies in results"
86 )]
87 pub with_bodies: bool,
88 #[arg(
90 long,
91 default_value_t = 50,
92 help = "Maximum results after deduplication"
93 )]
94 pub max_results: usize,
95 #[arg(
97 long,
98 default_value_t = 60.0,
99 help = "RRF k parameter (higher = less weight on top ranks)"
100 )]
101 pub rrf_k: f64,
102 #[arg(
104 long,
105 default_value_t = 0.7,
106 help = "Graph score decay factor per hop (0.0-1.0)"
107 )]
108 pub graph_decay: f64,
109 #[arg(
111 long,
112 default_value_t = 0.05,
113 help = "Minimum score threshold for graph-expanded results"
114 )]
115 pub graph_min_score: f64,
116 #[arg(
118 long,
119 help = "Limit neighbours per entity per hop for graph traversal (default: unlimited)"
120 )]
121 pub max_neighbors_per_hop: Option<usize>,
122 #[arg(
124 long,
125 help = "Namespace (env: SQLITE_GRAPHRAG_NAMESPACE, default: global)"
126 )]
127 pub namespace: Option<String>,
128 #[arg(long, default_value = "none", value_parser = ["none"], hide = true)]
130 pub mode: String,
131 #[arg(
133 long,
134 value_name = "USD",
135 help = "Max LLM cost in USD (effective with --mode claude-code/codex)"
136 )]
137 pub max_cost_usd: Option<f64>,
138 #[arg(long, hide = true)]
140 pub json: bool,
141 #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
143 pub db: Option<String>,
144}
145
146#[derive(Serialize)]
147struct SubQuery {
148 id: usize,
149 text: String,
150 source: &'static str,
151}
152
153#[derive(Serialize)]
154struct DeepResult {
155 name: String,
156 score: f64,
157 source: String,
158 sub_query_ids: Vec<usize>,
159 snippet: String,
160 #[serde(skip_serializing_if = "Option::is_none")]
161 body: Option<String>,
162 hop_distance: Option<usize>,
163}
164
165#[derive(Serialize, Clone)]
167struct EvidenceNode {
168 entity: String,
169 #[serde(skip_serializing_if = "Option::is_none")]
170 relation: Option<String>,
171 #[serde(skip_serializing_if = "Option::is_none")]
172 weight: Option<f64>,
173}
174
175#[derive(Serialize)]
184struct EvidenceChain {
185 from: String,
186 to: String,
187 path: Vec<EvidenceNode>,
188 total_weight: f64,
189 depth: usize,
190 sub_query_ids: Vec<usize>,
191}
192
193#[derive(Serialize)]
194struct ResearchStats {
195 sub_queries_total: usize,
196 sub_queries_completed: usize,
197 sub_queries_failed: usize,
198 sub_queries_timed_out: usize,
199 unique_memories_found: usize,
200 evidence_chains_found: usize,
201 elapsed_ms: u64,
202 vec_degraded: bool,
203}
204
205#[derive(Serialize)]
206struct GraphContextEntity {
207 name: String,
208 entity_type: String,
209 degree: u32,
210}
211
212#[derive(Serialize)]
213struct GraphContextRel {
214 from: String,
215 to: String,
216 relation: String,
217 weight: f64,
218}
219
220#[derive(Serialize)]
221struct GraphContext {
222 entities: Vec<GraphContextEntity>,
223 relationships: Vec<GraphContextRel>,
224}
225
226#[derive(Serialize)]
227struct DeepResearchResponse {
228 query: String,
229 sub_queries: Vec<SubQuery>,
230 results: Vec<DeepResult>,
231 evidence_chains: Vec<EvidenceChain>,
232 #[serde(skip_serializing_if = "Option::is_none")]
233 graph_context: Option<GraphContext>,
234 stats: ResearchStats,
235}
236
237type MergedHit = (f64, String, String, String, Option<usize>, Vec<usize>);
239
240struct SubQueryResult {
242 sub_query_id: usize,
243 hits: Vec<(i64, f64, String, String, String, Option<usize>)>,
245 chains: Vec<EvidenceChain>,
247}
248
249#[tracing::instrument(skip_all, level = "debug", name = "deep_research")]
251pub fn run(args: DeepResearchArgs, llm_backend: crate::cli::LlmBackendChoice) -> Result<(), AppError> {
252 tracing::debug!(target: "deep_research", query = %args.query, k = args.k, "starting deep research");
253 let rt = tokio::runtime::Builder::new_multi_thread()
254 .worker_threads(2)
255 .enable_all()
256 .build()
257 .map_err(|e| AppError::Internal(anyhow::anyhow!("failed to build tokio runtime: {e}")))?;
258 rt.block_on(run_async(args, llm_backend))
259}
260
261async fn run_async(args: DeepResearchArgs, llm_backend: crate::cli::LlmBackendChoice) -> Result<(), AppError> {
263 let start = std::time::Instant::now();
264
265 if args.query.trim().is_empty() {
266 return Err(AppError::Validation(crate::i18n::validation::empty_query()));
267 }
268
269 if args.max_cost_usd.is_some() && args.mode == "none" {
270 tracing::warn!(target: "deep_research", "--max-cost-usd has no effect without --mode claude-code/codex");
271 }
272
273 let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
274 let paths = AppPaths::resolve(args.db.as_deref())?;
275 crate::storage::connection::ensure_db_ready(&paths)?;
276
277 let sub_query_texts = decompose_query(&args.query, args.max_sub_queries);
279 let sub_queries: Vec<SubQuery> = sub_query_texts
280 .iter()
281 .enumerate()
282 .map(|(i, text)| SubQuery {
283 id: i,
284 text: text.clone(),
285 source: if sub_query_texts.len() == 1 {
286 "original"
287 } else {
288 "decomposed"
289 },
290 })
291 .collect();
292
293 output::emit_progress_i18n(
298 "Computing per-sub-query embeddings...",
299 "Calculando embeddings por sub-consulta...",
300 );
301 let mut sub_embeddings: Vec<Option<Arc<Vec<f32>>>> = Vec::with_capacity(sub_query_texts.len());
302 let mut vec_degraded = false;
303 for sq_text in &sub_query_texts {
304 match crate::embedder::try_embed_query_with_deterministic_fallback(
305 &paths.models,
306 sq_text,
307 Some(llm_backend),
308 ) {
309 Ok((v, _backend)) => sub_embeddings.push(Some(Arc::new(v))),
310 Err(reason) => {
311 tracing::warn!(target: "deep_research", fallback_reason = %reason, reason_code = %reason.reason_code(), "embedding failed for sub-query; falling back to FTS5");
312 sub_embeddings.push(None);
313 vec_degraded = true;
314 }
315 }
316 }
317
318 let cpu_count = std::thread::available_parallelism()
320 .map(|n| n.get())
321 .unwrap_or(4);
322 let permits = args
323 .max_concurrency
324 .unwrap_or_else(|| cpu_count.min(8))
325 .min(sub_queries.len())
326 .max(1);
327 let semaphore = Arc::new(Semaphore::new(permits));
328 let timeout_dur = std::time::Duration::from_secs(args.timeout);
329
330 let mut join_set: JoinSet<Result<SubQueryResult, (usize, String)>> = JoinSet::new();
331
332 for (idx, sq_text) in sub_query_texts.iter().enumerate() {
333 let sem = Arc::clone(&semaphore);
334 let emb = sub_embeddings[idx].clone();
336 let ns = namespace.clone();
337 let db_path = paths.db.clone();
338 let query_text = sq_text.clone();
339 let k = args.k;
340 let max_hops = args.max_hops;
341 let min_weight = args.min_weight;
342 let rrf_k = args.rrf_k;
343 let graph_decay = args.graph_decay;
344 let graph_min_score = args.graph_min_score;
345 let max_neighbors_per_hop = args.max_neighbors_per_hop;
346
347 join_set.spawn(async move {
348 let _permit = sem
349 .acquire_owned()
350 .await
351 .map_err(|e| (idx, format!("semaphore closed: {e}")))?;
352
353 let result = tokio::time::timeout(timeout_dur, async move {
355 execute_sub_query(
356 idx,
357 &query_text,
358 emb.as_ref().map(|v| v.as_slice()),
359 &ns,
360 &db_path,
361 k,
362 max_hops,
363 min_weight,
364 rrf_k,
365 graph_decay,
366 graph_min_score,
367 max_neighbors_per_hop,
368 )
369 })
370 .await;
371
372 match result {
373 Ok(inner) => inner.map_err(|e| (idx, e)),
374 Err(_) => Err((idx, "timeout".to_string())),
375 }
376 });
377 }
378
379 let mut sub_query_results: Vec<SubQueryResult> = Vec::with_capacity(sub_queries.len());
381 let mut failed_count = 0usize;
382 let mut timed_out_count = 0usize;
383
384 while let Some(join_result) = join_set.join_next().await {
385 match join_result {
386 Ok(Ok(sqr)) => sub_query_results.push(sqr),
387 Ok(Err((_idx, reason))) => {
388 if reason == "timeout" {
389 timed_out_count += 1;
390 } else {
391 failed_count += 1;
392 }
393 tracing::warn!(target: "deep_research", sub_query_id = _idx, reason = %reason, "sub-query failed");
394 }
395 Err(join_err) => {
396 failed_count += 1;
397 if join_err.is_panic() {
398 tracing::error!(target: "deep_research", error = %join_err, "sub-query task panicked");
399 } else {
400 tracing::warn!(target: "deep_research", error = %join_err, "sub-query task cancelled");
401 }
402 }
403 }
404 }
405
406 let mut merged: crate::hash::AHashMap<i64, MergedHit> =
409 crate::hash::AHashMap::with_capacity_and_hasher(
410 sub_query_results.len() * args.k,
411 Default::default(),
412 );
413
414 for sqr in &sub_query_results {
415 for (mem_id, score, source, snippet, body, hop) in &sqr.hits {
416 let entry = merged.entry(*mem_id).or_insert_with(|| {
417 (
418 *score,
419 source.clone(),
420 snippet.clone(),
421 body.clone(),
422 *hop,
423 Vec::new(),
424 )
425 });
426 if *score > entry.0 {
428 entry.0 = *score;
429 entry.1 = source.clone();
430 entry.2 = snippet.clone();
431 entry.3 = body.clone();
432 entry.4 = *hop;
433 }
434 if !entry.5.contains(&sqr.sub_query_id) {
435 entry.5.push(sqr.sub_query_id);
436 }
437 }
438 }
439
440 let conn = open_ro(&paths.db)?;
442 let mut results: Vec<DeepResult> = Vec::with_capacity(merged.len().min(args.max_results));
443
444 let mut ranked: Vec<(i64, MergedHit)> = merged.into_iter().collect();
446 ranked.sort_by(|a, b| {
447 b.1 .0
448 .partial_cmp(&a.1 .0)
449 .unwrap_or(std::cmp::Ordering::Equal)
450 });
451 ranked.truncate(args.max_results);
452
453 for (mem_id, (score, source, snippet, body, hop, sq_ids)) in ranked {
454 let name = match memories::read_full(&conn, mem_id)? {
455 Some(row) => row.name,
456 None => continue,
457 };
458 results.push(DeepResult {
459 name,
460 score,
461 source,
462 sub_query_ids: sq_ids,
463 snippet,
464 body: if args.with_bodies { Some(body) } else { None },
465 hop_distance: hop,
466 });
467 }
468
469 let completed_count = sub_query_results.len();
473 let mut evidence_chains: Vec<EvidenceChain> = Vec::with_capacity(completed_count * 2);
474 let mut seen_chain_keys: HashSet<String> = HashSet::with_capacity(completed_count * 2);
475
476 for sqr in sub_query_results {
477 for chain in sqr.chains {
478 let key = format!("{}->{}", chain.from, chain.to);
480 if seen_chain_keys.insert(key) {
481 evidence_chains.push(chain);
482 }
483 }
484 }
485
486 evidence_chains.retain(|c| c.depth >= 2);
488 evidence_chains.sort_by(|a, b| {
489 b.total_weight
490 .partial_cmp(&a.total_weight)
491 .unwrap_or(std::cmp::Ordering::Equal)
492 });
493
494 let unique_memories = results.len();
495 let evidence_count = evidence_chains.len();
496
497 let graph_context = if !results.is_empty() {
499 let result_names: Vec<&str> = results.iter().map(|r| r.name.as_str()).collect();
500 let mut ctx_entities: Vec<GraphContextEntity> = Vec::with_capacity(results.len());
501 let mut ctx_rels: Vec<GraphContextRel> = Vec::with_capacity(results.len() * 2);
502 let mut seen_entity_ids: crate::hash::AHashSet<i64> =
503 crate::hash::AHashSet::with_capacity_and_hasher(results.len(), Default::default());
504
505 for name in &result_names {
506 if let Ok(Some(eid)) = entities::find_entity_id(&conn, &namespace, name) {
507 if seen_entity_ids.insert(eid) {
508 let etype: String = conn
509 .query_row(
510 "SELECT COALESCE(type,'concept') FROM entities WHERE id = ?1",
511 rusqlite::params![eid],
512 |r| r.get(0),
513 )
514 .unwrap_or_else(|_| "concept".to_string());
515 let degree: u32 = conn
516 .query_row(
517 "SELECT COUNT(*) FROM relationships WHERE source_id = ?1 OR target_id = ?1",
518 rusqlite::params![eid],
519 |r| r.get(0),
520 )
521 .unwrap_or(0);
522 ctx_entities.push(GraphContextEntity {
523 name: name.to_string(),
524 entity_type: etype,
525 degree,
526 });
527 }
528 }
529 }
530
531 let entity_ids: Vec<i64> = seen_entity_ids.iter().copied().collect();
532 if entity_ids.len() >= 2 {
533 let placeholders: String = entity_ids.iter().map(|_| "?").collect::<Vec<_>>().join(",");
534 let sql = format!(
535 "SELECT s.name, t.name, r.relation, r.weight \
536 FROM relationships r \
537 JOIN entities s ON s.id = r.source_id \
538 JOIN entities t ON t.id = r.target_id \
539 WHERE r.source_id IN ({placeholders}) AND r.target_id IN ({placeholders}) \
540 LIMIT 50"
541 );
542 if let Ok(mut stmt) = conn.prepare(&sql) {
543 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> =
544 Vec::with_capacity(entity_ids.len() * 2);
545 for id in &entity_ids {
546 params.push(Box::new(*id));
547 }
548 for id in &entity_ids {
549 params.push(Box::new(*id));
550 }
551 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
552 params.iter().map(|p| p.as_ref()).collect();
553 if let Ok(rows) = stmt.query_map(param_refs.as_slice(), |r| {
554 Ok((
555 r.get::<_, String>(0)?,
556 r.get::<_, String>(1)?,
557 r.get::<_, String>(2)?,
558 r.get::<_, f64>(3)?,
559 ))
560 }) {
561 for row in rows.flatten() {
562 ctx_rels.push(GraphContextRel {
563 from: row.0,
564 to: row.1,
565 relation: row.2,
566 weight: row.3,
567 });
568 }
569 }
570 }
571 }
572
573 if ctx_entities.is_empty() {
574 None
575 } else {
576 Some(GraphContext {
577 entities: ctx_entities,
578 relationships: ctx_rels,
579 })
580 }
581 } else {
582 None
583 };
584
585 tracing::debug!(target: "deep_research",
586 total_results = results.len(),
587 total_chains = evidence_chains.len(),
588 "assembly complete"
589 );
590
591 output::emit_json(&DeepResearchResponse {
593 query: args.query,
594 sub_queries,
595 results,
596 evidence_chains,
597 graph_context,
598 stats: ResearchStats {
599 sub_queries_total: sub_query_texts.len(),
600 sub_queries_completed: completed_count,
601 sub_queries_failed: failed_count,
602 sub_queries_timed_out: timed_out_count,
603 unique_memories_found: unique_memories,
604 evidence_chains_found: evidence_count,
605 elapsed_ms: start.elapsed().as_millis() as u64,
606 vec_degraded,
607 },
608 })?;
609
610 Ok(())
611}
612
613fn decompose_query(query: &str, max: usize) -> Vec<String> {
616 if query.is_empty() {
617 return vec![query.to_string()];
618 }
619
620 let mut parts: Vec<String> = Vec::with_capacity(max);
621
622 let relational = [
624 " that caused ",
625 " depending on ",
626 " related to ",
627 " connected to ",
628 " linked to ",
629 " caused by ",
630 " followed by ",
631 ];
632 let mut text = query.to_string();
633 let mut did_relational_split = false;
634 for phrase in &relational {
635 if text.to_lowercase().contains(phrase) {
636 let lower = text.to_lowercase();
637 if let Some(pos) = lower.find(phrase) {
638 let left = text[..pos].trim().to_string();
639 let right = text[pos + phrase.len()..].trim().to_string();
640 if !left.is_empty() {
641 parts.push(left);
642 }
643 if !right.is_empty() {
644 text = right;
645 }
646 did_relational_split = true;
647 }
648 }
649 }
650 if did_relational_split && !text.is_empty() {
651 parts.push(text.clone());
652 }
653
654 if parts.is_empty() {
656 let semi_parts: Vec<&str> = query.split(';').collect();
658 if semi_parts.len() > 1 {
659 for p in &semi_parts {
660 let trimmed = p.trim();
661 if !trimmed.is_empty() {
662 parts.push(trimmed.to_string());
663 }
664 }
665 } else {
666 let normalized = query
669 .replace(" and ", ", ")
670 .replace(" AND ", ", ")
671 .replace(" e ", ", ")
672 .replace(" E ", ", ");
673 let comma_parts: Vec<&str> = normalized.split(',').collect();
674 if comma_parts.len() > 1 {
675 for p in &comma_parts {
676 let trimmed = p.trim();
677 if !trimmed.is_empty() {
678 parts.push(trimmed.to_string());
679 }
680 }
681 }
682 }
683 }
684
685 if parts.is_empty() {
687 let words: Vec<&str> = query.split_whitespace().filter(|w| w.len() > 2).collect();
688 if words.len() >= 3 {
689 parts.push(query.to_string());
690 parts.push(format!("{} {}", words[0], words[1]));
691 parts.push(format!(
692 "{} {}",
693 words[words.len() - 2],
694 words[words.len() - 1]
695 ));
696 }
697 }
698
699 if parts.is_empty() {
700 return vec![query.to_string()];
701 }
702
703 parts.truncate(max);
705 parts
706}
707
708fn reconstruct_path(
712 target_id: i64,
713 seed_entity_ids: &HashSet<i64>,
714 predecessor: &PredecessorMap,
715 entity_names: &crate::hash::AHashMap<i64, String>,
716) -> Option<(Vec<EvidenceNode>, f64)> {
717 let mut path_ids: Vec<(i64, Option<String>, Option<f64>)> = Vec::with_capacity(8);
718 let mut total_weight = 1.0_f64;
719 let mut current = target_id;
720
721 loop {
722 if seed_entity_ids.contains(¤t) {
723 break;
724 }
725 let (parent, relation, weight) = predecessor.get(¤t)?;
726 total_weight *= weight;
727 path_ids.push((current, Some(relation.clone()), Some(*weight)));
728 current = *parent;
729 }
730 path_ids.push((current, None, None));
732
733 path_ids.reverse();
735
736 let nodes: Vec<EvidenceNode> = path_ids
737 .into_iter()
738 .map(|(id, relation, weight)| EvidenceNode {
739 entity: entity_names
740 .get(&id)
741 .cloned()
742 .unwrap_or_else(|| format!("entity-{id}")),
743 relation,
744 weight,
745 })
746 .collect();
747
748 Some((nodes, total_weight))
749}
750
751#[allow(clippy::too_many_arguments)]
761fn execute_sub_query(
762 sub_query_id: usize,
763 query_text: &str,
764 embedding: Option<&[f32]>,
765 namespace: &str,
766 db_path: &std::path::Path,
767 k: usize,
768 max_hops: usize,
769 min_weight: f64,
770 rrf_k: f64,
771 graph_decay: f64,
772 graph_min_score: f64,
773 max_neighbors_per_hop: Option<usize>,
774) -> Result<SubQueryResult, String> {
775 let conn = open_ro(db_path).map_err(|e| format!("failed to open db: {e}"))?;
776
777 let mut hits: Vec<(i64, f64, String, String, String, Option<usize>)> =
778 Vec::with_capacity(k * 2);
779 let mut seen_ids: crate::hash::AHashSet<i64> =
780 crate::hash::AHashSet::with_capacity_and_hasher(k * 2, Default::default());
781
782 let (knn_ids, knn_distance_map) = if let Some(emb) = embedding {
786 let knn_results = memories::knn_search(&conn, emb, &[namespace.to_string()], None, k)
787 .map_err(|e| format!("knn_search failed: {e}"))?;
788 let ids: Vec<i64> = knn_results.iter().map(|(id, _)| *id).collect();
789 tracing::debug!(target: "deep_research", sub_query_id, knn_count = ids.len(), "KNN complete");
790 let dist_map: crate::hash::AHashMap<i64, f64> = knn_results
791 .iter()
792 .map(|(id, dist)| (*id, *dist as f64))
793 .collect();
794 (ids, dist_map)
795 } else {
796 tracing::debug!(target: "deep_research", sub_query_id, "KNN skipped (no embedding); FTS5-only");
797 (vec![], crate::hash::AHashMap::default())
798 };
799
800 let fts_results = match memories::fts_search(&conn, query_text, namespace, None, k) {
802 Ok(rows) => rows,
803 Err(e) => {
804 tracing::warn!(target: "deep_research",
805 sub_query_id,
806 "FTS5 search failed, continuing with KNN only: {e}"
807 );
808 vec![]
809 }
810 };
811 let fts_ids: Vec<i64> = fts_results.iter().map(|r| r.id).collect();
812 tracing::debug!(target: "deep_research", sub_query_id, fts_count = fts_ids.len(), "FTS complete");
813
814 let rrf_scores = rrf_fuse(&[(1.0, &knn_ids), (1.0, &fts_ids)], rrf_k);
816 let max_possible = rrf_max_possible(&[1.0, 1.0], rrf_k);
817
818 let mut fused: Vec<(i64, f64)> = rrf_scores.into_iter().collect();
820 fused.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
821 fused.truncate(k * 2);
822 tracing::debug!(target: "deep_research",
823 sub_query_id,
824 fused_count = fused.len(),
825 "RRF fusion complete"
826 );
827
828 if fused.is_empty() && !knn_ids.is_empty() {
829 tracing::warn!(target: "deep_research", sub_query_id, knn_count = knn_ids.len(), fts_count = fts_ids.len(),
830 "RRF fusion returned 0 results despite KNN/FTS hits; consider lowering --graph-min-score");
831 }
832
833 for (memory_id, combined_score) in &fused {
834 if seen_ids.insert(*memory_id) {
835 let normalized = if max_possible > 0.0 {
836 combined_score / max_possible
837 } else {
838 0.0
839 };
840 let score = normalized.clamp(0.0, 1.0);
841 let in_knn = knn_distance_map.contains_key(memory_id);
842 let in_fts = fts_ids.contains(memory_id);
843 let source = match (in_knn, in_fts) {
844 (true, true) => "hybrid",
845 (true, false) => "knn",
846 (false, true) => "fts",
847 (false, false) => "graph",
848 };
849 if let Ok(Some(row)) = memories::read_full(&conn, *memory_id) {
850 let snippet: String = row.body.chars().take(300).collect();
851 hits.push((
852 *memory_id,
853 score,
854 source.to_string(),
855 snippet,
856 row.body,
857 None,
858 ));
859 }
860 }
861 }
862
863 let memory_ids: Vec<i64> = hits.iter().map(|(id, ..)| *id).collect();
866 let mut chains: Vec<EvidenceChain> = Vec::with_capacity(memory_ids.len());
867
868 if !memory_ids.is_empty() && max_hops > 0 {
869 let entity_ids: Vec<i64> = if let Some(emb) = embedding {
871 entities::knn_search(&conn, emb, namespace, 5)
872 .inspect_err(|e| tracing::warn!(target: "deep_research", error = %e, "entity KNN search failed, skipping graph seed"))
873 .unwrap_or_default()
874 .iter()
875 .map(|(id, _)| *id)
876 .collect()
877 } else {
878 vec![]
879 };
880
881 let top_seed_count = 5.min(memory_ids.len());
884 let top_memory_ids = &memory_ids[..top_seed_count];
885 let mut seed_entity_ids: Vec<i64> = entity_ids.clone();
886 for &mem_id in top_memory_ids {
887 let mut stmt = conn
888 .prepare_cached("SELECT entity_id FROM memory_entities WHERE memory_id = ?1")
889 .map_err(|e| format!("prepare failed: {e}"))?;
890 let ids: Vec<i64> = stmt
891 .query_map(rusqlite::params![mem_id], |r| r.get(0))
892 .map_err(|e| format!("query failed: {e}"))?
893 .filter_map(|r| r.ok())
894 .collect();
895 seed_entity_ids.extend(ids);
896 }
897 seed_entity_ids.sort_unstable();
898 seed_entity_ids.dedup();
899 tracing::debug!(target: "deep_research",
900 sub_query_id,
901 seed_count = seed_entity_ids.len(),
902 "seed entities collected"
903 );
904
905 let all_seed_ids: Vec<i64> = memory_ids
906 .iter()
907 .chain(entity_ids.iter())
908 .copied()
909 .collect();
910
911 if let Ok(graph_results) = traverse_from_memories_with_hops_capped(
913 &conn,
914 &all_seed_ids,
915 namespace,
916 min_weight,
917 max_hops as u32,
918 max_neighbors_per_hop,
919 ) {
920 let seed_score_map: crate::hash::AHashMap<i64, f64> = fused
922 .iter()
923 .map(|(id, s)| {
924 let normalized = if max_possible > 0.0 {
925 s / max_possible
926 } else {
927 0.0
928 };
929 (*id, normalized.clamp(0.0, 1.0))
930 })
931 .collect();
932
933 for (graph_mem_id, hop) in graph_results {
934 if seen_ids.insert(graph_mem_id) {
935 let avg_seed_score: f64 = if seed_score_map.is_empty() {
940 0.5
941 } else {
942 let sum: f64 = seed_score_map.values().sum();
943 sum / seed_score_map.len() as f64
944 };
945 let graph_score =
946 (avg_seed_score * graph_decay.powi(hop as i32)).clamp(0.0, 1.0);
947
948 if graph_score < graph_min_score {
949 continue;
950 }
951
952 if let Ok(Some(row)) = memories::read_full(&conn, graph_mem_id) {
953 let snippet: String = row.body.chars().take(300).collect();
954 hits.push((
955 graph_mem_id,
956 graph_score,
957 "graph".to_string(),
958 snippet,
959 row.body,
960 Some(hop as usize),
961 ));
962 }
963 }
964 }
965 }
966
967 if !seed_entity_ids.is_empty() {
970 let (entity_depth, predecessor) = bfs_with_predecessors(
971 &conn,
972 &seed_entity_ids,
973 namespace,
974 min_weight,
975 max_hops as u32,
976 max_neighbors_per_hop,
977 )
978 .unwrap_or_default();
979
980 tracing::debug!(target: "deep_research",
981 sub_query_id,
982 bfs_nodes = entity_depth.len(),
983 predecessors = predecessor.len(),
984 "BFS complete"
985 );
986
987 let seed_entity_set: HashSet<i64> = seed_entity_ids.iter().copied().collect();
988
989 let all_entity_ids: Vec<i64> = entity_depth.keys().copied().collect();
991 let mut entity_names: crate::hash::AHashMap<i64, String> =
992 crate::hash::AHashMap::with_capacity_and_hasher(
993 all_entity_ids.len(),
994 ahash::RandomState::default(),
995 );
996 for &eid in &all_entity_ids {
997 let name_res: rusqlite::Result<String> = conn.query_row(
998 "SELECT name FROM entities WHERE id = ?1",
999 rusqlite::params![eid],
1000 |r| r.get(0),
1001 );
1002 if let Ok(name) = name_res {
1003 entity_names.insert(eid, name);
1004 }
1005 }
1006
1007 for (&target_id, &_hop) in &entity_depth {
1009 if seed_entity_set.contains(&target_id) {
1010 continue;
1011 }
1012 if !predecessor.contains_key(&target_id) {
1013 continue;
1014 }
1015 if let Some((path_nodes, total_weight)) =
1016 reconstruct_path(target_id, &seed_entity_set, &predecessor, &entity_names)
1017 {
1018 if path_nodes.len() < 2 {
1019 continue;
1020 }
1021 let from = path_nodes
1022 .first()
1023 .map(|n| n.entity.clone())
1024 .unwrap_or_default();
1025 let to = path_nodes
1026 .last()
1027 .map(|n| n.entity.clone())
1028 .unwrap_or_default();
1029 let depth = path_nodes.len();
1030 chains.push(EvidenceChain {
1031 from,
1032 to,
1033 path: path_nodes,
1034 total_weight,
1035 depth,
1036 sub_query_ids: vec![sub_query_id],
1037 });
1038 }
1039 }
1040
1041 chains.sort_by(|a, b| {
1043 b.total_weight
1044 .partial_cmp(&a.total_weight)
1045 .unwrap_or(std::cmp::Ordering::Equal)
1046 });
1047 chains.truncate(20);
1048 tracing::debug!(target: "deep_research",
1049 sub_query_id,
1050 chains_count = chains.len(),
1051 "evidence chains built"
1052 );
1053 }
1054 }
1055
1056 Ok(SubQueryResult {
1057 sub_query_id,
1058 hits,
1059 chains,
1060 })
1061}
1062
1063#[cfg(test)]
1069mod tests {
1070 use super::*;
1071
1072 #[test]
1073 fn test_decompose_and_conjunction() {
1074 let result = decompose_query("A and B", 7);
1075 assert_eq!(result, vec!["A", "B"]);
1076 }
1077
1078 #[test]
1079 fn test_decompose_no_split() {
1080 let result = decompose_query("simple query", 7);
1081 assert_eq!(result, vec!["simple query"]);
1082 }
1083
1084 #[test]
1085 fn test_decompose_three_parts() {
1086 let result = decompose_query("A, B and C", 7);
1087 assert_eq!(result, vec!["A", "B", "C"]);
1088 }
1089
1090 #[test]
1091 fn test_decompose_portuguese_conjunctions() {
1092 let result = decompose_query("A e B", 7);
1093 assert_eq!(result, vec!["A", "B"]);
1094 }
1095
1096 #[test]
1097 fn test_decompose_max_cap() {
1098 let parts: Vec<String> = (0..10).map(|i| format!("part{i}")).collect();
1099 let query = parts.join(", ");
1100 let result = decompose_query(&query, 7);
1101 assert!(
1102 result.len() <= 7,
1103 "expected at most 7 sub-queries, got {}",
1104 result.len()
1105 );
1106 }
1107
1108 #[test]
1109 fn test_decompose_empty_preserves_original() {
1110 let result = decompose_query("", 7);
1111 assert_eq!(result, vec![""]);
1112 }
1113
1114 #[test]
1115 fn test_decompose_semicolons() {
1116 let result = decompose_query("auth design; deployment config; logging", 7);
1117 assert_eq!(result, vec!["auth design", "deployment config", "logging"]);
1118 }
1119
1120 #[test]
1121 fn test_decompose_relational_phrase() {
1122 let result = decompose_query("auth that caused deployment failure", 7);
1123 assert_eq!(result, vec!["auth", "deployment failure"]);
1124 }
1125
1126 #[test]
1127 fn test_sub_query_serialization() {
1128 let sq = SubQuery {
1129 id: 0,
1130 text: "test query".to_string(),
1131 source: "original",
1132 };
1133 let json = serde_json::to_value(&sq).expect("serialization failed");
1134 assert_eq!(json["id"], 0);
1135 assert_eq!(json["text"], "test query");
1136 assert_eq!(json["source"], "original");
1137 }
1138
1139 #[test]
1140 fn test_deep_result_omits_body_when_none() {
1141 let result = DeepResult {
1142 name: "test".to_string(),
1143 score: 0.9,
1144 source: "knn".to_string(),
1145 sub_query_ids: vec![0],
1146 snippet: "snippet".to_string(),
1147 body: None,
1148 hop_distance: None,
1149 };
1150 let json = serde_json::to_string(&result).expect("serialization failed");
1151 assert!(!json.contains("\"body\""), "body must be omitted when None");
1152 }
1153
1154 #[test]
1155 fn test_deep_result_includes_body_when_some() {
1156 let result = DeepResult {
1157 name: "test".to_string(),
1158 score: 0.9,
1159 source: "knn".to_string(),
1160 sub_query_ids: vec![0, 1],
1161 snippet: "snippet".to_string(),
1162 body: Some("full body content".to_string()),
1163 hop_distance: Some(2),
1164 };
1165 let json = serde_json::to_string(&result).expect("serialization failed");
1166 assert!(json.contains("\"body\""), "body must be present when Some");
1167 assert!(json.contains("full body content"));
1168 }
1169
1170 #[test]
1171 fn test_evidence_node_omits_none_fields() {
1172 let node = EvidenceNode {
1173 entity: "auth-module".to_string(),
1174 relation: None,
1175 weight: None,
1176 };
1177 let json = serde_json::to_string(&node).expect("serialization failed");
1178 assert!(
1179 !json.contains("\"relation\""),
1180 "relation must be omitted when None"
1181 );
1182 assert!(
1183 !json.contains("\"weight\""),
1184 "weight must be omitted when None"
1185 );
1186 }
1187
1188 #[test]
1189 fn test_research_stats_serialization() {
1190 let stats = ResearchStats {
1191 sub_queries_total: 3,
1192 sub_queries_completed: 2,
1193 sub_queries_failed: 1,
1194 sub_queries_timed_out: 0,
1195 unique_memories_found: 10,
1196 evidence_chains_found: 2,
1197 elapsed_ms: 1234,
1198 vec_degraded: false,
1199 };
1200 let json = serde_json::to_value(&stats).expect("serialization failed");
1201 assert_eq!(json["sub_queries_total"], 3);
1202 assert_eq!(json["sub_queries_completed"], 2);
1203 assert_eq!(json["sub_queries_failed"], 1);
1204 assert_eq!(json["elapsed_ms"], 1234);
1205 }
1206
1207 #[test]
1208 fn test_deep_research_response_serialization() {
1209 let resp = DeepResearchResponse {
1210 query: "test query".to_string(),
1211 sub_queries: vec![SubQuery {
1212 id: 0,
1213 text: "test query".to_string(),
1214 source: "original",
1215 }],
1216 results: vec![],
1217 evidence_chains: vec![],
1218 graph_context: None,
1219 stats: ResearchStats {
1220 sub_queries_total: 1,
1221 sub_queries_completed: 1,
1222 sub_queries_failed: 0,
1223 sub_queries_timed_out: 0,
1224 unique_memories_found: 0,
1225 evidence_chains_found: 0,
1226 elapsed_ms: 42,
1227 vec_degraded: false,
1228 },
1229 };
1230 let json = serde_json::to_value(&resp).expect("serialization failed");
1231 assert_eq!(json["query"], "test query");
1232 assert!(json["sub_queries"].is_array());
1233 assert!(json["results"].is_array());
1234 assert!(json["evidence_chains"].is_array());
1235 assert_eq!(json["stats"]["elapsed_ms"], 42);
1236 }
1237
1238 #[test]
1242 fn test_distinct_sub_queries_produce_distinct_texts() {
1243 let queries = [
1244 "authentication design decisions",
1245 "deployment configuration and infrastructure",
1246 ];
1247 assert_ne!(queries[0], queries[1]);
1249
1250 let decomposed = decompose_query(
1252 "authentication design decisions; deployment configuration and infrastructure",
1253 7,
1254 );
1255 assert_eq!(decomposed.len(), 2);
1256 assert_ne!(decomposed[0], decomposed[1]);
1257 }
1258
1259 #[test]
1261 fn test_rrf_fuse_via_fusion_module() {
1262 use crate::storage::fusion::rrf_fuse;
1263
1264 let knn_ids: Vec<i64> = vec![1, 2, 3];
1265 let fts_ids: Vec<i64> = vec![2, 1, 4];
1266 let scores = rrf_fuse(&[(1.0, &knn_ids), (1.0, &fts_ids)], 60.0);
1267
1268 let score_1 = scores[&1];
1270 let score_2 = scores[&2];
1271 let score_3 = scores[&3]; let score_4 = scores[&4]; assert!(
1275 score_1 > score_3,
1276 "id 1 (both lists) must beat id 3 (knn-only rank 3)"
1277 );
1278 assert!(
1279 score_2 > score_4,
1280 "id 2 (both lists) must beat id 4 (fts-only rank 3)"
1281 );
1282 }
1283
1284 #[test]
1286 fn test_evidence_chain_has_from_to_and_path() {
1287 let chain = EvidenceChain {
1288 from: "auth-module".to_string(),
1289 to: "jwt-service".to_string(),
1290 path: vec![
1291 EvidenceNode {
1292 entity: "auth-module".to_string(),
1293 relation: None,
1294 weight: None,
1295 },
1296 EvidenceNode {
1297 entity: "token-validator".to_string(),
1298 relation: Some("depends-on".to_string()),
1299 weight: Some(0.9),
1300 },
1301 EvidenceNode {
1302 entity: "jwt-service".to_string(),
1303 relation: Some("uses".to_string()),
1304 weight: Some(0.8),
1305 },
1306 ],
1307 total_weight: 0.72,
1308 depth: 3,
1309 sub_query_ids: vec![0],
1310 };
1311
1312 let json = serde_json::to_value(&chain).expect("serialization failed");
1313 assert!(
1314 json["from"].is_string(),
1315 "evidence chain must have 'from' field"
1316 );
1317 assert!(
1318 json["to"].is_string(),
1319 "evidence chain must have 'to' field"
1320 );
1321 assert!(
1322 json["path"].is_array(),
1323 "evidence chain must have 'path' array"
1324 );
1325 assert_eq!(json["path"].as_array().unwrap().len(), 3);
1326 assert!(json["total_weight"].is_number(), "must have total_weight");
1327 assert_eq!(json["depth"], 3);
1328 }
1329
1330 #[test]
1332 fn test_reconstruct_path_root_to_target_order() {
1333 let seed_set: HashSet<i64> = [10i64].into_iter().collect();
1335 let mut predecessor: PredecessorMap = std::collections::HashMap::new();
1336 predecessor.insert(20, (10, "depends-on".to_string(), 0.9));
1337 predecessor.insert(30, (20, "uses".to_string(), 0.8));
1338 let mut entity_names: crate::hash::AHashMap<i64, String> = crate::hash::AHashMap::default();
1339 entity_names.insert(10, "seed-entity".to_string());
1340 entity_names.insert(20, "middle-entity".to_string());
1341 entity_names.insert(30, "target-entity".to_string());
1342
1343 let result = reconstruct_path(30, &seed_set, &predecessor, &entity_names);
1344 assert!(result.is_some(), "path must be reconstructed");
1345 let (nodes, weight) = result.unwrap();
1346 assert_eq!(nodes.len(), 3);
1348 assert_eq!(nodes[0].entity, "seed-entity");
1349 assert_eq!(nodes[1].entity, "middle-entity");
1350 assert_eq!(nodes[2].entity, "target-entity");
1351 assert!((weight - 0.72).abs() < 1e-6);
1353 }
1354
1355 #[test]
1357 fn test_evidence_chains_single_hop_filtered_out() {
1358 let chain = EvidenceChain {
1360 from: "a".to_string(),
1361 to: "a".to_string(),
1362 path: vec![EvidenceNode {
1363 entity: "a".to_string(),
1364 relation: None,
1365 weight: None,
1366 }],
1367 total_weight: 1.0,
1368 depth: 1,
1369 sub_query_ids: vec![0],
1370 };
1371 let chains = vec![chain];
1373 let retained: Vec<_> = chains.into_iter().filter(|c| c.depth >= 2).collect();
1374 assert!(retained.is_empty(), "depth-1 chains must be filtered out");
1375 }
1376
1377 #[test]
1379 fn test_bfs_with_predecessors_respects_neighbor_cap() {
1380 use crate::graph::bfs_with_predecessors;
1381 use rusqlite::Connection;
1382
1383 let conn = Connection::open_in_memory().unwrap();
1384 conn.execute_batch(
1385 "CREATE TABLE relationships (
1386 source_id INTEGER NOT NULL,
1387 target_id INTEGER NOT NULL,
1388 weight REAL NOT NULL,
1389 namespace TEXT NOT NULL,
1390 relation TEXT NOT NULL DEFAULT 'related'
1391 );",
1392 )
1393 .unwrap();
1394
1395 for target in 2i64..=6 {
1397 conn.execute(
1398 "INSERT INTO relationships (source_id, target_id, weight, namespace) VALUES (?1, ?2, ?3, 'ns')",
1399 rusqlite::params![1i64, target, 1.0f64],
1400 )
1401 .unwrap();
1402 }
1403
1404 let (depth_uncapped, _) = bfs_with_predecessors(&conn, &[1], "ns", 0.0, 1, None).unwrap();
1406 assert_eq!(
1407 depth_uncapped.len() - 1,
1408 5,
1409 "uncapped must discover all 5 neighbours (plus seed)"
1410 );
1411
1412 let (depth_capped, _) = bfs_with_predecessors(&conn, &[1], "ns", 0.0, 1, Some(2)).unwrap();
1414 assert_eq!(
1416 depth_capped.len(),
1417 3,
1418 "capped to 2 must yield seed + 2 neighbours"
1419 );
1420 }
1421}