mentedb_context/
budget.rs

1//! Token budget management for context assembly.
2
3/// Estimates token count for a text string.
4/// Uses word count * 1.3 as a rough approximation for English text.
5pub fn estimate_tokens(text: &str) -> usize {
6    let word_count = text.split_whitespace().count();
7    ((word_count as f64) * 1.3).ceil() as usize
8}
9
10/// Tracks token usage against a maximum budget.
11#[derive(Debug, Clone)]
12pub struct TokenBudget {
13    pub max_tokens: usize,
14    pub used_tokens: usize,
15}
16
17impl TokenBudget {
18    pub fn new(max_tokens: usize) -> Self {
19        Self {
20            max_tokens,
21            used_tokens: 0,
22        }
23    }
24
25    /// Remaining tokens available.
26    pub fn remaining(&self) -> usize {
27        self.max_tokens.saturating_sub(self.used_tokens)
28    }
29
30    /// Check if the given text fits within remaining budget.
31    pub fn can_fit(&self, text: &str) -> bool {
32        estimate_tokens(text) <= self.remaining()
33    }
34
35    /// Consume tokens for the given text. Returns the number of tokens consumed.
36    pub fn consume(&mut self, text: &str) -> usize {
37        let tokens = estimate_tokens(text);
38        let actual = tokens.min(self.remaining());
39        self.used_tokens += actual;
40        actual
41    }
42
43    /// Reset the budget to zero usage.
44    pub fn reset(&mut self) {
45        self.used_tokens = 0;
46    }
47}
48
49/// Configuration for zone budget percentages.
50#[derive(Debug, Clone)]
51pub struct ZoneBudgetConfig {
52    /// Fraction of total budget for system zone.
53    pub system_pct: f32,
54    /// Fraction of total budget for critical zone.
55    pub critical_pct: f32,
56    /// Fraction of total budget for primary zone.
57    pub primary_pct: f32,
58    /// Fraction of total budget for supporting zone.
59    pub supporting_pct: f32,
60    /// Fraction of total budget for reference zone.
61    pub reference_pct: f32,
62}
63
64impl Default for ZoneBudgetConfig {
65    fn default() -> Self {
66        Self {
67            system_pct: 0.10,
68            critical_pct: 0.25,
69            primary_pct: 0.35,
70            supporting_pct: 0.20,
71            reference_pct: 0.10,
72        }
73    }
74}
75
76/// Divides a total token budget across context zones.
77#[derive(Debug, Clone)]
78pub struct BudgetAllocation {
79    /// System zone
80    pub system: usize,
81    /// Critical zone
82    pub critical: usize,
83    /// Primary zone
84    pub primary: usize,
85    /// Supporting zone
86    pub supporting: usize,
87    /// Reference zone
88    pub reference: usize,
89}
90
91impl BudgetAllocation {
92    /// Allocate a total budget across zones with the given percentages.
93    pub fn from_total_with_config(total: usize, config: &ZoneBudgetConfig) -> Self {
94        Self {
95            system: (total as f32 * config.system_pct) as usize,
96            critical: (total as f32 * config.critical_pct) as usize,
97            primary: (total as f32 * config.primary_pct) as usize,
98            supporting: (total as f32 * config.supporting_pct) as usize,
99            reference: (total as f32 * config.reference_pct) as usize,
100        }
101    }
102
103    /// Allocate a total budget across zones with default percentages.
104    pub fn from_total(total: usize) -> Self {
105        Self::from_total_with_config(total, &ZoneBudgetConfig::default())
106    }
107
108    /// Total allocated tokens (may differ slightly from input due to rounding).
109    pub fn total(&self) -> usize {
110        self.system + self.critical + self.primary + self.supporting + self.reference
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn test_estimate_tokens() {
120        assert_eq!(estimate_tokens(""), 0);
121        assert_eq!(estimate_tokens("hello world"), 3); // 2 * 1.3 = 2.6 -> 3
122        assert_eq!(estimate_tokens("one"), 2); // 1 * 1.3 = 1.3 -> 2
123    }
124
125    #[test]
126    fn test_token_budget_lifecycle() {
127        let mut budget = TokenBudget::new(100);
128        assert_eq!(budget.remaining(), 100);
129        assert!(budget.can_fit("hello world"));
130
131        let used = budget.consume("hello world");
132        assert_eq!(used, 3);
133        assert_eq!(budget.remaining(), 97);
134
135        budget.reset();
136        assert_eq!(budget.remaining(), 100);
137    }
138
139    #[test]
140    fn test_budget_overflow_protection() {
141        let mut budget = TokenBudget::new(2);
142        // "a b c d e" = 5 words * 1.3 = 7 tokens, won't fit
143        assert!(!budget.can_fit("a b c d e"));
144        // consume should only take what's available
145        let used = budget.consume("a b c d e");
146        assert_eq!(used, 2);
147        assert_eq!(budget.remaining(), 0);
148    }
149
150    #[test]
151    fn test_budget_allocation() {
152        let alloc = BudgetAllocation::from_total(1000);
153        assert_eq!(alloc.system, 100);
154        assert_eq!(alloc.critical, 250);
155        assert_eq!(alloc.primary, 350);
156        assert_eq!(alloc.supporting, 200);
157        assert_eq!(alloc.reference, 100);
158        assert_eq!(alloc.total(), 1000);
159    }
160}
mentedb_context/budget.rs

mentedb_context/
budget.rs
