cairn_runtime/
budget.rs

1use serde::{Deserialize, Serialize};
2
3/// Hard limits enforced by the runtime before every model call.
4#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
5pub struct Budget {
6    /// Maximum number of model calls.
7    pub max_steps: Option<u32>,
8    /// Maximum spend in USD. Every model the agent may call must have
9    /// pricing, otherwise [`crate::AgentBuilder::build`] refuses the agent.
10    pub max_usd: Option<f64>,
11    /// Maximum input + output tokens.
12    pub max_tokens: Option<u64>,
13    /// Stop instead of calling the model when the remaining budget allows
14    /// fewer output tokens than this. Defaults to 1.
15    pub min_output_tokens: Option<u32>,
16}
17
18impl Budget {
19    pub fn max_steps(mut self, n: u32) -> Self {
20        self.max_steps = Some(n);
21        self
22    }
23
24    pub fn max_usd(mut self, usd: f64) -> Self {
25        self.max_usd = Some(usd);
26        self
27    }
28
29    pub fn max_tokens(mut self, n: u64) -> Self {
30        self.max_tokens = Some(n);
31        self
32    }
33
34    pub fn min_output_tokens(mut self, n: u32) -> Self {
35        self.min_output_tokens = Some(n);
36        self
37    }
38
39    /// Largest output-token limit for the next call that keeps the run
40    /// within budget, given an estimate of the call's input tokens.
41    ///
42    /// `call_cost(input, output)` is the model's worst-case USD for a call
43    /// (see [`crate::Model::max_call_cost`]). Returns `Err(reason)` when the
44    /// remaining budget cannot pay for `min_output_tokens`.
45    pub fn output_cap(
46        &self,
47        configured: u32,
48        used_tokens: u64,
49        spent_usd: f64,
50        est_input_tokens: u64,
51        call_cost: impl Fn(u64, u64) -> Option<f64>,
52    ) -> Result<u32, String> {
53        let mut cap = u64::from(configured);
54        if let Some(max) = self.max_tokens {
55            cap = cap.min(max.saturating_sub(used_tokens).saturating_sub(est_input_tokens));
56        }
57        if let (Some(max), Some(fixed)) = (self.max_usd, call_cost(est_input_tokens, 0)) {
58            // Cost is linear in output tokens: fixed + per_token * n.
59            let per_token = call_cost(0, 1_000_000).unwrap_or(0.0) / 1_000_000.0;
60            let left = max - spent_usd - fixed;
61            if per_token > 0.0 {
62                cap = cap.min((left / per_token).floor().max(0.0) as u64);
63            } else if left < 0.0 {
64                cap = 0;
65            }
66        }
67        let min = u64::from(self.min_output_tokens.unwrap_or(1));
68        if cap < min {
69            return Err(format!(
70                "remaining budget allows only {cap} output tokens for the next call (minimum {min}, \
71                 estimated input {est_input_tokens} tokens)"
72            ));
73        }
74        Ok(cap as u32)
75    }
76
77    /// Returns why the budget is exhausted, if it is.
78    pub fn exceeded(&self, steps: u32, tokens: u64, usd: f64) -> Option<String> {
79        if let Some(max) = self.max_steps {
80            if steps >= max {
81                return Some(format!("step limit reached ({steps}/{max})"));
82            }
83        }
84        if let Some(max) = self.max_tokens {
85            if tokens >= max {
86                return Some(format!("token limit reached ({tokens}/{max})"));
87            }
88        }
89        if let Some(max) = self.max_usd {
90            if usd >= max {
91                return Some(format!("spend limit reached (${usd:.4}/${max:.4})"));
92            }
93        }
94        None
95    }
96}
cairn_runtime/budget.rs

cairn_runtime/
budget.rs
