Skip to main content

lean_ctx/core/knowledge_router/
context_bundle.rs

1use super::planner::ContextCandidate;
2#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
3pub enum BundleStrategy {
4    #[default]
5    Minimal,
6    Enriched,
7    Governed,
8}
9#[derive(Debug, Clone, Default, PartialEq, Eq)]
10pub struct ContextBundle {
11    pub bundle_id: String,
12    pub task_id: String,
13    pub candidates: Vec<String>,
14    pub total_tokens: u64,
15    pub coverage_milli: u16,
16    pub strategy: BundleStrategy,
17}
18pub fn create_bundle(
19    task_id: &str,
20    candidates: &[ContextCandidate],
21    strategy: BundleStrategy,
22) -> ContextBundle {
23    let selected = if strategy == BundleStrategy::Minimal {
24        &candidates[..candidates.len().min(3)]
25    } else {
26        candidates
27    };
28    let total_tokens = selected
29        .iter()
30        .map(|candidate| candidate.estimated_tokens)
31        .sum();
32    let coverage_milli = selected
33        .iter()
34        .map(|candidate| u64::from(candidate.relevance_milli))
35        .sum::<u64>()
36        .checked_div(selected.len() as u64)
37        .unwrap_or(0) as u16;
38    ContextBundle {
39        bundle_id: format!("bundle-{task_id}-{}", selected.len()),
40        task_id: task_id.into(),
41        candidates: selected
42            .iter()
43            .map(|candidate| candidate.candidate_id.clone())
44            .collect(),
45        total_tokens,
46        coverage_milli,
47        strategy,
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54    fn candidates() -> Vec<ContextCandidate> {
55        (0..4)
56            .map(|id| ContextCandidate {
57                candidate_id: format!("c{id}"),
58                source_id: "local_files".into(),
59                kind: "file".into(),
60                relevance_milli: 800 + id,
61                estimated_tokens: 100,
62                reference: None,
63            })
64            .collect()
65    }
66    #[test]
67    fn minimal_has_at_most_three_candidates() {
68        assert_eq!(
69            create_bundle("task", &candidates(), BundleStrategy::Minimal)
70                .candidates
71                .len(),
72            3
73        );
74    }
75    #[test]
76    fn coverage_is_average_relevance() {
77        assert_eq!(
78            create_bundle("task", &candidates()[..2], BundleStrategy::Enriched).coverage_milli,
79            800
80        );
81    }
82}