agentic_reality/query/
budget.rs

1//! Token budget enforcement.
2
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct TokenBudget {
7    pub max_tokens: u64,
8    pub used_tokens: u64,
9}
10
11impl TokenBudget {
12    pub fn new(max_tokens: u64) -> Self {
13        Self {
14            max_tokens,
15            used_tokens: 0,
16        }
17    }
18
19    pub fn unlimited() -> Self {
20        Self {
21            max_tokens: u64::MAX,
22            used_tokens: 0,
23        }
24    }
25
26    pub fn remaining(&self) -> u64 {
27        self.max_tokens.saturating_sub(self.used_tokens)
28    }
29
30    pub fn is_exhausted(&self) -> bool {
31        self.used_tokens >= self.max_tokens
32    }
33
34    pub fn can_afford(&self, cost: u64) -> bool {
35        self.remaining() >= cost
36    }
37
38    pub fn spend(&mut self, tokens: u64) -> bool {
39        if self.can_afford(tokens) {
40            self.used_tokens += tokens;
41            true
42        } else {
43            false
44        }
45    }
46
47    pub fn force_spend(&mut self, tokens: u64) {
48        self.used_tokens += tokens;
49    }
50
51    pub fn utilization(&self) -> f64 {
52        if self.max_tokens == 0 || self.max_tokens == u64::MAX {
53            return 0.0;
54        }
55        self.used_tokens as f64 / self.max_tokens as f64
56    }
57}
58
59impl Default for TokenBudget {
60    fn default() -> Self {
61        Self::unlimited()
62    }
63}
agentic_reality/query/budget.rs

agentic_reality/query/
budget.rs
