1use crate::errors::AppError;
8use crate::output;
9use crate::paths::AppPaths;
10use crate::storage::connection::open_ro;
11use crate::storage::{entities, memories};
12
13use serde::Serialize;
14use std::collections::HashSet;
15use std::sync::Arc;
16use tokio::sync::Semaphore;
17use tokio::task::JoinSet;
18
19mod pipeline;
20use pipeline::{
21 compute_sub_embeddings, execute_sub_query, resolve_sub_queries,
22};
23
24#[cfg(test)]
25use pipeline::{{decompose_query, decompose_query_with_sources}};
26
27#[derive(clap::Args)]
29#[command(
30 about = "Deep parallel multi-hop GraphRAG research via query decomposition",
31 after_long_help = "CONTRACT:\n \
32 stdout = pretty JSON envelope only (machine-readable).\n \
33 stderr = tracing / progress / diagnostics only.\n \
34 Never redirect with `&>` or `2>&1` into the same file as stdout — that\n \
35 contaminates the JSON and breaks jaq/jq. Prefer:\n \
36 sqlite-graphrag deep-research \"q\" > out.json 2>/dev/null\n \
37 or --output out.json (atomic write via atomwrite algorithm).\n\n\
38EXAMPLES:\n \
39 # Basic deep research (single-token queries auto-expand into aspects)\n \
40 sqlite-graphrag deep-research \"danilo\"\n\n \
41 # With custom parameters\n \
42 sqlite-graphrag deep-research \"auth\" --k 20 --max-hops 3 --max-sub-queries 7\n\n \
43 # Include full memory bodies in output\n \
44 sqlite-graphrag deep-research \"auth\" --with-bodies\n\n \
45 # Manual sub-queries (one query per line)\n \
46 sqlite-graphrag deep-research \"danilo\" --sub-query-strategy manual \\\n \
47 --sub-queries-file aspects.txt\n\n \
48 # Atomic JSON file (crash-safe; preferred for large --with-bodies runs)\n \
49 sqlite-graphrag deep-research \"auth\" --output /tmp/dr.json\n\n \
50 # Tune RRF and graph scoring\n \
51 sqlite-graphrag deep-research \"auth and deployment\" --rrf-k 60 --graph-decay 0.7"
52)]
53pub struct DeepResearchArgs {
54 #[arg(
56 value_name = "QUERY",
57 allow_hyphen_values = true,
58 help = "Research query to decompose and search"
59 )]
60 pub query: String,
61 #[arg(
63 long,
64 short,
65 aliases = ["limit", "top-k"],
66 default_value_t = 20,
67 help = "Results per sub-query (Recall@20 captures 95%+ relevant hits)"
68 )]
69 pub k: usize,
70 #[arg(
72 long,
73 default_value_t = 7,
74 help = "Maximum sub-queries (covers complex multi-hop queries)"
75 )]
76 pub max_sub_queries: usize,
77 #[arg(
79 long,
80 default_value_t = 3,
81 help = "Multi-hop graph traversal depth (sweet spot: 2-3 hops)"
82 )]
83 pub max_hops: usize,
84 #[arg(
86 long,
87 default_value_t = 0.3,
88 help = "Minimum edge weight for graph traversal"
89 )]
90 pub min_weight: f64,
91 #[arg(long, help = "Maximum concurrent sub-queries (default: min(cpus, 8))")]
93 pub max_concurrency: Option<usize>,
94 #[arg(long, default_value_t = 30, help = "Timeout per sub-query in seconds")]
96 pub timeout: u64,
97 #[arg(
99 long,
100 default_value_t = false,
101 help = "Include full memory bodies in results"
102 )]
103 pub with_bodies: bool,
104 #[arg(
106 long,
107 default_value_t = 50,
108 help = "Maximum results after deduplication"
109 )]
110 pub max_results: usize,
111 #[arg(
113 long,
114 default_value_t = 60.0,
115 help = "RRF k parameter (higher = less weight on top ranks)"
116 )]
117 pub rrf_k: f64,
118 #[arg(
120 long,
121 default_value_t = 0.7,
122 help = "Graph score decay factor per hop (0.0-1.0)"
123 )]
124 pub graph_decay: f64,
125 #[arg(
127 long,
128 default_value_t = 0.05,
129 help = "Minimum score threshold for graph-expanded results"
130 )]
131 pub graph_min_score: f64,
132 #[arg(
134 long,
135 help = "Limit neighbours per entity per hop for graph traversal (default: unlimited)"
136 )]
137 pub max_neighbors_per_hop: Option<usize>,
138 #[arg(
140 long,
141 help = "Namespace (flag / XDG namespace.default / global)"
142 )]
143 pub namespace: Option<String>,
144 #[arg(long, default_value = "none", value_parser = ["none"], hide = true)]
146 pub mode: String,
147 #[arg(
149 long,
150 value_name = "USD",
151 help = "Max LLM cost in USD (effective with --mode claude-code/codex)"
152 )]
153 pub max_cost_usd: Option<f64>,
154 #[arg(long, hide = true)]
156 pub json: bool,
157 #[arg(long)]
159 pub db: Option<String>,
160 #[arg(
163 long,
164 default_value = "heuristic",
165 value_parser = ["heuristic", "manual"],
166 help = "Sub-query strategy: heuristic (default) or manual"
167 )]
168 pub sub_query_strategy: String,
169 #[arg(
172 long,
173 value_name = "PATH",
174 help = "File with one sub-query per line (manual strategy)"
175 )]
176 pub sub_queries_file: Option<std::path::PathBuf>,
177 #[arg(
182 short = 'o',
183 long,
184 value_name = "PATH",
185 help = "Atomic JSON output path (atomwrite algorithm; short -o)"
186 )]
187 pub output: Option<std::path::PathBuf>,
188}
189
190#[derive(Serialize)]
191pub(super) struct SubQuery {
192 pub(super) id: usize,
193 pub(super) text: String,
194 pub(super) source: &'static str,
195}
196
197#[derive(Serialize)]
198struct DeepResult {
199 name: String,
200 score: f64,
201 source: String,
202 sub_query_ids: Vec<usize>,
203 snippet: String,
204 #[serde(skip_serializing_if = "Option::is_none")]
205 body: Option<String>,
206 hop_distance: Option<usize>,
207}
208
209#[derive(Serialize, Clone)]
211pub(super) struct EvidenceNode {
212 pub(super) entity: String,
213 #[serde(skip_serializing_if = "Option::is_none")]
214 pub(super) relation: Option<String>,
215 #[serde(skip_serializing_if = "Option::is_none")]
216 pub(super) weight: Option<f64>,
217}
218
219#[derive(Serialize)]
228pub(super) struct EvidenceChain {
229 pub(super) from: String,
230 pub(super) to: String,
231 pub(super) path: Vec<EvidenceNode>,
232 pub(super) total_weight: f64,
233 pub(super) depth: usize,
234 pub(super) sub_query_ids: Vec<usize>,
235}
236
237#[derive(Serialize)]
238struct ResearchStats {
239 sub_queries_total: usize,
240 sub_queries_completed: usize,
241 sub_queries_failed: usize,
242 sub_queries_timed_out: usize,
243 unique_memories_found: usize,
244 evidence_chains_found: usize,
245 elapsed_ms: u64,
246 vec_degraded: bool,
247}
248
249#[derive(Serialize)]
250struct GraphContextEntity {
251 name: String,
252 entity_type: String,
253 degree: u32,
254}
255
256#[derive(Serialize)]
257struct GraphContextRel {
258 from: String,
259 to: String,
260 relation: String,
261 weight: f64,
262}
263
264#[derive(Serialize)]
265struct GraphContext {
266 entities: Vec<GraphContextEntity>,
267 relationships: Vec<GraphContextRel>,
268}
269
270#[derive(Serialize)]
271struct DeepResearchResponse {
272 query: String,
273 sub_queries: Vec<SubQuery>,
274 results: Vec<DeepResult>,
275 evidence_chains: Vec<EvidenceChain>,
276 #[serde(skip_serializing_if = "Option::is_none")]
277 graph_context: Option<GraphContext>,
278 stats: ResearchStats,
279}
280
281type MergedHit = (f64, String, String, String, Option<usize>, Vec<usize>);
283
284pub(super) struct SubQueryResult {
286 pub(super) sub_query_id: usize,
287 pub(super) hits: Vec<(i64, f64, String, String, String, Option<usize>)>,
289 pub(super) chains: Vec<EvidenceChain>,
291}
292
293#[tracing::instrument(skip_all, level = "debug", name = "deep_research")]
295pub fn run(
296 args: DeepResearchArgs,
297 llm_backend: crate::cli::LlmBackendChoice,
298 embedding_backend: crate::cli::EmbeddingBackendChoice,
299) -> Result<(), AppError> {
300 tracing::debug!(target: "deep_research", query = %args.query, k = args.k, "starting deep research");
301
302 let paths = AppPaths::resolve(args.db.as_deref())?;
310 crate::storage::connection::ensure_db_ready(&paths)?;
311 let sub_query_plan = resolve_sub_queries(&args)?;
313 let sub_query_texts: Vec<String> = sub_query_plan.iter().map(|s| s.text.clone()).collect();
314 let (sub_embeddings, vec_degraded) =
315 compute_sub_embeddings(&paths, &sub_query_texts, embedding_backend, llm_backend);
316
317 let rt = tokio::runtime::Builder::new_multi_thread()
318 .worker_threads(2)
319 .enable_all()
320 .build()
321 .map_err(|e| AppError::Internal(anyhow::anyhow!("failed to build tokio runtime: {e}")))?;
322 rt.block_on(run_async(
323 args,
324 llm_backend,
325 embedding_backend,
326 sub_query_plan,
327 sub_embeddings,
328 vec_degraded,
329 ))
330}
331
332async fn run_async(
340 args: DeepResearchArgs,
341 _llm_backend: crate::cli::LlmBackendChoice,
342 _embedding_backend: crate::cli::EmbeddingBackendChoice,
343 sub_queries: Vec<SubQuery>,
344 sub_embeddings: Vec<Option<Arc<Vec<f32>>>>,
345 vec_degraded: bool,
346) -> Result<(), AppError> {
347 let start = std::time::Instant::now();
348
349 if args.query.trim().is_empty() {
350 return Err(AppError::Validation(crate::i18n::validation::empty_query()));
351 }
352
353 if args.max_cost_usd.is_some() && args.mode == "none" {
354 tracing::warn!(target: "deep_research", "--max-cost-usd has no effect without --mode claude-code/codex");
355 }
356
357 let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
358 let paths = AppPaths::resolve(args.db.as_deref())?;
359 crate::storage::connection::ensure_db_ready(&paths)?;
360
361 let sub_query_texts: Vec<String> = sub_queries.iter().map(|s| s.text.clone()).collect();
363
364 if vec_degraded {
369 tracing::debug!(target: "deep_research", "vector degraded: at least one sub-query fell back to FTS5");
370 }
371
372 let cpu_count = std::thread::available_parallelism()
374 .map(|n| n.get())
375 .unwrap_or(4);
376 let permits = args
377 .max_concurrency
378 .unwrap_or_else(|| cpu_count.min(8))
379 .min(sub_queries.len())
380 .max(1);
381 let semaphore = Arc::new(Semaphore::new(permits));
382 let timeout_dur = std::time::Duration::from_secs(args.timeout);
383
384 let mut join_set: JoinSet<Result<SubQueryResult, (usize, String)>> = JoinSet::new();
385
386 for (idx, sq_text) in sub_query_texts.iter().enumerate() {
387 let sem = Arc::clone(&semaphore);
388 let emb = sub_embeddings[idx].clone();
390 let ns = namespace.clone();
391 let db_path = paths.db.clone();
392 let query_text = sq_text.clone();
393 let k = args.k;
394 let max_hops = args.max_hops;
395 let min_weight = args.min_weight;
396 let rrf_k = args.rrf_k;
397 let graph_decay = args.graph_decay;
398 let graph_min_score = args.graph_min_score;
399 let max_neighbors_per_hop = args.max_neighbors_per_hop;
400
401 join_set.spawn(async move {
402 let _permit = sem
403 .acquire_owned()
404 .await
405 .map_err(|e| (idx, format!("semaphore closed: {e}")))?;
406
407 let result = tokio::time::timeout(timeout_dur, async move {
409 execute_sub_query(
410 idx,
411 &query_text,
412 emb.as_ref().map(|v| v.as_slice()),
413 &ns,
414 &db_path,
415 k,
416 max_hops,
417 min_weight,
418 rrf_k,
419 graph_decay,
420 graph_min_score,
421 max_neighbors_per_hop,
422 )
423 })
424 .await;
425
426 match result {
427 Ok(inner) => inner.map_err(|e| (idx, e)),
428 Err(_) => Err((idx, "timeout".to_string())),
429 }
430 });
431 }
432
433 let mut sub_query_results: Vec<SubQueryResult> = Vec::with_capacity(sub_queries.len());
435 let mut failed_count = 0usize;
436 let mut timed_out_count = 0usize;
437
438 while let Some(join_result) = join_set.join_next().await {
439 match join_result {
440 Ok(Ok(sqr)) => sub_query_results.push(sqr),
441 Ok(Err((_idx, reason))) => {
442 if reason == "timeout" {
443 timed_out_count += 1;
444 } else {
445 failed_count += 1;
446 }
447 tracing::warn!(target: "deep_research", sub_query_id = _idx, reason = %reason, "sub-query failed");
448 }
449 Err(join_err) => {
450 failed_count += 1;
451 if join_err.is_panic() {
452 tracing::error!(target: "deep_research", error = %join_err, "sub-query task panicked");
453 } else {
454 tracing::warn!(target: "deep_research", error = %join_err, "sub-query task cancelled");
455 }
456 }
457 }
458 }
459
460 let mut merged: crate::hash::AHashMap<i64, MergedHit> =
463 crate::hash::AHashMap::with_capacity_and_hasher(
464 sub_query_results.len() * args.k,
465 Default::default(),
466 );
467
468 for sqr in &sub_query_results {
469 for (mem_id, score, source, snippet, body, hop) in &sqr.hits {
470 let entry = merged.entry(*mem_id).or_insert_with(|| {
471 (
472 *score,
473 source.clone(),
474 snippet.clone(),
475 body.clone(),
476 *hop,
477 Vec::new(),
478 )
479 });
480 if *score > entry.0 {
482 entry.0 = *score;
483 entry.1 = source.clone();
484 entry.2 = snippet.clone();
485 entry.3 = body.clone();
486 entry.4 = *hop;
487 }
488 if !entry.5.contains(&sqr.sub_query_id) {
489 entry.5.push(sqr.sub_query_id);
490 }
491 }
492 }
493
494 let conn = open_ro(&paths.db)?;
496 let mut results: Vec<DeepResult> = Vec::with_capacity(merged.len().min(args.max_results));
497
498 let mut ranked: Vec<(i64, MergedHit)> = merged.into_iter().collect();
500 ranked.sort_by(|a, b| {
501 b.1 .0
502 .partial_cmp(&a.1 .0)
503 .unwrap_or(std::cmp::Ordering::Equal)
504 });
505 ranked.truncate(args.max_results);
506
507 for (mem_id, (score, source, snippet, body, hop, sq_ids)) in ranked {
508 let name = match memories::read_full(&conn, mem_id)? {
509 Some(row) => row.name,
510 None => continue,
511 };
512 results.push(DeepResult {
513 name,
514 score,
515 source,
516 sub_query_ids: sq_ids,
517 snippet,
518 body: if args.with_bodies { Some(body) } else { None },
519 hop_distance: hop,
520 });
521 }
522
523 let completed_count = sub_query_results.len();
527 let mut evidence_chains: Vec<EvidenceChain> = Vec::with_capacity(completed_count * 2);
528 let mut seen_chain_keys: HashSet<String> = HashSet::with_capacity(completed_count * 2);
529
530 for sqr in sub_query_results {
531 for chain in sqr.chains {
532 let key = format!("{}->{}", chain.from, chain.to);
534 if seen_chain_keys.insert(key) {
535 evidence_chains.push(chain);
536 }
537 }
538 }
539
540 evidence_chains.retain(|c| c.depth >= 2);
542 evidence_chains.sort_by(|a, b| {
543 b.total_weight
544 .partial_cmp(&a.total_weight)
545 .unwrap_or(std::cmp::Ordering::Equal)
546 });
547
548 let unique_memories = results.len();
549 let evidence_count = evidence_chains.len();
550
551 let graph_context = if !results.is_empty() {
553 let result_names: Vec<&str> = results.iter().map(|r| r.name.as_str()).collect();
554 let mut ctx_entities: Vec<GraphContextEntity> = Vec::with_capacity(results.len());
555 let mut ctx_rels: Vec<GraphContextRel> = Vec::with_capacity(results.len() * 2);
556 let mut seen_entity_ids: crate::hash::AHashSet<i64> =
557 crate::hash::AHashSet::with_capacity_and_hasher(results.len(), Default::default());
558
559 for name in &result_names {
560 if let Ok(Some(eid)) = entities::find_entity_id(&conn, &namespace, name) {
561 if seen_entity_ids.insert(eid) {
562 let etype: String = conn
563 .query_row(
564 "SELECT COALESCE(type,'concept') FROM entities WHERE id = ?1",
565 rusqlite::params![eid],
566 |r| r.get(0),
567 )
568 .unwrap_or_else(|_| "concept".to_string());
569 let degree: u32 = conn
570 .query_row(
571 "SELECT COUNT(*) FROM relationships WHERE source_id = ?1 OR target_id = ?1",
572 rusqlite::params![eid],
573 |r| r.get(0),
574 )
575 .unwrap_or(0);
576 ctx_entities.push(GraphContextEntity {
577 name: name.to_string(),
578 entity_type: etype,
579 degree,
580 });
581 }
582 }
583 }
584
585 let entity_ids: Vec<i64> = seen_entity_ids.iter().copied().collect();
586 if entity_ids.len() >= 2 {
587 let placeholders: String = entity_ids.iter().map(|_| "?").collect::<Vec<_>>().join(",");
588 let sql = format!(
589 "SELECT s.name, t.name, r.relation, r.weight \
590 FROM relationships r \
591 JOIN entities s ON s.id = r.source_id \
592 JOIN entities t ON t.id = r.target_id \
593 WHERE r.source_id IN ({placeholders}) AND r.target_id IN ({placeholders}) \
594 LIMIT 50"
595 );
596 if let Ok(mut stmt) = conn.prepare(&sql) {
597 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> =
598 Vec::with_capacity(entity_ids.len() * 2);
599 for id in &entity_ids {
600 params.push(Box::new(*id));
601 }
602 for id in &entity_ids {
603 params.push(Box::new(*id));
604 }
605 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
606 params.iter().map(|p| p.as_ref()).collect();
607 if let Ok(rows) = stmt.query_map(param_refs.as_slice(), |r| {
608 Ok((
609 r.get::<_, String>(0)?,
610 r.get::<_, String>(1)?,
611 r.get::<_, String>(2)?,
612 r.get::<_, f64>(3)?,
613 ))
614 }) {
615 for row in rows.flatten() {
616 ctx_rels.push(GraphContextRel {
617 from: row.0,
618 to: row.1,
619 relation: row.2,
620 weight: row.3,
621 });
622 }
623 }
624 }
625 }
626
627 if ctx_entities.is_empty() {
628 None
629 } else {
630 Some(GraphContext {
631 entities: ctx_entities,
632 relationships: ctx_rels,
633 })
634 }
635 } else {
636 None
637 };
638
639 tracing::debug!(target: "deep_research",
640 total_results = results.len(),
641 total_chains = evidence_chains.len(),
642 "assembly complete"
643 );
644
645 let response = DeepResearchResponse {
647 query: args.query,
648 sub_queries,
649 results,
650 evidence_chains,
651 graph_context,
652 stats: ResearchStats {
653 sub_queries_total: sub_query_texts.len(),
654 sub_queries_completed: completed_count,
655 sub_queries_failed: failed_count,
656 sub_queries_timed_out: timed_out_count,
657 unique_memories_found: unique_memories,
658 evidence_chains_found: evidence_count,
659 elapsed_ms: start.elapsed().as_millis() as u64,
660 vec_degraded,
661 },
662 };
663
664 if let Some(path) = args.output.as_ref() {
665 crate::atomic_io::write_json_atomic(path, &response)?;
671 if !path.exists() {
672 return Err(AppError::Validation(crate::i18n::validation::deep_research_output_missing(
673 &path.display().to_string(),
674 )));
675 }
676 let meta = std::fs::metadata(path).map_err(AppError::Io)?;
677 if meta.len() == 0 {
678 return Err(AppError::Validation(crate::i18n::validation::deep_research_output_empty(&path.display().to_string())));
679 }
680 let file_bytes = std::fs::read(path).map_err(AppError::Io)?;
681 let digest = blake3::hash(&file_bytes).to_hex().to_string();
682 #[derive(Serialize)]
683 struct WrittenAck {
684 written: String,
685 bytes: u64,
686 blake3: String,
687 sub_queries_total: usize,
688 unique_memories_found: usize,
689 elapsed_ms: u64,
690 }
691 output::emit_json(&WrittenAck {
692 written: path.display().to_string(),
693 bytes: meta.len(),
694 blake3: digest,
695 sub_queries_total: response.stats.sub_queries_total,
696 unique_memories_found: response.stats.unique_memories_found,
697 elapsed_ms: response.stats.elapsed_ms,
698 })?;
699 } else {
700 output::emit_json(&response)?;
701 }
702
703 Ok(())
704}
705
706#[cfg(test)]
710mod tests;