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(
252 args: DeepResearchArgs,
253 llm_backend: crate::cli::LlmBackendChoice,
254 embedding_backend: crate::cli::EmbeddingBackendChoice,
255) -> Result<(), AppError> {
256 tracing::debug!(target: "deep_research", query = %args.query, k = args.k, "starting deep research");
257
258 let paths = AppPaths::resolve(args.db.as_deref())?;
266 crate::storage::connection::ensure_db_ready(&paths)?;
267 let sub_query_texts = decompose_query(&args.query, args.max_sub_queries);
268 let (sub_embeddings, vec_degraded) =
269 compute_sub_embeddings(&paths, &sub_query_texts, embedding_backend, llm_backend);
270
271 let rt = tokio::runtime::Builder::new_multi_thread()
272 .worker_threads(2)
273 .enable_all()
274 .build()
275 .map_err(|e| AppError::Internal(anyhow::anyhow!("failed to build tokio runtime: {e}")))?;
276 rt.block_on(run_async(
277 args,
278 llm_backend,
279 embedding_backend,
280 sub_embeddings,
281 vec_degraded,
282 ))
283}
284
285fn compute_sub_embeddings(
294 paths: &crate::paths::AppPaths,
295 sub_query_texts: &[String],
296 embedding_backend: crate::cli::EmbeddingBackendChoice,
297 llm_backend: crate::cli::LlmBackendChoice,
298) -> (Vec<Option<Arc<Vec<f32>>>>, bool) {
299 output::emit_progress_i18n(
300 "Computing per-sub-query embeddings...",
301 "Calculando embeddings por sub-consulta...",
302 );
303 let mut sub_embeddings: Vec<Option<Arc<Vec<f32>>>> = Vec::with_capacity(sub_query_texts.len());
304 let mut vec_degraded = false;
305 for sq_text in sub_query_texts {
306 match crate::embedder::try_embed_query_with_embedding_choice(
307 &paths.models,
308 sq_text,
309 embedding_backend,
310 llm_backend,
311 ) {
312 Ok((v, _backend)) => sub_embeddings.push(Some(Arc::new(v))),
313 Err(reason) => {
314 tracing::warn!(target: "deep_research", fallback_reason = %reason, reason_code = %reason.reason_code(), "embedding failed for sub-query; falling back to FTS5");
315 sub_embeddings.push(None);
316 vec_degraded = true;
317 }
318 }
319 }
320 (sub_embeddings, vec_degraded)
321}
322
323async fn run_async(
329 args: DeepResearchArgs,
330 _llm_backend: crate::cli::LlmBackendChoice,
331 _embedding_backend: crate::cli::EmbeddingBackendChoice,
332 sub_embeddings: Vec<Option<Arc<Vec<f32>>>>,
333 vec_degraded: bool,
334) -> Result<(), AppError> {
335 let start = std::time::Instant::now();
336
337 if args.query.trim().is_empty() {
338 return Err(AppError::Validation(crate::i18n::validation::empty_query()));
339 }
340
341 if args.max_cost_usd.is_some() && args.mode == "none" {
342 tracing::warn!(target: "deep_research", "--max-cost-usd has no effect without --mode claude-code/codex");
343 }
344
345 let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
346 let paths = AppPaths::resolve(args.db.as_deref())?;
347 crate::storage::connection::ensure_db_ready(&paths)?;
348
349 let sub_query_texts = decompose_query(&args.query, args.max_sub_queries);
351 let sub_queries: Vec<SubQuery> = sub_query_texts
352 .iter()
353 .enumerate()
354 .map(|(i, text)| SubQuery {
355 id: i,
356 text: text.clone(),
357 source: if sub_query_texts.len() == 1 {
358 "original"
359 } else {
360 "decomposed"
361 },
362 })
363 .collect();
364
365 if vec_degraded {
370 tracing::debug!(target: "deep_research", "vector degraded: at least one sub-query fell back to FTS5");
371 }
372
373 let cpu_count = std::thread::available_parallelism()
375 .map(|n| n.get())
376 .unwrap_or(4);
377 let permits = args
378 .max_concurrency
379 .unwrap_or_else(|| cpu_count.min(8))
380 .min(sub_queries.len())
381 .max(1);
382 let semaphore = Arc::new(Semaphore::new(permits));
383 let timeout_dur = std::time::Duration::from_secs(args.timeout);
384
385 let mut join_set: JoinSet<Result<SubQueryResult, (usize, String)>> = JoinSet::new();
386
387 for (idx, sq_text) in sub_query_texts.iter().enumerate() {
388 let sem = Arc::clone(&semaphore);
389 let emb = sub_embeddings[idx].clone();
391 let ns = namespace.clone();
392 let db_path = paths.db.clone();
393 let query_text = sq_text.clone();
394 let k = args.k;
395 let max_hops = args.max_hops;
396 let min_weight = args.min_weight;
397 let rrf_k = args.rrf_k;
398 let graph_decay = args.graph_decay;
399 let graph_min_score = args.graph_min_score;
400 let max_neighbors_per_hop = args.max_neighbors_per_hop;
401
402 join_set.spawn(async move {
403 let _permit = sem
404 .acquire_owned()
405 .await
406 .map_err(|e| (idx, format!("semaphore closed: {e}")))?;
407
408 let result = tokio::time::timeout(timeout_dur, async move {
410 execute_sub_query(
411 idx,
412 &query_text,
413 emb.as_ref().map(|v| v.as_slice()),
414 &ns,
415 &db_path,
416 k,
417 max_hops,
418 min_weight,
419 rrf_k,
420 graph_decay,
421 graph_min_score,
422 max_neighbors_per_hop,
423 )
424 })
425 .await;
426
427 match result {
428 Ok(inner) => inner.map_err(|e| (idx, e)),
429 Err(_) => Err((idx, "timeout".to_string())),
430 }
431 });
432 }
433
434 let mut sub_query_results: Vec<SubQueryResult> = Vec::with_capacity(sub_queries.len());
436 let mut failed_count = 0usize;
437 let mut timed_out_count = 0usize;
438
439 while let Some(join_result) = join_set.join_next().await {
440 match join_result {
441 Ok(Ok(sqr)) => sub_query_results.push(sqr),
442 Ok(Err((_idx, reason))) => {
443 if reason == "timeout" {
444 timed_out_count += 1;
445 } else {
446 failed_count += 1;
447 }
448 tracing::warn!(target: "deep_research", sub_query_id = _idx, reason = %reason, "sub-query failed");
449 }
450 Err(join_err) => {
451 failed_count += 1;
452 if join_err.is_panic() {
453 tracing::error!(target: "deep_research", error = %join_err, "sub-query task panicked");
454 } else {
455 tracing::warn!(target: "deep_research", error = %join_err, "sub-query task cancelled");
456 }
457 }
458 }
459 }
460
461 let mut merged: crate::hash::AHashMap<i64, MergedHit> =
464 crate::hash::AHashMap::with_capacity_and_hasher(
465 sub_query_results.len() * args.k,
466 Default::default(),
467 );
468
469 for sqr in &sub_query_results {
470 for (mem_id, score, source, snippet, body, hop) in &sqr.hits {
471 let entry = merged.entry(*mem_id).or_insert_with(|| {
472 (
473 *score,
474 source.clone(),
475 snippet.clone(),
476 body.clone(),
477 *hop,
478 Vec::new(),
479 )
480 });
481 if *score > entry.0 {
483 entry.0 = *score;
484 entry.1 = source.clone();
485 entry.2 = snippet.clone();
486 entry.3 = body.clone();
487 entry.4 = *hop;
488 }
489 if !entry.5.contains(&sqr.sub_query_id) {
490 entry.5.push(sqr.sub_query_id);
491 }
492 }
493 }
494
495 let conn = open_ro(&paths.db)?;
497 let mut results: Vec<DeepResult> = Vec::with_capacity(merged.len().min(args.max_results));
498
499 let mut ranked: Vec<(i64, MergedHit)> = merged.into_iter().collect();
501 ranked.sort_by(|a, b| {
502 b.1 .0
503 .partial_cmp(&a.1 .0)
504 .unwrap_or(std::cmp::Ordering::Equal)
505 });
506 ranked.truncate(args.max_results);
507
508 for (mem_id, (score, source, snippet, body, hop, sq_ids)) in ranked {
509 let name = match memories::read_full(&conn, mem_id)? {
510 Some(row) => row.name,
511 None => continue,
512 };
513 results.push(DeepResult {
514 name,
515 score,
516 source,
517 sub_query_ids: sq_ids,
518 snippet,
519 body: if args.with_bodies { Some(body) } else { None },
520 hop_distance: hop,
521 });
522 }
523
524 let completed_count = sub_query_results.len();
528 let mut evidence_chains: Vec<EvidenceChain> = Vec::with_capacity(completed_count * 2);
529 let mut seen_chain_keys: HashSet<String> = HashSet::with_capacity(completed_count * 2);
530
531 for sqr in sub_query_results {
532 for chain in sqr.chains {
533 let key = format!("{}->{}", chain.from, chain.to);
535 if seen_chain_keys.insert(key) {
536 evidence_chains.push(chain);
537 }
538 }
539 }
540
541 evidence_chains.retain(|c| c.depth >= 2);
543 evidence_chains.sort_by(|a, b| {
544 b.total_weight
545 .partial_cmp(&a.total_weight)
546 .unwrap_or(std::cmp::Ordering::Equal)
547 });
548
549 let unique_memories = results.len();
550 let evidence_count = evidence_chains.len();
551
552 let graph_context = if !results.is_empty() {
554 let result_names: Vec<&str> = results.iter().map(|r| r.name.as_str()).collect();
555 let mut ctx_entities: Vec<GraphContextEntity> = Vec::with_capacity(results.len());
556 let mut ctx_rels: Vec<GraphContextRel> = Vec::with_capacity(results.len() * 2);
557 let mut seen_entity_ids: crate::hash::AHashSet<i64> =
558 crate::hash::AHashSet::with_capacity_and_hasher(results.len(), Default::default());
559
560 for name in &result_names {
561 if let Ok(Some(eid)) = entities::find_entity_id(&conn, &namespace, name) {
562 if seen_entity_ids.insert(eid) {
563 let etype: String = conn
564 .query_row(
565 "SELECT COALESCE(type,'concept') FROM entities WHERE id = ?1",
566 rusqlite::params![eid],
567 |r| r.get(0),
568 )
569 .unwrap_or_else(|_| "concept".to_string());
570 let degree: u32 = conn
571 .query_row(
572 "SELECT COUNT(*) FROM relationships WHERE source_id = ?1 OR target_id = ?1",
573 rusqlite::params![eid],
574 |r| r.get(0),
575 )
576 .unwrap_or(0);
577 ctx_entities.push(GraphContextEntity {
578 name: name.to_string(),
579 entity_type: etype,
580 degree,
581 });
582 }
583 }
584 }
585
586 let entity_ids: Vec<i64> = seen_entity_ids.iter().copied().collect();
587 if entity_ids.len() >= 2 {
588 let placeholders: String = entity_ids.iter().map(|_| "?").collect::<Vec<_>>().join(",");
589 let sql = format!(
590 "SELECT s.name, t.name, r.relation, r.weight \
591 FROM relationships r \
592 JOIN entities s ON s.id = r.source_id \
593 JOIN entities t ON t.id = r.target_id \
594 WHERE r.source_id IN ({placeholders}) AND r.target_id IN ({placeholders}) \
595 LIMIT 50"
596 );
597 if let Ok(mut stmt) = conn.prepare(&sql) {
598 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> =
599 Vec::with_capacity(entity_ids.len() * 2);
600 for id in &entity_ids {
601 params.push(Box::new(*id));
602 }
603 for id in &entity_ids {
604 params.push(Box::new(*id));
605 }
606 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
607 params.iter().map(|p| p.as_ref()).collect();
608 if let Ok(rows) = stmt.query_map(param_refs.as_slice(), |r| {
609 Ok((
610 r.get::<_, String>(0)?,
611 r.get::<_, String>(1)?,
612 r.get::<_, String>(2)?,
613 r.get::<_, f64>(3)?,
614 ))
615 }) {
616 for row in rows.flatten() {
617 ctx_rels.push(GraphContextRel {
618 from: row.0,
619 to: row.1,
620 relation: row.2,
621 weight: row.3,
622 });
623 }
624 }
625 }
626 }
627
628 if ctx_entities.is_empty() {
629 None
630 } else {
631 Some(GraphContext {
632 entities: ctx_entities,
633 relationships: ctx_rels,
634 })
635 }
636 } else {
637 None
638 };
639
640 tracing::debug!(target: "deep_research",
641 total_results = results.len(),
642 total_chains = evidence_chains.len(),
643 "assembly complete"
644 );
645
646 output::emit_json(&DeepResearchResponse {
648 query: args.query,
649 sub_queries,
650 results,
651 evidence_chains,
652 graph_context,
653 stats: ResearchStats {
654 sub_queries_total: sub_query_texts.len(),
655 sub_queries_completed: completed_count,
656 sub_queries_failed: failed_count,
657 sub_queries_timed_out: timed_out_count,
658 unique_memories_found: unique_memories,
659 evidence_chains_found: evidence_count,
660 elapsed_ms: start.elapsed().as_millis() as u64,
661 vec_degraded,
662 },
663 })?;
664
665 Ok(())
666}
667
668fn decompose_query(query: &str, max: usize) -> Vec<String> {
671 if query.is_empty() {
672 return vec![query.to_string()];
673 }
674
675 let mut parts: Vec<String> = Vec::with_capacity(max);
676
677 let relational = [
679 " that caused ",
680 " depending on ",
681 " related to ",
682 " connected to ",
683 " linked to ",
684 " caused by ",
685 " followed by ",
686 ];
687 let mut text = query.to_string();
688 let mut did_relational_split = false;
689 for phrase in &relational {
690 if text.to_lowercase().contains(phrase) {
691 let lower = text.to_lowercase();
692 if let Some(pos) = lower.find(phrase) {
693 let left = text[..pos].trim().to_string();
694 let right = text[pos + phrase.len()..].trim().to_string();
695 if !left.is_empty() {
696 parts.push(left);
697 }
698 if !right.is_empty() {
699 text = right;
700 }
701 did_relational_split = true;
702 }
703 }
704 }
705 if did_relational_split && !text.is_empty() {
706 parts.push(text.clone());
707 }
708
709 if parts.is_empty() {
711 let semi_parts: Vec<&str> = query.split(';').collect();
713 if semi_parts.len() > 1 {
714 for p in &semi_parts {
715 let trimmed = p.trim();
716 if !trimmed.is_empty() {
717 parts.push(trimmed.to_string());
718 }
719 }
720 } else {
721 let normalized = query
724 .replace(" and ", ", ")
725 .replace(" AND ", ", ")
726 .replace(" e ", ", ")
727 .replace(" E ", ", ");
728 let comma_parts: Vec<&str> = normalized.split(',').collect();
729 if comma_parts.len() > 1 {
730 for p in &comma_parts {
731 let trimmed = p.trim();
732 if !trimmed.is_empty() {
733 parts.push(trimmed.to_string());
734 }
735 }
736 }
737 }
738 }
739
740 if parts.is_empty() {
742 let words: Vec<&str> = query.split_whitespace().filter(|w| w.len() > 2).collect();
743 if words.len() >= 3 {
744 parts.push(query.to_string());
745 parts.push(format!("{} {}", words[0], words[1]));
746 parts.push(format!(
747 "{} {}",
748 words[words.len() - 2],
749 words[words.len() - 1]
750 ));
751 }
752 }
753
754 if parts.is_empty() {
755 return vec![query.to_string()];
756 }
757
758 parts.truncate(max);
760 parts
761}
762
763fn reconstruct_path(
767 target_id: i64,
768 seed_entity_ids: &HashSet<i64>,
769 predecessor: &PredecessorMap,
770 entity_names: &crate::hash::AHashMap<i64, String>,
771) -> Option<(Vec<EvidenceNode>, f64)> {
772 let mut path_ids: Vec<(i64, Option<String>, Option<f64>)> = Vec::with_capacity(8);
773 let mut total_weight = 1.0_f64;
774 let mut current = target_id;
775
776 loop {
777 if seed_entity_ids.contains(¤t) {
778 break;
779 }
780 let (parent, relation, weight) = predecessor.get(¤t)?;
781 total_weight *= weight;
782 path_ids.push((current, Some(relation.clone()), Some(*weight)));
783 current = *parent;
784 }
785 path_ids.push((current, None, None));
787
788 path_ids.reverse();
790
791 let nodes: Vec<EvidenceNode> = path_ids
792 .into_iter()
793 .map(|(id, relation, weight)| EvidenceNode {
794 entity: entity_names
795 .get(&id)
796 .cloned()
797 .unwrap_or_else(|| format!("entity-{id}")),
798 relation,
799 weight,
800 })
801 .collect();
802
803 Some((nodes, total_weight))
804}
805
806#[allow(clippy::too_many_arguments)]
816fn execute_sub_query(
817 sub_query_id: usize,
818 query_text: &str,
819 embedding: Option<&[f32]>,
820 namespace: &str,
821 db_path: &std::path::Path,
822 k: usize,
823 max_hops: usize,
824 min_weight: f64,
825 rrf_k: f64,
826 graph_decay: f64,
827 graph_min_score: f64,
828 max_neighbors_per_hop: Option<usize>,
829) -> Result<SubQueryResult, String> {
830 let conn = open_ro(db_path).map_err(|e| format!("failed to open db: {e}"))?;
831
832 let mut hits: Vec<(i64, f64, String, String, String, Option<usize>)> =
833 Vec::with_capacity(k * 2);
834 let mut seen_ids: crate::hash::AHashSet<i64> =
835 crate::hash::AHashSet::with_capacity_and_hasher(k * 2, Default::default());
836
837 let (knn_ids, knn_distance_map) = if let Some(emb) = embedding {
841 let knn_results = memories::knn_search(&conn, emb, &[namespace.to_string()], None, k)
842 .map_err(|e| format!("knn_search failed: {e}"))?;
843 let ids: Vec<i64> = knn_results.iter().map(|(id, _)| *id).collect();
844 tracing::debug!(target: "deep_research", sub_query_id, knn_count = ids.len(), "KNN complete");
845 let dist_map: crate::hash::AHashMap<i64, f64> = knn_results
846 .iter()
847 .map(|(id, dist)| (*id, *dist as f64))
848 .collect();
849 (ids, dist_map)
850 } else {
851 tracing::debug!(target: "deep_research", sub_query_id, "KNN skipped (no embedding); FTS5-only");
852 (vec![], crate::hash::AHashMap::default())
853 };
854
855 let fts_results = match memories::fts_search(&conn, query_text, namespace, None, k) {
857 Ok(rows) => rows,
858 Err(e) => {
859 tracing::warn!(target: "deep_research",
860 sub_query_id,
861 "FTS5 search failed, continuing with KNN only: {e}"
862 );
863 vec![]
864 }
865 };
866 let fts_ids: Vec<i64> = fts_results.iter().map(|r| r.id).collect();
867 tracing::debug!(target: "deep_research", sub_query_id, fts_count = fts_ids.len(), "FTS complete");
868
869 let rrf_scores = rrf_fuse(&[(1.0, &knn_ids), (1.0, &fts_ids)], rrf_k);
871 let max_possible = rrf_max_possible(&[1.0, 1.0], rrf_k);
872
873 let mut fused: Vec<(i64, f64)> = rrf_scores.into_iter().collect();
875 fused.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
876 fused.truncate(k * 2);
877 tracing::debug!(target: "deep_research",
878 sub_query_id,
879 fused_count = fused.len(),
880 "RRF fusion complete"
881 );
882
883 if fused.is_empty() && !knn_ids.is_empty() {
884 tracing::warn!(target: "deep_research", sub_query_id, knn_count = knn_ids.len(), fts_count = fts_ids.len(),
885 "RRF fusion returned 0 results despite KNN/FTS hits; consider lowering --graph-min-score");
886 }
887
888 for (memory_id, combined_score) in &fused {
889 if seen_ids.insert(*memory_id) {
890 let normalized = if max_possible > 0.0 {
891 combined_score / max_possible
892 } else {
893 0.0
894 };
895 let score = normalized.clamp(0.0, 1.0);
896 let in_knn = knn_distance_map.contains_key(memory_id);
897 let in_fts = fts_ids.contains(memory_id);
898 let source = match (in_knn, in_fts) {
899 (true, true) => "hybrid",
900 (true, false) => "knn",
901 (false, true) => "fts",
902 (false, false) => "graph",
903 };
904 if let Ok(Some(row)) = memories::read_full(&conn, *memory_id) {
905 let snippet: String = row.body.chars().take(300).collect();
906 hits.push((
907 *memory_id,
908 score,
909 source.to_string(),
910 snippet,
911 row.body,
912 None,
913 ));
914 }
915 }
916 }
917
918 let memory_ids: Vec<i64> = hits.iter().map(|(id, ..)| *id).collect();
921 let mut chains: Vec<EvidenceChain> = Vec::with_capacity(memory_ids.len());
922
923 if !memory_ids.is_empty() && max_hops > 0 {
924 let entity_ids: Vec<i64> = if let Some(emb) = embedding {
926 entities::knn_search(&conn, emb, namespace, 5)
927 .inspect_err(|e| tracing::warn!(target: "deep_research", error = %e, "entity KNN search failed, skipping graph seed"))
928 .unwrap_or_default()
929 .iter()
930 .map(|(id, _)| *id)
931 .collect()
932 } else {
933 vec![]
934 };
935
936 let top_seed_count = 5.min(memory_ids.len());
939 let top_memory_ids = &memory_ids[..top_seed_count];
940 let mut seed_entity_ids: Vec<i64> = entity_ids.clone();
941 for &mem_id in top_memory_ids {
942 let mut stmt = conn
943 .prepare_cached("SELECT entity_id FROM memory_entities WHERE memory_id = ?1")
944 .map_err(|e| format!("prepare failed: {e}"))?;
945 let ids: Vec<i64> = stmt
946 .query_map(rusqlite::params![mem_id], |r| r.get(0))
947 .map_err(|e| format!("query failed: {e}"))?
948 .filter_map(|r| r.ok())
949 .collect();
950 seed_entity_ids.extend(ids);
951 }
952 seed_entity_ids.sort_unstable();
953 seed_entity_ids.dedup();
954 tracing::debug!(target: "deep_research",
955 sub_query_id,
956 seed_count = seed_entity_ids.len(),
957 "seed entities collected"
958 );
959
960 let all_seed_ids: Vec<i64> = memory_ids
961 .iter()
962 .chain(entity_ids.iter())
963 .copied()
964 .collect();
965
966 if let Ok(graph_results) = traverse_from_memories_with_hops_capped(
968 &conn,
969 &all_seed_ids,
970 namespace,
971 min_weight,
972 max_hops as u32,
973 max_neighbors_per_hop,
974 ) {
975 let seed_score_map: crate::hash::AHashMap<i64, f64> = fused
977 .iter()
978 .map(|(id, s)| {
979 let normalized = if max_possible > 0.0 {
980 s / max_possible
981 } else {
982 0.0
983 };
984 (*id, normalized.clamp(0.0, 1.0))
985 })
986 .collect();
987
988 for (graph_mem_id, hop) in graph_results {
989 if seen_ids.insert(graph_mem_id) {
990 let avg_seed_score: f64 = if seed_score_map.is_empty() {
995 0.5
996 } else {
997 let sum: f64 = seed_score_map.values().sum();
998 sum / seed_score_map.len() as f64
999 };
1000 let graph_score =
1001 (avg_seed_score * graph_decay.powi(hop as i32)).clamp(0.0, 1.0);
1002
1003 if graph_score < graph_min_score {
1004 continue;
1005 }
1006
1007 if let Ok(Some(row)) = memories::read_full(&conn, graph_mem_id) {
1008 let snippet: String = row.body.chars().take(300).collect();
1009 hits.push((
1010 graph_mem_id,
1011 graph_score,
1012 "graph".to_string(),
1013 snippet,
1014 row.body,
1015 Some(hop as usize),
1016 ));
1017 }
1018 }
1019 }
1020 }
1021
1022 if !seed_entity_ids.is_empty() {
1025 let (entity_depth, predecessor) = bfs_with_predecessors(
1026 &conn,
1027 &seed_entity_ids,
1028 namespace,
1029 min_weight,
1030 max_hops as u32,
1031 max_neighbors_per_hop,
1032 )
1033 .unwrap_or_default();
1034
1035 tracing::debug!(target: "deep_research",
1036 sub_query_id,
1037 bfs_nodes = entity_depth.len(),
1038 predecessors = predecessor.len(),
1039 "BFS complete"
1040 );
1041
1042 let seed_entity_set: HashSet<i64> = seed_entity_ids.iter().copied().collect();
1043
1044 let all_entity_ids: Vec<i64> = entity_depth.keys().copied().collect();
1046 let mut entity_names: crate::hash::AHashMap<i64, String> =
1047 crate::hash::AHashMap::with_capacity_and_hasher(
1048 all_entity_ids.len(),
1049 ahash::RandomState::default(),
1050 );
1051 for &eid in &all_entity_ids {
1052 let name_res: rusqlite::Result<String> = conn.query_row(
1053 "SELECT name FROM entities WHERE id = ?1",
1054 rusqlite::params![eid],
1055 |r| r.get(0),
1056 );
1057 if let Ok(name) = name_res {
1058 entity_names.insert(eid, name);
1059 }
1060 }
1061
1062 for (&target_id, &_hop) in &entity_depth {
1064 if seed_entity_set.contains(&target_id) {
1065 continue;
1066 }
1067 if !predecessor.contains_key(&target_id) {
1068 continue;
1069 }
1070 if let Some((path_nodes, total_weight)) =
1071 reconstruct_path(target_id, &seed_entity_set, &predecessor, &entity_names)
1072 {
1073 if path_nodes.len() < 2 {
1074 continue;
1075 }
1076 let from = path_nodes
1077 .first()
1078 .map(|n| n.entity.clone())
1079 .unwrap_or_default();
1080 let to = path_nodes
1081 .last()
1082 .map(|n| n.entity.clone())
1083 .unwrap_or_default();
1084 let depth = path_nodes.len();
1085 chains.push(EvidenceChain {
1086 from,
1087 to,
1088 path: path_nodes,
1089 total_weight,
1090 depth,
1091 sub_query_ids: vec![sub_query_id],
1092 });
1093 }
1094 }
1095
1096 chains.sort_by(|a, b| {
1098 b.total_weight
1099 .partial_cmp(&a.total_weight)
1100 .unwrap_or(std::cmp::Ordering::Equal)
1101 });
1102 chains.truncate(20);
1103 tracing::debug!(target: "deep_research",
1104 sub_query_id,
1105 chains_count = chains.len(),
1106 "evidence chains built"
1107 );
1108 }
1109 }
1110
1111 Ok(SubQueryResult {
1112 sub_query_id,
1113 hits,
1114 chains,
1115 })
1116}
1117
1118#[cfg(test)]
1124mod tests {
1125 use super::*;
1126
1127 #[test]
1128 fn test_decompose_and_conjunction() {
1129 let result = decompose_query("A and B", 7);
1130 assert_eq!(result, vec!["A", "B"]);
1131 }
1132
1133 #[test]
1134 fn test_decompose_no_split() {
1135 let result = decompose_query("simple query", 7);
1136 assert_eq!(result, vec!["simple query"]);
1137 }
1138
1139 #[test]
1140 fn test_decompose_three_parts() {
1141 let result = decompose_query("A, B and C", 7);
1142 assert_eq!(result, vec!["A", "B", "C"]);
1143 }
1144
1145 #[test]
1146 fn test_decompose_portuguese_conjunctions() {
1147 let result = decompose_query("A e B", 7);
1148 assert_eq!(result, vec!["A", "B"]);
1149 }
1150
1151 #[test]
1152 fn test_decompose_max_cap() {
1153 let parts: Vec<String> = (0..10).map(|i| format!("part{i}")).collect();
1154 let query = parts.join(", ");
1155 let result = decompose_query(&query, 7);
1156 assert!(
1157 result.len() <= 7,
1158 "expected at most 7 sub-queries, got {}",
1159 result.len()
1160 );
1161 }
1162
1163 #[test]
1164 fn test_decompose_empty_preserves_original() {
1165 let result = decompose_query("", 7);
1166 assert_eq!(result, vec![""]);
1167 }
1168
1169 #[test]
1170 fn test_decompose_semicolons() {
1171 let result = decompose_query("auth design; deployment config; logging", 7);
1172 assert_eq!(result, vec!["auth design", "deployment config", "logging"]);
1173 }
1174
1175 #[test]
1176 fn test_decompose_relational_phrase() {
1177 let result = decompose_query("auth that caused deployment failure", 7);
1178 assert_eq!(result, vec!["auth", "deployment failure"]);
1179 }
1180
1181 #[test]
1182 fn test_sub_query_serialization() {
1183 let sq = SubQuery {
1184 id: 0,
1185 text: "test query".to_string(),
1186 source: "original",
1187 };
1188 let json = serde_json::to_value(&sq).expect("serialization failed");
1189 assert_eq!(json["id"], 0);
1190 assert_eq!(json["text"], "test query");
1191 assert_eq!(json["source"], "original");
1192 }
1193
1194 #[test]
1195 fn test_deep_result_omits_body_when_none() {
1196 let result = DeepResult {
1197 name: "test".to_string(),
1198 score: 0.9,
1199 source: "knn".to_string(),
1200 sub_query_ids: vec![0],
1201 snippet: "snippet".to_string(),
1202 body: None,
1203 hop_distance: None,
1204 };
1205 let json = serde_json::to_string(&result).expect("serialization failed");
1206 assert!(!json.contains("\"body\""), "body must be omitted when None");
1207 }
1208
1209 #[test]
1210 fn test_deep_result_includes_body_when_some() {
1211 let result = DeepResult {
1212 name: "test".to_string(),
1213 score: 0.9,
1214 source: "knn".to_string(),
1215 sub_query_ids: vec![0, 1],
1216 snippet: "snippet".to_string(),
1217 body: Some("full body content".to_string()),
1218 hop_distance: Some(2),
1219 };
1220 let json = serde_json::to_string(&result).expect("serialization failed");
1221 assert!(json.contains("\"body\""), "body must be present when Some");
1222 assert!(json.contains("full body content"));
1223 }
1224
1225 #[test]
1226 fn test_evidence_node_omits_none_fields() {
1227 let node = EvidenceNode {
1228 entity: "auth-module".to_string(),
1229 relation: None,
1230 weight: None,
1231 };
1232 let json = serde_json::to_string(&node).expect("serialization failed");
1233 assert!(
1234 !json.contains("\"relation\""),
1235 "relation must be omitted when None"
1236 );
1237 assert!(
1238 !json.contains("\"weight\""),
1239 "weight must be omitted when None"
1240 );
1241 }
1242
1243 #[test]
1244 fn test_research_stats_serialization() {
1245 let stats = ResearchStats {
1246 sub_queries_total: 3,
1247 sub_queries_completed: 2,
1248 sub_queries_failed: 1,
1249 sub_queries_timed_out: 0,
1250 unique_memories_found: 10,
1251 evidence_chains_found: 2,
1252 elapsed_ms: 1234,
1253 vec_degraded: false,
1254 };
1255 let json = serde_json::to_value(&stats).expect("serialization failed");
1256 assert_eq!(json["sub_queries_total"], 3);
1257 assert_eq!(json["sub_queries_completed"], 2);
1258 assert_eq!(json["sub_queries_failed"], 1);
1259 assert_eq!(json["elapsed_ms"], 1234);
1260 }
1261
1262 #[test]
1263 fn test_deep_research_response_serialization() {
1264 let resp = DeepResearchResponse {
1265 query: "test query".to_string(),
1266 sub_queries: vec![SubQuery {
1267 id: 0,
1268 text: "test query".to_string(),
1269 source: "original",
1270 }],
1271 results: vec![],
1272 evidence_chains: vec![],
1273 graph_context: None,
1274 stats: ResearchStats {
1275 sub_queries_total: 1,
1276 sub_queries_completed: 1,
1277 sub_queries_failed: 0,
1278 sub_queries_timed_out: 0,
1279 unique_memories_found: 0,
1280 evidence_chains_found: 0,
1281 elapsed_ms: 42,
1282 vec_degraded: false,
1283 },
1284 };
1285 let json = serde_json::to_value(&resp).expect("serialization failed");
1286 assert_eq!(json["query"], "test query");
1287 assert!(json["sub_queries"].is_array());
1288 assert!(json["results"].is_array());
1289 assert!(json["evidence_chains"].is_array());
1290 assert_eq!(json["stats"]["elapsed_ms"], 42);
1291 }
1292
1293 #[test]
1297 fn test_distinct_sub_queries_produce_distinct_texts() {
1298 let queries = [
1299 "authentication design decisions",
1300 "deployment configuration and infrastructure",
1301 ];
1302 assert_ne!(queries[0], queries[1]);
1304
1305 let decomposed = decompose_query(
1307 "authentication design decisions; deployment configuration and infrastructure",
1308 7,
1309 );
1310 assert_eq!(decomposed.len(), 2);
1311 assert_ne!(decomposed[0], decomposed[1]);
1312 }
1313
1314 #[test]
1316 fn test_rrf_fuse_via_fusion_module() {
1317 use crate::storage::fusion::rrf_fuse;
1318
1319 let knn_ids: Vec<i64> = vec![1, 2, 3];
1320 let fts_ids: Vec<i64> = vec![2, 1, 4];
1321 let scores = rrf_fuse(&[(1.0, &knn_ids), (1.0, &fts_ids)], 60.0);
1322
1323 let score_1 = scores[&1];
1325 let score_2 = scores[&2];
1326 let score_3 = scores[&3]; let score_4 = scores[&4]; assert!(
1330 score_1 > score_3,
1331 "id 1 (both lists) must beat id 3 (knn-only rank 3)"
1332 );
1333 assert!(
1334 score_2 > score_4,
1335 "id 2 (both lists) must beat id 4 (fts-only rank 3)"
1336 );
1337 }
1338
1339 #[test]
1341 fn test_evidence_chain_has_from_to_and_path() {
1342 let chain = EvidenceChain {
1343 from: "auth-module".to_string(),
1344 to: "jwt-service".to_string(),
1345 path: vec![
1346 EvidenceNode {
1347 entity: "auth-module".to_string(),
1348 relation: None,
1349 weight: None,
1350 },
1351 EvidenceNode {
1352 entity: "token-validator".to_string(),
1353 relation: Some("depends-on".to_string()),
1354 weight: Some(0.9),
1355 },
1356 EvidenceNode {
1357 entity: "jwt-service".to_string(),
1358 relation: Some("uses".to_string()),
1359 weight: Some(0.8),
1360 },
1361 ],
1362 total_weight: 0.72,
1363 depth: 3,
1364 sub_query_ids: vec![0],
1365 };
1366
1367 let json = serde_json::to_value(&chain).expect("serialization failed");
1368 assert!(
1369 json["from"].is_string(),
1370 "evidence chain must have 'from' field"
1371 );
1372 assert!(
1373 json["to"].is_string(),
1374 "evidence chain must have 'to' field"
1375 );
1376 assert!(
1377 json["path"].is_array(),
1378 "evidence chain must have 'path' array"
1379 );
1380 assert_eq!(json["path"].as_array().unwrap().len(), 3);
1381 assert!(json["total_weight"].is_number(), "must have total_weight");
1382 assert_eq!(json["depth"], 3);
1383 }
1384
1385 #[test]
1387 fn test_reconstruct_path_root_to_target_order() {
1388 let seed_set: HashSet<i64> = [10i64].into_iter().collect();
1390 let mut predecessor: PredecessorMap = std::collections::HashMap::new();
1391 predecessor.insert(20, (10, "depends-on".to_string(), 0.9));
1392 predecessor.insert(30, (20, "uses".to_string(), 0.8));
1393 let mut entity_names: crate::hash::AHashMap<i64, String> = crate::hash::AHashMap::default();
1394 entity_names.insert(10, "seed-entity".to_string());
1395 entity_names.insert(20, "middle-entity".to_string());
1396 entity_names.insert(30, "target-entity".to_string());
1397
1398 let result = reconstruct_path(30, &seed_set, &predecessor, &entity_names);
1399 assert!(result.is_some(), "path must be reconstructed");
1400 let (nodes, weight) = result.unwrap();
1401 assert_eq!(nodes.len(), 3);
1403 assert_eq!(nodes[0].entity, "seed-entity");
1404 assert_eq!(nodes[1].entity, "middle-entity");
1405 assert_eq!(nodes[2].entity, "target-entity");
1406 assert!((weight - 0.72).abs() < 1e-6);
1408 }
1409
1410 #[test]
1412 fn test_evidence_chains_single_hop_filtered_out() {
1413 let chain = EvidenceChain {
1415 from: "a".to_string(),
1416 to: "a".to_string(),
1417 path: vec![EvidenceNode {
1418 entity: "a".to_string(),
1419 relation: None,
1420 weight: None,
1421 }],
1422 total_weight: 1.0,
1423 depth: 1,
1424 sub_query_ids: vec![0],
1425 };
1426 let chains = vec![chain];
1428 let retained: Vec<_> = chains.into_iter().filter(|c| c.depth >= 2).collect();
1429 assert!(retained.is_empty(), "depth-1 chains must be filtered out");
1430 }
1431
1432 #[test]
1434 fn test_bfs_with_predecessors_respects_neighbor_cap() {
1435 use crate::graph::bfs_with_predecessors;
1436 use rusqlite::Connection;
1437
1438 let conn = Connection::open_in_memory().unwrap();
1439 conn.execute_batch(
1440 "CREATE TABLE relationships (
1441 source_id INTEGER NOT NULL,
1442 target_id INTEGER NOT NULL,
1443 weight REAL NOT NULL,
1444 namespace TEXT NOT NULL,
1445 relation TEXT NOT NULL DEFAULT 'related'
1446 );",
1447 )
1448 .unwrap();
1449
1450 for target in 2i64..=6 {
1452 conn.execute(
1453 "INSERT INTO relationships (source_id, target_id, weight, namespace) VALUES (?1, ?2, ?3, 'ns')",
1454 rusqlite::params![1i64, target, 1.0f64],
1455 )
1456 .unwrap();
1457 }
1458
1459 let (depth_uncapped, _) = bfs_with_predecessors(&conn, &[1], "ns", 0.0, 1, None).unwrap();
1461 assert_eq!(
1462 depth_uncapped.len() - 1,
1463 5,
1464 "uncapped must discover all 5 neighbours (plus seed)"
1465 );
1466
1467 let (depth_capped, _) = bfs_with_predecessors(&conn, &[1], "ns", 0.0, 1, Some(2)).unwrap();
1469 assert_eq!(
1471 depth_capped.len(),
1472 3,
1473 "capped to 2 must yield seed + 2 neighbours"
1474 );
1475 }
1476}