atomr_agents_context/
lib.rs

1//! Context assembler.
2//!
3//! Strategies produce `ContextFragment`s with priorities and token
4//! estimates. The assembler orders by priority and packs greedily
5//! into the remaining `TokenBudget` — high-priority fragments
6//! survive; low-priority ones are dropped first under pressure.
7
8use atomr_agents_core::{Result, TokenBudget};
9
10/// One fragment contributed by a strategy.
11#[derive(Debug, Clone)]
12pub struct ContextFragment {
13    pub source: &'static str,
14    pub priority: u8,
15    pub estimated_tokens: u32,
16    pub text: String,
17}
18
19/// Output of `ContextAssembler::assemble`.
20#[derive(Debug, Clone, Default)]
21pub struct RenderedContext {
22    pub fragments: Vec<ContextFragment>,
23    pub total_tokens: u32,
24}
25
26impl RenderedContext {
27    pub fn join(&self, sep: &str) -> String {
28        self.fragments
29            .iter()
30            .map(|f| f.text.as_str())
31            .collect::<Vec<_>>()
32            .join(sep)
33    }
34}
35
36pub struct ContextAssembler;
37
38impl ContextAssembler {
39    /// Pack the fragments into the remaining budget, highest priority
40    /// first. Returns the fragments that fit, in the order they were
41    /// originally given (priority is used for *eviction*, not
42    /// reordering).
43    pub fn assemble(
44        mut fragments: Vec<ContextFragment>,
45        budget: &mut TokenBudget,
46    ) -> Result<RenderedContext> {
47        // Stable indexed sort: keep original positions for tie-break,
48        // but evict lowest priority first when over budget.
49        let mut indexed: Vec<(usize, ContextFragment)> = fragments.drain(..).enumerate().collect();
50
51        let total: u64 = indexed.iter().map(|(_, f)| f.estimated_tokens as u64).sum();
52        if total <= budget.remaining as u64 {
53            // Everything fits; restore original order.
54            let mut out: Vec<ContextFragment> = indexed.into_iter().map(|(_, f)| f).collect();
55            let total_tokens = out.iter().map(|f| f.estimated_tokens).sum();
56            budget.consume(total_tokens)?;
57            return Ok(RenderedContext {
58                fragments: std::mem::take(&mut out),
59                total_tokens,
60            });
61        }
62
63        // Otherwise, evict lowest-priority fragments until we fit.
64        // Priority: higher number = more important.
65        indexed.sort_by(|a, b| b.1.priority.cmp(&a.1.priority).then_with(|| a.0.cmp(&b.0)));
66        let mut kept: Vec<(usize, ContextFragment)> = Vec::new();
67        let mut acc: u64 = 0;
68        for entry in indexed {
69            let cost = entry.1.estimated_tokens as u64;
70            if acc + cost <= budget.remaining as u64 {
71                acc += cost;
72                kept.push(entry);
73            }
74        }
75        kept.sort_by_key(|(i, _)| *i);
76        let out: Vec<ContextFragment> = kept.into_iter().map(|(_, f)| f).collect();
77        let total_tokens = out.iter().map(|f| f.estimated_tokens).sum();
78        budget.consume(total_tokens)?;
79        Ok(RenderedContext {
80            fragments: out,
81            total_tokens,
82        })
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    fn frag(source: &'static str, prio: u8, tokens: u32) -> ContextFragment {
91        ContextFragment {
92            source,
93            priority: prio,
94            estimated_tokens: tokens,
95            text: source.to_string(),
96        }
97    }
98
99    #[test]
100    fn assemble_fits_under_budget() {
101        let mut b = TokenBudget::new(100);
102        let r =
103            ContextAssembler::assemble(vec![frag("a", 5, 30), frag("b", 5, 30), frag("c", 5, 30)], &mut b)
104                .unwrap();
105        assert_eq!(r.fragments.len(), 3);
106        assert_eq!(r.total_tokens, 90);
107        assert_eq!(b.remaining, 10);
108    }
109
110    #[test]
111    fn assemble_evicts_lowest_priority_first() {
112        let mut b = TokenBudget::new(60);
113        let r = ContextAssembler::assemble(
114            vec![frag("low", 1, 30), frag("hi", 9, 30), frag("med", 5, 30)],
115            &mut b,
116        )
117        .unwrap();
118        let kept: Vec<&str> = r.fragments.iter().map(|f| f.source).collect();
119        assert_eq!(kept, vec!["hi", "med"]);
120    }
121}
atomr_agents_context/lib.rs

atomr_agents_context/
lib.rs
