1use crate::error::Result;
14use crate::fts::{is_generic_topic, prepare_search_query, tokenize_query};
15use crate::query::{QueryOptions, RankedHit};
16use crate::storage::Database;
17use crate::types::{ContextBundle, ContextNode, ContextRole, Node, NodeType};
18use serde::{Deserialize, Serialize};
19use std::collections::{HashMap, HashSet};
20use std::path::Path;
21use std::time::Instant;
22
23const MAX_EXCERPT_CHARS: usize = 900;
25const MAX_SYMBOL_EXCERPT_CHARS: usize = 280;
27const MAX_PACKED_SYMBOL_NEIGHBORS: usize = 8;
29
30#[derive(Debug, Clone)]
32pub struct ContextOptions {
33 pub max_tokens: usize,
35 pub hop_depth: usize,
37 pub max_seeds: usize,
39 pub max_nodes: usize,
41 pub hop_decay: f32,
43 pub no_symbols: bool,
45 pub include_types: Vec<crate::types::NodeType>,
47 pub exclude_types: Vec<crate::types::NodeType>,
49 pub hop_to_symbols: bool,
52 pub hop_from_docs_only: bool,
54 pub include_excerpts: bool,
56}
57
58impl Default for ContextOptions {
59 fn default() -> Self {
60 Self {
61 max_tokens: 2048,
62 hop_depth: 1,
63 max_seeds: 8,
64 max_nodes: 24,
65 hop_decay: 0.65,
66 no_symbols: true,
68 include_types: Vec::new(),
69 exclude_types: Vec::new(),
70 hop_to_symbols: true,
71 hop_from_docs_only: true,
72 include_excerpts: true,
73 }
74 }
75}
76
77impl ContextOptions {
78 pub fn agent() -> Self {
80 Self::default()
81 }
82
83 pub fn with_symbols(mut self) -> Self {
85 self.no_symbols = false;
86 self
87 }
88
89 fn seed_query_opts(&self) -> QueryOptions {
90 QueryOptions {
91 limit: self.max_seeds.saturating_mul(2).max(8),
92 no_symbols: self.no_symbols,
93 include_types: self.include_types.clone(),
94 exclude_types: self.exclude_types.clone(),
95 ..QueryOptions::default()
96 }
97 }
98
99 fn allows_pack(&self, ty: &crate::types::NodeType, role: ContextRole) -> bool {
100 use crate::types::NodeType;
101 if !self.include_types.is_empty() && !self.include_types.contains(ty) {
102 if !(self.hop_to_symbols && role == ContextRole::Neighbor && *ty == NodeType::Symbol) {
104 return false;
105 }
106 }
107 if self.exclude_types.contains(ty) {
108 return false;
109 }
110 if *ty == NodeType::Symbol {
111 if role == ContextRole::Seed && self.no_symbols {
112 return false;
113 }
114 if role == ContextRole::Neighbor && self.no_symbols && !self.hop_to_symbols {
115 return false;
116 }
117 }
118 true
119 }
120}
121
122#[derive(Debug, Clone)]
124struct Candidate {
125 node: Node,
126 score: f32,
127 role: ContextRole,
128 hop: u8,
129 excerpt: Option<String>,
130}
131
132pub fn assemble_context(
142 db: &Database,
143 brain_dir: &Path,
144 prompt: &str,
145 opts: &ContextOptions,
146) -> Result<ContextBundle> {
147 let start = Instant::now();
148 let char_budget = opts.max_tokens.saturating_mul(4).max(256);
149 let query_tokens = prepare_search_query(prompt)
150 .map(|p| p.tokens)
151 .unwrap_or_else(|_| tokenize_query(prompt));
152
153 let qopts = opts.seed_query_opts();
154 let mut seeds = db.search_ranked(prompt, &qopts)?;
155 let generic = is_generic_topic(&query_tokens);
156 if seeds.is_empty() || generic {
158 inject_hub_seeds(db, &mut seeds, opts)?;
159 }
160
161 let mut candidates: Vec<Candidate> = Vec::new();
162 let mut seed_ids: HashSet<String> = HashSet::new();
163
164 for (i, hit) in seeds.into_iter().take(opts.max_seeds).enumerate() {
165 if !opts.allows_pack(&hit.node.node_type, ContextRole::Seed) {
166 continue;
167 }
168 if is_template_stub(&hit.node) {
170 continue;
171 }
172 seed_ids.insert(hit.node.id.clone());
173 let mut score = hit.score * (1.0 / (1.0 + i as f32 * 0.05));
175 if is_generated_path(&hit.node) {
177 score *= 0.85;
178 }
179 if hit.node.node_type == NodeType::Adr {
180 score *= 1.08;
181 }
182 let excerpt = if opts.include_excerpts {
183 load_excerpt(db, &hit.node, MAX_EXCERPT_CHARS)
184 } else {
185 None
186 };
187 candidates.push(Candidate {
188 node: hit.node,
189 score,
190 role: ContextRole::Seed,
191 hop: 0,
192 excerpt,
193 });
194 }
195
196 if candidates.is_empty() {
198 let mut hub_hits = Vec::new();
199 inject_hub_seeds(db, &mut hub_hits, opts)?;
200 for hit in hub_hits.into_iter().take(4) {
201 if !opts.allows_pack(&hit.node.node_type, ContextRole::Seed) {
202 continue;
203 }
204 if seed_ids.contains(&hit.node.id) {
205 continue;
206 }
207 seed_ids.insert(hit.node.id.clone());
208 let excerpt = if opts.include_excerpts {
209 load_excerpt(db, &hit.node, MAX_EXCERPT_CHARS)
210 } else {
211 None
212 };
213 candidates.push(Candidate {
214 node: hit.node,
215 score: hit.score,
216 role: ContextRole::Seed,
217 hop: 0,
218 excerpt,
219 });
220 }
221 }
222
223 let mut graph_nodes = 0usize;
224 let mut graph_edges = 0usize;
225 let mut neighbor_ids: Vec<String> = Vec::new();
226
227 #[cfg(feature = "mmap")]
228 if opts.hop_depth > 0 {
229 let mmap_path = brain_dir.join("graph.mmap");
230 if mmap_path.exists() {
231 if let Ok(graph) = crate::mmap::CsrMmapGraph::open(&mmap_path) {
232 graph_nodes = graph.node_count;
233 graph_edges = graph.edge_count;
234
235 let mut best_neighbor: HashMap<String, (f32, u8)> = HashMap::new();
236
237 for cand in candidates.iter().filter(|c| c.role == ContextRole::Seed) {
238 if opts.hop_from_docs_only && cand.node.node_type == NodeType::Symbol {
239 continue;
240 }
241 if let Some(idx) = graph.index_of(&cand.node.id) {
242 for (nidx, edge_w) in
243 graph.k_hop_neighborhood(idx as usize, opts.hop_depth)
244 {
245 if let Some(id) = graph.node_id(nidx as usize) {
246 if seed_ids.contains(id) {
247 continue;
248 }
249 let hop: u8 = if opts.hop_depth <= 1 || edge_w >= 0.85 {
250 1
251 } else {
252 2
253 };
254 let nscore = cand.score * edge_w * opts.hop_decay.powi(hop as i32);
255 best_neighbor
256 .entry(id.to_string())
257 .and_modify(|(s, h)| {
258 if nscore > *s {
259 *s = nscore;
260 *h = hop;
261 }
262 })
263 .or_insert((nscore, hop));
264 }
265 }
266 }
267 }
268
269 let mut neigh_sorted: Vec<_> = best_neighbor.into_iter().collect();
270 neigh_sorted.sort_by(|a, b| {
271 b.1 .0
272 .partial_cmp(&a.1 .0)
273 .unwrap_or(std::cmp::Ordering::Equal)
274 });
275
276 for (id, (score, hop)) in neigh_sorted {
277 neighbor_ids.push(id.clone());
278 if let Some(node) = db.get_node(&id)? {
279 if !opts.allows_pack(&node.node_type, ContextRole::Neighbor) {
280 continue;
281 }
282 let mut nscore = score;
283 if node.node_type == NodeType::Symbol {
284 match symbol_neighbor_quality(&node, &query_tokens) {
285 SymbolQuality::Drop => continue,
286 SymbolQuality::Keep { boost } => nscore *= boost,
287 }
288 }
289 let max_ex = if node.node_type == NodeType::Symbol {
290 MAX_SYMBOL_EXCERPT_CHARS
291 } else {
292 MAX_EXCERPT_CHARS
293 };
294 let excerpt = if opts.include_excerpts {
295 load_excerpt(db, &node, max_ex)
296 } else {
297 None
298 };
299 candidates.push(Candidate {
300 node,
301 score: nscore,
302 role: ContextRole::Neighbor,
303 hop,
304 excerpt,
305 });
306 }
307 }
308 }
309 }
310 }
311
312 candidates.sort_by(|a, b| {
314 let a_bonus = if a.node.node_type == NodeType::Symbol {
316 0.0
317 } else {
318 0.05
319 };
320 let b_bonus = if b.node.node_type == NodeType::Symbol {
321 0.0
322 } else {
323 0.05
324 };
325 (b.score + b_bonus)
326 .partial_cmp(&(a.score + a_bonus))
327 .unwrap_or(std::cmp::Ordering::Equal)
328 });
329
330 let mut packed: Vec<ContextNode> = Vec::new();
331 let mut used_chars = estimate_overhead(prompt);
332 let mut seen = HashSet::new();
333 let mut symbol_neighbors = 0usize;
334 let mut packed_excerpts: Vec<String> = Vec::new();
335
336 for cand in candidates {
337 if packed.len() >= opts.max_nodes {
338 break;
339 }
340 if !seen.insert(cand.node.id.clone()) {
341 continue;
342 }
343 if cand.role == ContextRole::Neighbor
344 && cand.node.node_type == NodeType::Symbol
345 && symbol_neighbors >= MAX_PACKED_SYMBOL_NEIGHBORS
346 {
347 continue;
348 }
349 if is_readme_family(&cand.node.id)
351 && packed.iter().any(|n| is_readme_family(&n.id))
352 {
353 continue;
354 }
355 if let Some(ex) = &cand.excerpt {
357 if packed_excerpts
358 .iter()
359 .any(|prev| excerpt_jaccard(prev, ex) > 0.55)
360 {
361 continue;
362 }
363 }
364 let ctx_node = to_context_node(&cand);
365 let cost = estimate_node_chars(&ctx_node);
366 if used_chars + cost > char_budget && !packed.is_empty() {
367 if cand.role == ContextRole::Neighbor {
369 continue;
370 }
371 if packed.iter().any(|n| n.role == ContextRole::Seed) {
372 continue;
373 }
374 }
375 if used_chars + cost > char_budget && !packed.is_empty() {
376 break;
377 }
378 used_chars += cost;
379 if cand.role == ContextRole::Neighbor && cand.node.node_type == NodeType::Symbol {
380 symbol_neighbors += 1;
381 }
382 if let Some(ex) = &ctx_node.excerpt {
383 packed_excerpts.push(ex.clone());
384 }
385 packed.push(ctx_node);
386 }
387
388 neighbor_ids.retain(|id| !seed_ids.contains(id));
389 neighbor_ids.truncate(48);
390
391 let latency_us = start.elapsed().as_micros() as u64;
392 let tokens_used = used_chars.div_ceil(4);
393
394 Ok(ContextBundle {
395 prompt: prompt.to_string(),
396 max_tokens: opts.max_tokens,
397 tokens_used,
398 nodes: packed,
399 neighbor_ids,
400 latency_us,
401 graph_nodes,
402 graph_edges,
403 })
404}
405
406enum SymbolQuality {
407 Drop,
408 Keep { boost: f32 },
409}
410
411fn symbol_neighbor_quality(node: &Node, query_tokens: &[String]) -> SymbolQuality {
413 let title = node.title.to_lowercase();
414 let id = node.id.to_lowercase();
415 let leaf = id.rsplit('/').next().unwrap_or(&id);
416
417 for t in query_tokens {
419 if title.contains(t.as_str()) || id.contains(t.as_str()) {
420 return SymbolQuality::Keep { boost: 1.8 };
421 }
422 }
423
424 if title.contains("::") || leaf.contains("::") {
426 return SymbolQuality::Keep { boost: 1.15 };
427 }
428
429 let name = title.split("::").last().unwrap_or(&title);
431 if name.len() <= 2 {
432 return SymbolQuality::Drop;
433 }
434 if name.chars().all(|c| c.is_ascii_uppercase() || c == '_') && name.len() <= 12 {
435 return SymbolQuality::Drop;
437 }
438 if !name.contains('_') && !title.contains("::") && name.len() <= 8 {
440 let has_upper = name.chars().any(|c| c.is_ascii_uppercase());
442 let has_lower = name.chars().any(|c| c.is_ascii_lowercase());
443 if has_upper && has_lower {
444 return SymbolQuality::Keep { boost: 1.0 };
445 }
446 return SymbolQuality::Drop;
447 }
448
449 SymbolQuality::Keep { boost: 1.0 }
450}
451
452fn inject_hub_seeds(
454 db: &Database,
455 seeds: &mut Vec<RankedHit>,
456 opts: &ContextOptions,
457) -> Result<()> {
458 let existing: HashSet<String> = seeds.iter().map(|h| h.node.id.clone()).collect();
459 let hubs = [
460 ("readme", 4.5_f32),
461 ("docs/goals/from-readme", 3.8),
462 ("docs/implementation/module-map.generated", 2.2),
463 ("docs/goals/readme", 1.5),
464 ];
465 for (id, score) in hubs {
466 if existing.contains(id) {
467 continue;
468 }
469 if let Some(node) = db.get_node(id)? {
470 if !opts.allows_pack(&node.node_type, ContextRole::Seed) {
471 continue;
472 }
473 if is_template_stub(&node) {
474 continue;
475 }
476 seeds.push(RankedHit {
477 node,
478 score,
479 reasons: vec!["hub-fallback".into()],
480 });
481 }
482 }
483 seeds.sort_by(|a, b| {
485 b.score
486 .partial_cmp(&a.score)
487 .unwrap_or(std::cmp::Ordering::Equal)
488 });
489 Ok(())
490}
491
492fn is_template_stub(node: &Node) -> bool {
493 let id = node.id.to_lowercase();
494 let path = node.file_path.as_deref().unwrap_or("").to_lowercase();
495 let title = node.title.to_lowercase();
496 id.contains("template")
497 || path.ends_with("template.md")
498 || title.contains("template")
499 || title == "adr template"
500}
501
502fn is_generated_path(node: &Node) -> bool {
503 let path = node.file_path.as_deref().unwrap_or("");
504 path.contains("from-readme")
505 || path.contains(".generated.")
506 || path.ends_with("module-map.generated.md")
507}
508
509fn is_readme_family(id: &str) -> bool {
510 id == "readme" || id.contains("from-readme")
511}
512
513fn excerpt_jaccard(a: &str, b: &str) -> f32 {
514 let ta: HashSet<&str> = a
515 .split_whitespace()
516 .filter(|w| w.len() > 2)
517 .collect();
518 let tb: HashSet<&str> = b
519 .split_whitespace()
520 .filter(|w| w.len() > 2)
521 .collect();
522 if ta.is_empty() || tb.is_empty() {
523 return 0.0;
524 }
525 let inter = ta.intersection(&tb).count() as f32;
526 let uni = ta.union(&tb).count() as f32;
527 if uni == 0.0 {
528 0.0
529 } else {
530 inter / uni
531 }
532}
533
534fn load_excerpt(db: &Database, node: &Node, max_chars: usize) -> Option<String> {
535 let raw = db
536 .get_fts_content(&node.id)
537 .ok()
538 .flatten()
539 .filter(|s| !s.trim().is_empty())
540 .or_else(|| node.summary.clone().filter(|s| !s.trim().is_empty()))?;
541 Some(truncate_excerpt(&raw, max_chars))
542}
543
544fn truncate_excerpt(s: &str, max_chars: usize) -> String {
545 let s = s.trim();
546 if s.chars().count() <= max_chars {
547 return s.to_string();
548 }
549 let mut out = String::new();
550 for (i, ch) in s.chars().enumerate() {
551 if i >= max_chars.saturating_sub(1) {
552 break;
553 }
554 out.push(ch);
555 }
556 out.push('…');
557 out
558}
559
560fn to_context_node(c: &Candidate) -> ContextNode {
561 ContextNode {
562 id: c.node.id.clone(),
563 node_type: c.node.node_type.clone(),
564 title: c.node.title.clone(),
565 summary: c.node.summary.clone(),
566 excerpt: c.excerpt.clone(),
567 file_path: c.node.file_path.clone(),
568 score_hint: c.score,
569 role: c.role,
570 hop: c.hop,
571 }
572}
573
574fn estimate_overhead(prompt: &str) -> usize {
575 120 + prompt.len()
576}
577
578fn estimate_node_chars(n: &ContextNode) -> usize {
579 let mut c = n.id.len() + n.title.len() + 48;
580 if let Some(s) = &n.summary {
581 c += s.len();
582 }
583 if let Some(ex) = &n.excerpt {
584 c += ex.len();
585 }
586 if let Some(p) = &n.file_path {
587 c += p.len();
588 }
589 c
590}
591
592pub fn hits_to_nodes(hits: Vec<RankedHit>) -> Vec<Node> {
594 hits.into_iter().map(|h| h.node).collect()
595}
596
597#[derive(Debug, Clone, Serialize, Deserialize)]
602pub struct WorkspaceHit {
603 pub workspace: String,
605 pub hit: RankedHit,
607}
608
609#[cfg(test)]
610mod tests {
611 use super::*;
612 use crate::brain::Brain;
613 use tempfile::tempdir;
614
615 #[test]
616 fn context_includes_graph_neighbors_within_budget() {
617 let dir = tempdir().unwrap();
618 let docs = dir.path().join("docs");
619 std::fs::create_dir_all(&docs).unwrap();
620 std::fs::write(
622 docs.join("alpha.md"),
623 "---\ntags: [alphaunique]\nnode_type: concept\n---\n# AlphaUnique\nLinks [[betaunique]].\nalphaunique body.\n",
624 )
625 .unwrap();
626 std::fs::write(
627 docs.join("beta.md"),
628 "---\ntags: [other]\nnode_type: concept\n---\n# BetaUnique\nNo alphaunique here.\n",
629 )
630 .unwrap();
631
632 let mut brain = Brain::create(dir.path()).unwrap();
633 brain.sync().unwrap();
634
635 let opts = ContextOptions {
636 max_tokens: 2048,
637 hop_depth: 1,
638 max_seeds: 4,
639 max_nodes: 12,
640 hop_decay: 0.7,
641 ..ContextOptions::default()
642 };
643 let ctx = assemble_context(brain.database(), brain.brain_dir(), "alphaunique", &opts)
644 .unwrap();
645 assert!(!ctx.nodes.is_empty());
646 assert!(ctx.tokens_used > 0);
647 assert!(ctx.tokens_used <= opts.max_tokens + 64); let has_beta = ctx.nodes.iter().any(|n| n.id.contains("beta"))
651 || ctx.neighbor_ids.iter().any(|id| id.contains("beta"));
652 assert!(
653 has_beta,
654 "expected beta neighbor via graph; nodes={:?} neigh={:?}",
655 ctx.nodes.iter().map(|n| &n.id).collect::<Vec<_>>(),
656 ctx.neighbor_ids
657 );
658 }
659
660 #[test]
661 fn tight_budget_limits_nodes() {
662 let dir = tempdir().unwrap();
663 let docs = dir.path().join("docs");
664 std::fs::create_dir_all(&docs).unwrap();
665 for i in 0..8 {
666 std::fs::write(
667 docs.join(format!("n{i}.md")),
668 format!("---\ntags: [topic]\n---\n# Note {i}\nshared topic content {i}\n"),
669 )
670 .unwrap();
671 }
672 let mut brain = Brain::create(dir.path()).unwrap();
673 brain.sync().unwrap();
674 let opts = ContextOptions {
675 max_tokens: 40, hop_depth: 0,
677 max_seeds: 8,
678 max_nodes: 20,
679 hop_decay: 0.5,
680 ..ContextOptions::default()
681 };
682 let ctx =
683 assemble_context(brain.database(), brain.brain_dir(), "topic", &opts).unwrap();
684 assert!(!ctx.nodes.is_empty());
685 assert!(ctx.nodes.len() < 8, "budget should clip nodes");
686 }
687
688 #[test]
689 fn natural_language_prompt_seeds_and_packs_excerpt() {
690 let dir = tempdir().unwrap();
691 std::fs::write(
692 dir.path().join("README.md"),
693 "# demotool\n\nLightweight **egui** explorer powered by DuckDB.\n\
694 Inspired by Duckling but avoids Tauri/WebView.\n",
695 )
696 .unwrap();
697 let mut brain = Brain::create(dir.path()).unwrap();
698 brain.sync().unwrap();
699
700 let ctx = assemble_context(
701 brain.database(),
702 brain.brain_dir(),
703 "why egui not tauri",
704 &ContextOptions::default(),
705 )
706 .unwrap();
707 assert!(
708 !ctx.nodes.is_empty(),
709 "expected seeds for NL prompt; packed=0"
710 );
711 let has_excerpt = ctx.nodes.iter().any(|n| {
712 n.excerpt
713 .as_ref()
714 .map(|e| e.to_lowercase().contains("egui") || e.to_lowercase().contains("tauri"))
715 .unwrap_or(false)
716 });
717 assert!(
718 has_excerpt,
719 "expected body excerpt mentioning egui/tauri; nodes={:?}",
720 ctx.nodes
721 .iter()
722 .map(|n| (&n.id, &n.excerpt))
723 .collect::<Vec<_>>()
724 );
725 }
726
727 #[test]
728 fn generic_overview_prompt_falls_back_to_hub() {
729 let dir = tempdir().unwrap();
730 std::fs::write(
731 dir.path().join("README.md"),
732 "# demotool\n\nNative egui + DuckDB CLI. Avoids Tauri.\n",
733 )
734 .unwrap();
735 let mut brain = Brain::create(dir.path()).unwrap();
736 brain.sync().unwrap();
737
738 for prompt in [
739 "summarize architecture",
740 "what is this project about",
741 "give an overview",
742 ] {
743 let ctx = assemble_context(
744 brain.database(),
745 brain.brain_dir(),
746 prompt,
747 &ContextOptions::default(),
748 )
749 .unwrap();
750 assert!(
751 !ctx.nodes.is_empty(),
752 "expected hub fallback for `{prompt}`; packed=0"
753 );
754 let has_hub = ctx.nodes.iter().any(|n| {
755 n.id == "readme"
756 || n.id.contains("from-readme")
757 || n.excerpt
758 .as_ref()
759 .map(|e| e.to_lowercase().contains("egui") || e.to_lowercase().contains("duckdb"))
760 .unwrap_or(false)
761 });
762 assert!(
763 has_hub,
764 "expected README hub content for `{prompt}`; nodes={:?}",
765 ctx.nodes.iter().map(|n| &n.id).collect::<Vec<_>>()
766 );
767 }
768 }
769
770 #[test]
771 fn dedups_near_identical_excerpts() {
772 let dir = tempdir().unwrap();
773 let body = "Lightweight egui explorer with DuckDB CLI. Avoids Tauri completely.\n".repeat(3);
774 std::fs::write(
775 dir.path().join("README.md"),
776 format!("# demotool\n\n{body}"),
777 )
778 .unwrap();
779 std::fs::create_dir_all(dir.path().join("docs/goals")).unwrap();
780 std::fs::write(
781 dir.path().join("docs/goals/from-readme.md"),
782 format!(
783 "---\nnode_type: goal\n---\n# Goals harvested from README\n\n{body}"
784 ),
785 )
786 .unwrap();
787 let mut brain = Brain::create(dir.path()).unwrap();
788 brain.sync().unwrap();
789 let ctx = assemble_context(
790 brain.database(),
791 brain.brain_dir(),
792 "egui duckdb tauri",
793 &ContextOptions {
794 max_tokens: 900,
795 hop_depth: 0,
796 ..ContextOptions::default()
797 },
798 )
799 .unwrap();
800 let ids: Vec<_> = ctx.nodes.iter().map(|n| n.id.as_str()).collect();
802 let both = ids.contains(&"readme") && ids.iter().any(|i| i.contains("from-readme"));
803 assert!(
804 !both || ctx.nodes.len() == 1,
805 "expected dedup of near-identical hubs; packed={ids:?}"
806 );
807 }
808}