aivcs_core/memory/
context.rs

1//! Token-budgeted context assembly from memory entries.
2
3use super::index::MemoryEntry;
4
5/// Budget constraints for context window assembly.
6#[derive(Debug, Clone)]
7pub struct ContextBudget {
8    pub max_tokens: usize,
9    pub reserved_tokens: usize,
10}
11
12impl ContextBudget {
13    pub fn new(max_tokens: usize, reserved_tokens: usize) -> Result<Self, String> {
14        if reserved_tokens >= max_tokens {
15            return Err(format!(
16                "reserved_tokens ({reserved_tokens}) must be less than max_tokens ({max_tokens})"
17            ));
18        }
19        Ok(Self {
20            max_tokens,
21            reserved_tokens,
22        })
23    }
24
25    /// Available tokens after reserving space.
26    pub fn available(&self) -> usize {
27        self.max_tokens - self.reserved_tokens
28    }
29}
30
31impl Default for ContextBudget {
32    fn default() -> Self {
33        Self {
34            max_tokens: 128_000,
35            reserved_tokens: 4_000,
36        }
37    }
38}
39
40/// A single item in the assembled context window.
41#[derive(Debug, Clone)]
42pub struct ContextItem {
43    pub entry_id: String,
44    pub text: String,
45    pub tokens: usize,
46}
47
48/// The assembled context window.
49#[derive(Debug, Clone)]
50pub struct ContextWindow {
51    pub items: Vec<ContextItem>,
52    pub total_tokens: usize,
53    pub dropped_count: usize,
54    pub budget: ContextBudget,
55}
56
57/// Assemble a context window from candidate entries, respecting the token budget.
58///
59/// Candidates are sorted by relevance (descending), then greedily packed
60/// until the budget is exhausted. Entries that don't fit are dropped.
61pub fn assemble_context(candidates: &[MemoryEntry], budget: &ContextBudget) -> ContextWindow {
62    let mut sorted: Vec<&MemoryEntry> = candidates.iter().collect();
63    sorted.sort_by(|a, b| {
64        let a_score = if a.relevance.is_finite() {
65            a.relevance
66        } else {
67            f64::NEG_INFINITY
68        };
69        let b_score = if b.relevance.is_finite() {
70            b.relevance
71        } else {
72            f64::NEG_INFINITY
73        };
74        b_score
75            .total_cmp(&a_score)
76            .then_with(|| b.created_at.cmp(&a.created_at))
77            .then_with(|| a.id.cmp(&b.id))
78    });
79
80    let available = budget.available();
81    let mut items = Vec::new();
82    let mut total_tokens = 0;
83    let mut dropped_count = 0;
84
85    for entry in sorted {
86        if total_tokens + entry.token_estimate <= available {
87            items.push(ContextItem {
88                entry_id: entry.id.clone(),
89                text: entry.summary.clone(),
90                tokens: entry.token_estimate,
91            });
92            total_tokens += entry.token_estimate;
93        } else {
94            dropped_count += 1;
95        }
96    }
97
98    ContextWindow {
99        items,
100        total_tokens,
101        dropped_count,
102        budget: budget.clone(),
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use crate::memory::index::MemoryEntryKind;
110    use chrono::Utc;
111
112    fn make(id: &str, tokens: usize, relevance: f64) -> MemoryEntry {
113        MemoryEntry {
114            id: id.into(),
115            kind: MemoryEntryKind::RunTrace,
116            summary: format!("summary {id}"),
117            content_digest: format!("d_{id}"),
118            created_at: Utc::now(),
119            tags: Vec::new(),
120            token_estimate: tokens,
121            relevance,
122        }
123    }
124
125    #[test]
126    fn test_empty_candidates() {
127        let budget = ContextBudget::new(1000, 100).unwrap();
128        let w = assemble_context(&[], &budget);
129        assert!(w.items.is_empty());
130        assert_eq!(w.total_tokens, 0);
131    }
132
133    #[test]
134    fn test_all_fit() {
135        let entries = vec![make("a", 100, 0.5), make("b", 200, 0.8)];
136        let budget = ContextBudget::new(1000, 100).unwrap();
137        let w = assemble_context(&entries, &budget);
138        assert_eq!(w.items.len(), 2);
139        assert_eq!(w.total_tokens, 300);
140        // Higher relevance first
141        assert_eq!(w.items[0].entry_id, "b");
142    }
143
144    #[test]
145    fn test_budget_drops() {
146        let entries = vec![make("a", 500, 0.9), make("b", 500, 0.5)];
147        let budget = ContextBudget::new(700, 100).unwrap();
148        let w = assemble_context(&entries, &budget);
149        assert_eq!(w.items.len(), 1);
150        assert_eq!(w.dropped_count, 1);
151    }
152
153    #[test]
154    fn test_budget_validation() {
155        assert!(ContextBudget::new(100, 200).is_err());
156        assert!(ContextBudget::new(100, 100).is_err());
157        assert!(ContextBudget::new(100, 99).is_ok());
158    }
159
160    #[test]
161    fn test_nan_relevance_does_not_panic_or_win_sorting() {
162        let entries = vec![make("good", 100, 0.8), make("nan", 100, f64::NAN)];
163        let budget = ContextBudget::new(150, 0).unwrap();
164        let w = assemble_context(&entries, &budget);
165        assert_eq!(w.items.len(), 1);
166        assert_eq!(w.items[0].entry_id, "good");
167    }
168}
aivcs_core/memory/context.rs

aivcs_core/memory/
context.rs
