Skip to main content

codetether_rlm/context_trace/
trace_budget.rs

1use super::ContextTrace;
2
3impl ContextTrace {
4    /// Increment the iteration counter.
5    pub fn next_iteration(&mut self) {
6        self.iteration += 1;
7    }
8
9    /// Get the current iteration number.
10    pub fn iteration(&self) -> usize {
11        self.iteration
12    }
13
14    /// Get the total tokens used.
15    pub fn total_tokens(&self) -> usize {
16        self.total_tokens
17    }
18
19    /// Get the remaining token budget.
20    pub fn remaining_tokens(&self) -> usize {
21        self.max_tokens.saturating_sub(self.total_tokens)
22    }
23
24    /// Get the percentage of budget used.
25    pub fn budget_used_percent(&self) -> f32 {
26        if self.max_tokens == 0 {
27            return 0.0;
28        }
29        let total = self.total_tokens as u64;
30        let max = self.max_tokens as u64;
31        let tenths = (total.saturating_mul(1000) + (max / 2)).saturating_div(max);
32        (tenths as f32) / 10.0
33    }
34
35    /// Check if the budget is exceeded.
36    pub fn is_over_budget(&self) -> bool {
37        self.total_tokens > self.max_tokens
38    }
39}