cortexkit_model_batch/
estimate.rs

1//! The conservative cost estimate a runner checks against a batch's
2//! `max_cost_usd` before each provider call, and the money arithmetic it
3//! needs (`RULES.md` §5).
4//!
5//! Text tokens are bounded by UTF-8 byte length: a tokenizer that works on
6//! bytes never makes more tokens than there are bytes, so the bound never
7//! underestimates. A module that has the provider's tokenizer may pass
8//! exact counts instead ([`Tokens::Exact`]). The fixed costs a role cannot
9//! see (the module's prompt template, the tokens an image costs) are
10//! required inputs, never defaulted to zero. Money is integer nanodollars,
11//! the catalog's own unit ([`RateNanosPerMtok`]), with checked arithmetic:
12//! an overflow refuses, it never wraps.
13
14use serde::Serialize;
15use serde_json::Number;
16
17/// A price in integer nanodollars per million tokens. It is the same type as
18/// the model catalog's rate (`cortexkit-model-catalog`'s `RateNanosPerMtok`,
19/// an `i64`), so a catalog price passes in as it is, without this crate
20/// depending on the catalog.
21pub type RateNanosPerMtok = i64;
22
23use crate::identity::{canonical_json, IdentityError};
24
25/// Nanodollars in one dollar.
26pub const NANOS_PER_USD: u64 = 1_000_000_000;
27
28/// Tokens in the unit a rate is priced per (one million).
29pub const TOKENS_PER_MTOK: u64 = 1_000_000;
30
31/// A token count for one part of a request.
32#[derive(Clone, Copy, Debug, PartialEq, Eq)]
33pub enum Tokens {
34    /// Counted with the provider's own tokenizer.
35    Exact(u64),
36    /// An upper bound, from the part's byte length.
37    Bound(u64),
38}
39
40impl Tokens {
41    /// The byte bound of `text`.
42    pub fn bound_of_text(text: &str) -> Self {
43        Self::Bound(text.len() as u64)
44    }
45
46    /// The byte bound of `value`'s RFC 8785 (JCS) canonical JSON: the bound
47    /// of a schema (a role's fields or questions) as the provider receives
48    /// it.
49    pub fn bound_of_json<T: Serialize + ?Sized>(value: &T) -> Result<Self, IdentityError> {
50        Ok(Self::Bound(canonical_json(value)?.len() as u64))
51    }
52
53    pub fn count(self) -> u64 {
54        match self {
55            Self::Exact(count) | Self::Bound(count) => count,
56        }
57    }
58
59    pub fn is_exact(self) -> bool {
60        matches!(self, Self::Exact(_))
61    }
62}
63
64/// What one item's provider call is estimated from.
65#[derive(Clone, Copy, Debug, PartialEq, Eq)]
66#[non_exhaustive]
67pub struct ItemCostInput {
68    /// The item's state.
69    pub state: Tokens,
70    /// The schema sent with every item (fields or questions).
71    pub schema: Tokens,
72    /// How many images the item carries.
73    pub images: u64,
74    /// The module's bound on the tokens one image costs. Required; a
75    /// module that does not know it states a bound, never zero.
76    pub image_tokens_each: u64,
77    /// The module's bound on its per-request overhead (prompt template,
78    /// system text). Required; never defaulted to zero.
79    pub overhead_tokens: u64,
80    /// The provider's per-pass token budget, when it splits a schema into
81    /// passes; `None` for a provider that answers in one pass.
82    pub pass_tokens: Option<u64>,
83}
84
85impl ItemCostInput {
86    pub fn new(
87        state: Tokens,
88        schema: Tokens,
89        image_tokens_each: u64,
90        overhead_tokens: u64,
91    ) -> Self {
92        Self {
93            state,
94            schema,
95            images: 0,
96            image_tokens_each,
97            overhead_tokens,
98            pass_tokens: None,
99        }
100    }
101
102    pub fn with_images(mut self, images: u64) -> Self {
103        self.images = images;
104        self
105    }
106
107    pub fn with_pass_tokens(mut self, pass_tokens: u64) -> Self {
108        self.pass_tokens = Some(pass_tokens);
109        self
110    }
111}
112
113/// An item's estimate.
114#[derive(Clone, Copy, Debug, PartialEq, Eq)]
115pub enum Estimate {
116    /// An upper bound on the item's input cost.
117    Bounded {
118        /// Passes the provider is assumed to make, at least 1.
119        passes: u64,
120        /// Input tokens billed over every pass.
121        input_tokens: u64,
122        /// Their cost, rounded up to a whole nanodollar.
123        cost_nanos: u64,
124    },
125    /// The state's byte bound reaches a pass's token budget, so the number
126    /// of passes, and with it the cost, has no upper bound. Against a
127    /// ceiling the item is `cost_exceeded`; with no ceiling it is sent.
128    Unbounded,
129}
130
131/// Why an item could not be estimated.
132#[derive(Clone, Copy, Debug, PartialEq, Eq)]
133pub enum EstimateError {
134    /// The state's exact count reaches the pass budget: the item cannot fit
135    /// in a pass, and is refused `invalid_item` before any call.
136    DoesNotFit { state_tokens: u64, pass_tokens: u64 },
137    /// The arithmetic overflowed. The item is refused; a runner never wraps.
138    Overflow,
139    /// The catalog rate is negative, which no catalog publishes.
140    NegativeRate,
141}
142
143impl std::fmt::Display for EstimateError {
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        match self {
146            Self::DoesNotFit {
147                state_tokens,
148                pass_tokens,
149            } => write!(
150                f,
151                "the state is {state_tokens} tokens; a pass holds {pass_tokens}"
152            ),
153            Self::Overflow => f.write_str("the cost estimate overflows"),
154            Self::NegativeRate => f.write_str("the catalog rate is negative"),
155        }
156    }
157}
158
159impl std::error::Error for EstimateError {}
160
161/// How many passes the provider makes for a schema of `schema` tokens with a
162/// state of `state` tokens, at `pass_tokens` per pass:
163/// `ceil(schema / (pass_tokens - state))`, at least 1. `None` when the state
164/// leaves no room in a pass (`state >= pass_tokens`).
165pub fn passes(schema: u64, state: u64, pass_tokens: u64) -> Option<u64> {
166    let room = pass_tokens.checked_sub(state).filter(|room| *room > 0)?;
167    Some(schema.div_ceil(room).max(1))
168}
169
170/// The cost of `tokens` input tokens at `rate` nanodollars per million,
171/// rounded up to a whole nanodollar.
172pub fn cost_nanos(tokens: u64, rate: RateNanosPerMtok) -> Result<u64, EstimateError> {
173    let rate = u64::try_from(rate).map_err(|_| EstimateError::NegativeRate)?;
174    let product = u128::from(tokens) * u128::from(rate);
175    let nanos = product.div_ceil(u128::from(TOKENS_PER_MTOK));
176    u64::try_from(nanos).map_err(|_| EstimateError::Overflow)
177}
178
179/// Estimate one item's input cost: `(state + schema + image tokens +
180/// overhead) x passes x input rate`. Each pass is assumed to bill the whole
181/// request again, which is the conservative reading while the provider's
182/// own billing of later passes is not yet measured.
183pub fn estimate_item(
184    input: &ItemCostInput,
185    input_rate: RateNanosPerMtok,
186) -> Result<Estimate, EstimateError> {
187    let state = input.state.count();
188    let schema = input.schema.count();
189    let passes = match input.pass_tokens {
190        None => 1,
191        Some(pass_tokens) => match passes(schema, state, pass_tokens) {
192            Some(passes) => passes,
193            None if input.state.is_exact() => {
194                return Err(EstimateError::DoesNotFit {
195                    state_tokens: state,
196                    pass_tokens,
197                })
198            }
199            None => return Ok(Estimate::Unbounded),
200        },
201    };
202    let image_tokens = input
203        .images
204        .checked_mul(input.image_tokens_each)
205        .ok_or(EstimateError::Overflow)?;
206    let per_pass = [schema, image_tokens, input.overhead_tokens]
207        .into_iter()
208        .try_fold(state, u64::checked_add)
209        .ok_or(EstimateError::Overflow)?;
210    let input_tokens = per_pass
211        .checked_mul(passes)
212        .ok_or(EstimateError::Overflow)?;
213    Ok(Estimate::Bounded {
214        passes,
215        input_tokens,
216        cost_nanos: cost_nanos(input_tokens, input_rate)?,
217    })
218}
219
220/// Whether an item estimated at `estimate` would cross `ceiling_nanos` after
221/// `spent_nanos` already recorded for the batch. An unbounded estimate, or a
222/// sum that overflows, crosses any ceiling.
223pub fn crosses_ceiling(spent_nanos: u64, estimate: &Estimate, ceiling_nanos: u64) -> bool {
224    match estimate {
225        Estimate::Unbounded => true,
226        Estimate::Bounded { cost_nanos, .. } => spent_nanos
227            .checked_add(*cost_nanos)
228            .is_none_or(|total| total > ceiling_nanos),
229    }
230}
231
232/// Which way [`usd_to_nanos`] rounds a value finer than a nanodollar.
233#[derive(Clone, Copy, Debug, PartialEq, Eq)]
234pub enum Rounding {
235    /// Toward zero: for a ceiling, so it is never raised.
236    Down,
237    /// Away from zero: for a cost, so it is never lowered.
238    Up,
239}
240
241/// A JSON number of USD as whole nanodollars, by decimal string scaling of
242/// the number's text (never through floating-point arithmetic). `None` for
243/// a negative number or one too large for `u64` nanodollars.
244pub fn usd_to_nanos(usd: &Number, rounding: Rounding) -> Option<u64> {
245    decimal_to_nanos(&usd.to_string(), rounding)
246}
247
248fn decimal_to_nanos(text: &str, rounding: Rounding) -> Option<u64> {
249    if let Some(magnitude) = text.strip_prefix('-') {
250        // Only a negative zero is a valid amount.
251        let zero = magnitude
252            .split(['e', 'E'])
253            .next()
254            .is_some_and(|mantissa| mantissa.bytes().all(|b| matches!(b, b'0' | b'.')));
255        return zero.then_some(0);
256    }
257    let (mantissa, exponent) = match text.find(['e', 'E']) {
258        Some(at) => (&text[..at], text[at + 1..].parse::<i64>().ok()?),
259        None => (text, 0),
260    };
261    let (int, frac) = mantissa.split_once('.').unwrap_or((mantissa, ""));
262    let digits: String = format!("{int}{frac}");
263    if digits.is_empty() || !digits.bytes().all(|b| b.is_ascii_digit()) {
264        return None;
265    }
266    let digits = digits.trim_start_matches('0');
267    if digits.is_empty() {
268        return Some(0);
269    }
270    // value = digits x 10^(exponent - frac.len()); nanos = value x 10^9.
271    let scale = exponent.checked_sub(frac.len() as i64)?.checked_add(9)?;
272    let whole: u128 = if digits.len() > 38 {
273        return None;
274    } else {
275        digits.parse().ok()?
276    };
277    let nanos = if scale >= 0 {
278        let factor = 10u128.checked_pow(u32::try_from(scale).ok()?)?;
279        whole.checked_mul(factor)?
280    } else {
281        let shift = u32::try_from(-scale).ok()?;
282        match 10u128.checked_pow(shift) {
283            Some(divisor) => {
284                let (quotient, remainder) = (whole / divisor, whole % divisor);
285                if rounding == Rounding::Up && remainder > 0 {
286                    quotient + 1
287                } else {
288                    quotient
289                }
290            }
291            // Smaller than any nanodollar fraction `u128` can hold.
292            None => u128::from(rounding == Rounding::Up),
293        }
294    };
295    u64::try_from(nanos).ok()
296}
297
298/// Whole nanodollars as an exact decimal number of USD, for a message:
299/// `0.00110396`, never an exponent form.
300pub fn nanos_to_usd_text(nanos: u64) -> String {
301    let whole = nanos / NANOS_PER_USD;
302    let frac = nanos % NANOS_PER_USD;
303    if frac == 0 {
304        format!("{whole}")
305    } else {
306        let frac = format!("{frac:09}");
307        format!("{whole}.{}", frac.trim_end_matches('0'))
308    }
309}
310
311/// Whole nanodollars as a JSON number of USD, for a refusal's
312/// `estimate_usd`: the exact decimal text, read as a JSON number.
313pub fn nanos_to_usd(nanos: u64) -> Number {
314    serde_json::from_str(&nanos_to_usd_text(nanos)).expect("a plain decimal is a JSON number")
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320
321    fn num(text: &str) -> Number {
322        serde_json::from_str(text).unwrap()
323    }
324
325    #[test]
326    fn passes_are_at_least_one_and_never_divide_by_zero() {
327        assert_eq!(passes(0, 10, 100), Some(1));
328        assert_eq!(passes(90, 10, 100), Some(1));
329        assert_eq!(passes(91, 10, 100), Some(2));
330        assert_eq!(passes(180, 10, 100), Some(2));
331        assert_eq!(passes(181, 10, 100), Some(3));
332        assert_eq!(passes(1, 100, 100), None);
333        assert_eq!(passes(1, 200, 100), None);
334    }
335
336    #[test]
337    fn a_byte_bound_past_the_pass_is_unbounded_and_an_exact_count_does_not_fit() {
338        let bound =
339            ItemCostInput::new(Tokens::Bound(9000), Tokens::Bound(10), 0, 0).with_pass_tokens(8192);
340        assert_eq!(estimate_item(&bound, 1000), Ok(Estimate::Unbounded));
341        let exact =
342            ItemCostInput::new(Tokens::Exact(9000), Tokens::Bound(10), 0, 0).with_pass_tokens(8192);
343        assert_eq!(
344            estimate_item(&exact, 1000),
345            Err(EstimateError::DoesNotFit {
346                state_tokens: 9000,
347                pass_tokens: 8192
348            })
349        );
350        assert!(crosses_ceiling(0, &Estimate::Unbounded, u64::MAX));
351    }
352
353    #[test]
354    fn every_pass_bills_the_whole_request() {
355        let input = ItemCostInput::new(Tokens::Bound(1000), Tokens::Bound(9000), 300, 50)
356            .with_images(2)
357            .with_pass_tokens(8192);
358        // passes = ceil(9000 / 7192) = 2; per pass 1000 + 9000 + 600 + 50.
359        assert_eq!(
360            estimate_item(&input, 110_000_000),
361            Ok(Estimate::Bounded {
362                passes: 2,
363                input_tokens: 21_300,
364                cost_nanos: 2_343_000,
365            })
366        );
367    }
368
369    #[test]
370    fn overflow_refuses() {
371        let input = ItemCostInput::new(Tokens::Exact(u64::MAX), Tokens::Exact(1), 0, 0);
372        assert_eq!(estimate_item(&input, 1), Err(EstimateError::Overflow));
373        assert_eq!(cost_nanos(1, -1), Err(EstimateError::NegativeRate));
374        assert!(crosses_ceiling(
375            u64::MAX,
376            &Estimate::Bounded {
377                passes: 1,
378                input_tokens: 1,
379                cost_nanos: 1
380            },
381            u64::MAX
382        ));
383    }
384
385    #[test]
386    fn cost_rounds_up_to_a_nanodollar() {
387        assert_eq!(cost_nanos(1, 1), Ok(1));
388        assert_eq!(cost_nanos(1_000_000, 3), Ok(3));
389        assert_eq!(cost_nanos(0, 5), Ok(0));
390    }
391
392    #[test]
393    fn usd_scales_by_decimal_text() {
394        assert_eq!(
395            usd_to_nanos(&num("0.25"), Rounding::Down),
396            Some(250_000_000)
397        );
398        assert_eq!(usd_to_nanos(&num("1e-7"), Rounding::Down), Some(100));
399        assert_eq!(usd_to_nanos(&num("3"), Rounding::Down), Some(3_000_000_000));
400        assert_eq!(usd_to_nanos(&num("1e-10"), Rounding::Down), Some(0));
401        assert_eq!(usd_to_nanos(&num("1e-10"), Rounding::Up), Some(1));
402        assert_eq!(usd_to_nanos(&num("0"), Rounding::Up), Some(0));
403        assert_eq!(usd_to_nanos(&num("-1"), Rounding::Down), None);
404        assert_eq!(usd_to_nanos(&num("1e300"), Rounding::Down), None);
405        assert_eq!(nanos_to_usd(420_000_000).to_string(), "0.42");
406        assert_eq!(nanos_to_usd(3_000_000_000).to_string(), "3");
407        assert_eq!(nanos_to_usd(2_343_000).to_string(), "0.002343");
408        assert_eq!(nanos_to_usd_text(36_960), "0.00003696");
409    }
410}
cortexkit_model_batch/estimate.rs

cortexkit_model_batch/
estimate.rs
