use std::collections::{BTreeSet, HashMap};
use std::path::Path;
use crate::analytics::{Analytics, CallGraphNode, ImpactNode};
use crate::db::Database;
use crate::embeddings::{semantic_search, Embedding, EmbeddingProvider, SearchResult};
use crate::error::{CtxError, Result};
use crate::tokens::{count_file_tokens, select_by_token_budget, Encoding, HasTokenCount};
use crate::utils::lexical_tokens;
#[derive(Debug, Clone)]
pub struct SmartConfig {
pub max_tokens: usize,
pub depth: i32,
pub top: usize,
pub encoding: Encoding,
}
impl Default for SmartConfig {
fn default() -> Self {
Self {
max_tokens: 8000,
depth: 2,
top: 10,
encoding: Encoding::default(),
}
}
}
#[derive(Debug, Clone)]
pub enum SelectionReason {
SemanticMatch {
symbol: String,
score: f32,
},
CalledBy {
caller: String,
depth: i32,
},
Calls {
callee: String,
depth: i32,
},
#[allow(dead_code)] SameModule {
symbol: String,
},
#[allow(dead_code)] Explicit,
}
impl SelectionReason {
pub fn description(&self) -> String {
match self {
SelectionReason::SemanticMatch { symbol, score } => {
format!("SemanticMatch: \"{}\" (score: {:.2})", symbol, score)
}
SelectionReason::CalledBy { caller, depth } => {
format!("CalledBy: {} (depth {})", caller, depth)
}
SelectionReason::Calls { callee, depth } => {
format!("Calls: {} (depth {})", callee, depth)
}
SelectionReason::SameModule { symbol } => {
format!("SameModule: shares context with {}", symbol)
}
SelectionReason::Explicit => "Explicit: user requested".to_string(),
}
}
fn weight(&self) -> f32 {
match self {
SelectionReason::SemanticMatch { score, .. } => *score,
SelectionReason::CalledBy { depth, .. } => 0.8 / (*depth as f32).max(1.0),
SelectionReason::Calls { depth, .. } => 0.7 / (*depth as f32).max(1.0),
SelectionReason::SameModule { .. } => 0.5,
SelectionReason::Explicit => 1.0,
}
}
}
#[derive(Debug, Clone)]
pub struct FileSelection {
pub path: String,
pub relevance_score: f32,
pub reasons: Vec<SelectionReason>,
pub token_count: usize,
pub lexical_score: f32,
}
impl HasTokenCount for FileSelection {
fn token_count(&self) -> usize {
self.token_count
}
}
impl FileSelection {
fn new(path: String, reason: SelectionReason) -> Self {
let score = reason.weight();
Self {
path,
relevance_score: score,
reasons: vec![reason],
token_count: 0,
lexical_score: 0.0,
}
}
fn add_reason(&mut self, reason: SelectionReason) {
let new_weight = reason.weight();
if new_weight > self.relevance_score {
self.relevance_score = new_weight;
}
self.reasons.push(reason);
}
fn best_semantic_score(&self) -> Option<f32> {
self.reasons
.iter()
.filter_map(|r| match r {
SelectionReason::SemanticMatch { score, .. } => Some(*score),
_ => None,
})
.reduce(f32::max)
}
fn compute_lexical_score(&self, task_tokens: &BTreeSet<String>) -> f32 {
if task_tokens.is_empty() {
return 0.0;
}
let path_tokens = lexical_tokens(&self.path);
task_tokens.intersection(&path_tokens).count() as f32
}
}
#[derive(Debug, Clone)]
pub struct SmartContext {
#[allow(dead_code)] pub task: String,
pub selected_files: Vec<FileSelection>,
pub total_tokens: usize,
pub truncated: bool,
pub omitted_count: usize,
}
pub fn smart_context(
db: &Database,
analytics: &Analytics,
provider: &dyn EmbeddingProvider,
task: &str,
config: SmartConfig,
) -> Result<SmartContext> {
let task_embedding = provider.embed(task)?;
smart_context_with_embedding(db, analytics, task, &task_embedding, config)
}
pub fn smart_context_with_embedding(
db: &Database,
analytics: &Analytics,
task: &str,
task_embedding: &Embedding,
config: SmartConfig,
) -> Result<SmartContext> {
let matches = semantic_search(db, task_embedding, config.top)?;
if matches.is_empty() {
return Err(CtxError::NoMatches);
}
let mut files: HashMap<String, FileSelection> = HashMap::new();
for result in &matches {
add_file(
&mut files,
&result.file_path,
SelectionReason::SemanticMatch {
symbol: result.name.clone(),
score: result.score,
},
);
expand_symbol(&mut files, analytics, result, config.depth)?;
}
let mut selections: Vec<FileSelection> = files.into_values().collect();
for selection in &mut selections {
selection.token_count = count_file_token_safe(&selection.path, config.encoding);
}
let task_tokens = lexical_tokens(task);
for selection in &mut selections {
selection.lexical_score = selection.compute_lexical_score(&task_tokens);
}
rank_files(&mut selections);
let (selected, total_tokens, omitted) =
select_with_guaranteed_top(selections, config.max_tokens);
Ok(SmartContext {
task: task.to_string(),
selected_files: selected,
total_tokens,
truncated: omitted > 0,
omitted_count: omitted,
})
}
fn select_with_guaranteed_top(
ranked: Vec<FileSelection>,
max_tokens: usize,
) -> (Vec<FileSelection>, usize, usize) {
let mut iter = ranked.into_iter();
let Some(top) = iter.next() else {
return (Vec::new(), 0, 0);
};
let top_tokens = top.token_count;
let rest: Vec<FileSelection> = iter.collect();
let remaining_budget = max_tokens.saturating_sub(top_tokens);
let (mut selected_rest, rest_tokens, omitted) = select_by_token_budget(rest, remaining_budget);
let mut selected = Vec::with_capacity(selected_rest.len() + 1);
selected.push(top);
selected.append(&mut selected_rest);
(selected, top_tokens + rest_tokens, omitted)
}
fn add_file(files: &mut HashMap<String, FileSelection>, path: &str, reason: SelectionReason) {
if let Some(existing) = files.get_mut(path) {
existing.add_reason(reason);
} else {
files.insert(
path.to_string(),
FileSelection::new(path.to_string(), reason),
);
}
}
fn expand_symbol(
files: &mut HashMap<String, FileSelection>,
analytics: &Analytics,
result: &SearchResult,
depth: i32,
) -> Result<()> {
if let Ok(callers) = analytics.impact_analysis(&result.symbol_id, depth) {
for caller in callers {
add_file_from_impact(files, &caller, &result.name);
}
}
if let Ok(callees) = analytics.call_graph(&result.symbol_id, depth) {
for callee in callees {
add_file_from_call_graph(files, &callee, &result.name);
}
}
Ok(())
}
fn add_file_from_impact(
files: &mut HashMap<String, FileSelection>,
node: &ImpactNode,
callee_name: &str,
) {
add_file(
files,
&node.file_path,
SelectionReason::CalledBy {
caller: node.name.clone(),
depth: node.distance,
},
);
if let Some(existing) = files.get(&node.file_path) {
if existing
.reasons
.iter()
.any(|r| matches!(r, SelectionReason::SemanticMatch { .. }))
{
} else {
let _ = callee_name; }
}
}
fn add_file_from_call_graph(
files: &mut HashMap<String, FileSelection>,
node: &CallGraphNode,
caller_name: &str,
) {
add_file(
files,
&node.file_path,
SelectionReason::Calls {
callee: node.name.clone(),
depth: node.depth,
},
);
let _ = caller_name; }
fn rank_key(f: &FileSelection) -> (bool, f32, f32) {
let sem = f.best_semantic_score();
let is_low_relevance = sem.is_none() && f.lexical_score == 0.0;
let tier_score = sem.unwrap_or(f.relevance_score);
(is_low_relevance, f.lexical_score, tier_score)
}
fn rank_files(files: &mut [FileSelection]) {
files.sort_by(|a, b| {
let a_key = rank_key(a);
let b_key = rank_key(b);
a_key
.0
.cmp(&b_key.0) .then_with(|| {
b_key
.1
.partial_cmp(&a_key.1) .unwrap_or(std::cmp::Ordering::Equal)
})
.then_with(|| {
b_key
.2
.partial_cmp(&a_key.2) .unwrap_or(std::cmp::Ordering::Equal)
})
.then_with(|| a.path.cmp(&b.path))
});
}
fn count_file_token_safe(path: &str, encoding: Encoding) -> usize {
count_file_tokens(Path::new(path), encoding)
.map(|tc| tc.count)
.unwrap_or(0)
}
pub fn format_explain(result: &SmartContext) -> String {
let mut output = String::new();
output.push_str(&format!(
"Selected {} files ({} tokens):\n\n",
result.selected_files.len(),
result.total_tokens
));
for (i, file) in result.selected_files.iter().enumerate() {
output.push_str(&format!(
"{}. {} ({} tokens)\n",
i + 1,
file.path,
file.token_count
));
for reason in &file.reasons {
output.push_str(&format!(" - {}\n", reason.description()));
}
output.push('\n');
}
if result.truncated {
output.push_str(&format!(
"({} files omitted due to token limit)\n",
result.omitted_count
));
}
output
}
pub fn format_dry_run(result: &SmartContext) -> String {
let mut output = String::new();
output.push_str(&format!(
"Would select {} files ({} tokens):\n",
result.selected_files.len(),
result.total_tokens
));
for file in &result.selected_files {
let primary_reason = file
.reasons
.first()
.map(|r| match r {
SelectionReason::SemanticMatch { .. } => "SemanticMatch",
SelectionReason::CalledBy { .. } => "CalledBy",
SelectionReason::Calls { .. } => "Calls",
SelectionReason::SameModule { .. } => "SameModule",
SelectionReason::Explicit => "Explicit",
})
.unwrap_or("Unknown");
output.push_str(&format!(
" {} ({} tokens) - {}\n",
file.path, file.token_count, primary_reason
));
}
if result.truncated {
output.push_str(&format!(
"\n({} files would be omitted due to token limit)\n",
result.omitted_count
));
}
output
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_selection_reason_weight() {
let semantic = SelectionReason::SemanticMatch {
symbol: "test".to_string(),
score: 0.9,
};
assert!((semantic.weight() - 0.9).abs() < 0.001);
let called_by = SelectionReason::CalledBy {
caller: "main".to_string(),
depth: 1,
};
assert!((called_by.weight() - 0.8).abs() < 0.001);
let called_by_depth2 = SelectionReason::CalledBy {
caller: "main".to_string(),
depth: 2,
};
assert!((called_by_depth2.weight() - 0.4).abs() < 0.001);
let calls = SelectionReason::Calls {
callee: "helper".to_string(),
depth: 1,
};
assert!((calls.weight() - 0.7).abs() < 0.001);
let same_module = SelectionReason::SameModule {
symbol: "related".to_string(),
};
assert!((same_module.weight() - 0.5).abs() < 0.001);
let explicit = SelectionReason::Explicit;
assert!((explicit.weight() - 1.0).abs() < 0.001);
}
#[test]
fn test_file_selection_add_reason() {
let mut selection = FileSelection::new(
"src/main.rs".to_string(),
SelectionReason::Calls {
callee: "helper".to_string(),
depth: 1,
},
);
assert!((selection.relevance_score - 0.7).abs() < 0.001);
assert_eq!(selection.reasons.len(), 1);
selection.add_reason(SelectionReason::SemanticMatch {
symbol: "main".to_string(),
score: 0.95,
});
assert!((selection.relevance_score - 0.95).abs() < 0.001);
assert_eq!(selection.reasons.len(), 2);
}
#[test]
fn test_select_by_token_budget() {
let files = vec![
FileSelection {
path: "a.rs".to_string(),
relevance_score: 1.0,
reasons: vec![SelectionReason::Explicit],
token_count: 100,
lexical_score: 0.0,
},
FileSelection {
path: "b.rs".to_string(),
relevance_score: 0.8,
reasons: vec![SelectionReason::Explicit],
token_count: 200,
lexical_score: 0.0,
},
FileSelection {
path: "c.rs".to_string(),
relevance_score: 0.5,
reasons: vec![SelectionReason::Explicit],
token_count: 150,
lexical_score: 0.0,
},
];
let (selected, total, omitted) = select_by_token_budget(files.clone(), 500);
assert_eq!(selected.len(), 3);
assert_eq!(total, 450);
assert_eq!(omitted, 0);
let (selected, total, omitted) = select_by_token_budget(files.clone(), 300);
assert_eq!(selected.len(), 2);
assert_eq!(total, 300);
assert_eq!(omitted, 1);
let (selected, total, omitted) = select_by_token_budget(files, 150);
assert_eq!(selected.len(), 1);
assert_eq!(total, 100);
assert_eq!(omitted, 2);
}
#[test]
fn test_format_dry_run() {
let result = SmartContext {
task: "add caching".to_string(),
selected_files: vec![
FileSelection {
path: "src/main.rs".to_string(),
relevance_score: 0.9,
reasons: vec![SelectionReason::SemanticMatch {
symbol: "main".to_string(),
score: 0.9,
}],
token_count: 500,
lexical_score: 0.0,
},
FileSelection {
path: "src/lib.rs".to_string(),
relevance_score: 0.7,
reasons: vec![SelectionReason::Calls {
callee: "helper".to_string(),
depth: 1,
}],
token_count: 300,
lexical_score: 0.0,
},
],
total_tokens: 800,
truncated: false,
omitted_count: 0,
};
let output = format_dry_run(&result);
assert!(output.contains("Would select 2 files"));
assert!(output.contains("src/main.rs"));
assert!(output.contains("SemanticMatch"));
}
#[test]
fn test_smart_config_default() {
let config = SmartConfig::default();
assert_eq!(config.max_tokens, 8000);
assert_eq!(config.depth, 2);
assert_eq!(config.top, 10);
}
#[test]
fn test_add_file_merges_reasons() {
let mut files: HashMap<String, FileSelection> = HashMap::new();
add_file(
&mut files,
"src/main.rs",
SelectionReason::SemanticMatch {
symbol: "main".to_string(),
score: 0.9,
},
);
assert_eq!(files.len(), 1);
assert!((files.get("src/main.rs").unwrap().relevance_score - 0.9).abs() < 0.001);
add_file(
&mut files,
"src/main.rs",
SelectionReason::CalledBy {
caller: "run".to_string(),
depth: 1,
},
);
assert_eq!(files.len(), 1); assert_eq!(files.get("src/main.rs").unwrap().reasons.len(), 2);
assert!((files.get("src/main.rs").unwrap().relevance_score - 0.9).abs() < 0.001);
}
#[test]
fn test_rank_files_sorts_by_relevance() {
let mut files = vec![
FileSelection {
path: "low.rs".to_string(),
relevance_score: 0.3,
reasons: vec![SelectionReason::Explicit],
token_count: 100,
lexical_score: 0.0,
},
FileSelection {
path: "high.rs".to_string(),
relevance_score: 0.9,
reasons: vec![SelectionReason::Explicit],
token_count: 100,
lexical_score: 0.0,
},
FileSelection {
path: "mid.rs".to_string(),
relevance_score: 0.5,
reasons: vec![SelectionReason::Explicit],
token_count: 100,
lexical_score: 0.0,
},
];
rank_files(&mut files);
assert_eq!(files[0].path, "high.rs");
assert_eq!(files[1].path, "mid.rs");
assert_eq!(files[2].path, "low.rs");
}
#[test]
fn test_semantic_match_ranks_above_graph_only() {
let mut files = vec![
FileSelection {
path: "graph_hub.rs".to_string(),
relevance_score: 0.8, reasons: vec![SelectionReason::CalledBy {
caller: "run".to_string(),
depth: 1,
}],
token_count: 100,
lexical_score: 0.0,
},
FileSelection {
path: "on_topic.rs".to_string(),
relevance_score: 0.5,
reasons: vec![SelectionReason::SemanticMatch {
symbol: "run_sql".to_string(),
score: 0.5,
}],
token_count: 100,
lexical_score: 0.0,
},
];
rank_files(&mut files);
assert_eq!(
files[0].path, "on_topic.rs",
"semantic match must outrank a higher-weighted graph-only file"
);
assert_eq!(files[1].path, "graph_hub.rs");
}
#[test]
fn test_semantic_tier_orders_by_score_then_path() {
let mk = |path: &str, score: f32| FileSelection {
path: path.to_string(),
relevance_score: score,
reasons: vec![SelectionReason::SemanticMatch {
symbol: "s".to_string(),
score,
}],
token_count: 100,
lexical_score: 0.0,
};
let mut files = vec![mk("b.rs", 0.6), mk("c.rs", 0.4), mk("a.rs", 0.6)];
rank_files(&mut files);
assert_eq!(
files.iter().map(|f| f.path.as_str()).collect::<Vec<_>>(),
vec!["a.rs", "b.rs", "c.rs"]
);
}
#[test]
fn test_rank_files_is_deterministic() {
let make_set = || {
vec![
FileSelection {
path: "z_graph.rs".to_string(),
relevance_score: 0.8,
reasons: vec![SelectionReason::CalledBy {
caller: "run".to_string(),
depth: 1,
}],
token_count: 100,
lexical_score: 0.0,
},
FileSelection {
path: "a_graph.rs".to_string(),
relevance_score: 0.8, reasons: vec![SelectionReason::CalledBy {
caller: "run".to_string(),
depth: 1,
}],
token_count: 100,
lexical_score: 0.0,
},
FileSelection {
path: "seed.rs".to_string(),
relevance_score: 0.5,
reasons: vec![SelectionReason::SemanticMatch {
symbol: "seed".to_string(),
score: 0.5,
}],
token_count: 100,
lexical_score: 0.0,
},
]
};
let mut a = make_set();
rank_files(&mut a);
let order_a: Vec<String> = a.iter().map(|f| f.path.clone()).collect();
let mut b = make_set();
b.reverse();
rank_files(&mut b);
let order_b: Vec<String> = b.iter().map(|f| f.path.clone()).collect();
assert_eq!(order_a, order_b);
assert_eq!(order_a, vec!["seed.rs", "a_graph.rs", "z_graph.rs"]);
}
#[test]
fn test_lexical_promotes_graph_only_file() {
let mut files = vec![
FileSelection {
path: "off_topic_semantic.rs".to_string(),
relevance_score: 0.6,
reasons: vec![SelectionReason::SemanticMatch {
symbol: "unrelated".to_string(),
score: 0.6,
}],
token_count: 100,
lexical_score: 0.0,
},
FileSelection {
path: "embeddings/openai.rs".to_string(),
relevance_score: 0.7, reasons: vec![SelectionReason::Calls {
callee: "embed".to_string(),
depth: 1,
}],
token_count: 100,
lexical_score: 2.0, },
];
rank_files(&mut files);
assert_eq!(
files[0].path, "embeddings/openai.rs",
"a lexical hit must outrank a semantic match with no task-token overlap"
);
}
#[test]
fn test_lexical_orders_within_tier() {
let mk = |path: &str, score: f32, lex: f32| FileSelection {
path: path.to_string(),
relevance_score: score,
reasons: vec![SelectionReason::SemanticMatch {
symbol: "s".to_string(),
score,
}],
token_count: 100,
lexical_score: lex,
};
let mut files = vec![
mk("high_score_no_lexical.rs", 0.9, 0.0),
mk("two_hits.rs", 0.4, 2.0),
mk("one_hit.rs", 0.4, 1.0),
];
rank_files(&mut files);
assert_eq!(
files.iter().map(|f| f.path.as_str()).collect::<Vec<_>>(),
vec!["two_hits.rs", "one_hit.rs", "high_score_no_lexical.rs"]
);
}
#[test]
fn test_compute_lexical_score() {
let sel = FileSelection {
path: "src/embeddings/openai.rs".to_string(),
relevance_score: 0.5,
reasons: vec![SelectionReason::SemanticMatch {
symbol: "ctx_embed".to_string(),
score: 0.5,
}],
token_count: 100,
lexical_score: 0.0,
};
let task = lexical_tokens("generate embeddings with openai");
assert!((sel.compute_lexical_score(&task) - 2.0).abs() < 0.001);
let ctx_task = lexical_tokens("ctx tooling");
assert!((sel.compute_lexical_score(&ctx_task)).abs() < 0.001);
let other = lexical_tokens("parse solidity contracts");
assert!((sel.compute_lexical_score(&other)).abs() < 0.001);
}
#[test]
fn test_budget_includes_oversized_top_file() {
let ranked = vec![
FileSelection {
path: "huge_top.rs".to_string(),
relevance_score: 0.9,
reasons: vec![SelectionReason::SemanticMatch {
symbol: "s".to_string(),
score: 0.9,
}],
token_count: 9000, lexical_score: 1.0,
},
FileSelection {
path: "small_next.rs".to_string(),
relevance_score: 0.5,
reasons: vec![SelectionReason::SemanticMatch {
symbol: "s".to_string(),
score: 0.5,
}],
token_count: 500,
lexical_score: 0.0,
},
];
let (selected, total, omitted) = select_with_guaranteed_top(ranked, 8000);
assert_eq!(selected.len(), 1, "only the oversized top file is included");
assert_eq!(selected[0].path, "huge_top.rs");
assert_eq!(total, 9000);
assert_eq!(omitted, 1, "the smaller file is omitted for lack of budget");
}
#[test]
fn test_budget_first_fits_remainder() {
let ranked = vec![
FileSelection {
path: "top.rs".to_string(),
relevance_score: 0.9,
reasons: vec![SelectionReason::Explicit],
token_count: 3000,
lexical_score: 0.0,
},
FileSelection {
path: "mid.rs".to_string(),
relevance_score: 0.8,
reasons: vec![SelectionReason::Explicit],
token_count: 4000,
lexical_score: 0.0,
},
FileSelection {
path: "tail.rs".to_string(),
relevance_score: 0.7,
reasons: vec![SelectionReason::Explicit],
token_count: 2000, lexical_score: 0.0,
},
];
let (selected, total, omitted) = select_with_guaranteed_top(ranked, 8000);
assert_eq!(
selected.iter().map(|f| f.path.as_str()).collect::<Vec<_>>(),
vec!["top.rs", "mid.rs"]
);
assert_eq!(total, 7000);
assert_eq!(omitted, 1);
}
#[test]
fn test_selection_reason_description() {
let semantic = SelectionReason::SemanticMatch {
symbol: "test".to_string(),
score: 0.9,
};
assert!(semantic.description().contains("SemanticMatch"));
assert!(semantic.description().contains("test"));
let called_by = SelectionReason::CalledBy {
caller: "main".to_string(),
depth: 2,
};
assert!(called_by.description().contains("CalledBy"));
assert!(called_by.description().contains("main"));
assert!(called_by.description().contains("2"));
}
}