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