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

atomr_agents_context/
lib.rs
