Skip to main content

jev_repl/
cost.rs

1//! What a session costs: the tokens a request carries, the tokens its answers bring back, and the
2//! money that is at rates you supply.
3//!
4//! Nothing here calls the API. Token counts are an estimate — roughly four characters to a token,
5//! JSON punctuation in pairs — so treat them as an order of magnitude, not an invoice. The answer
6//! side is not guesswork about length: a System One answer has the shape the question asks for, so
7//! the estimate prices the real shape, which is why a choice over eight labels costs more to
8//! answer than a noul.
9//!
10//! Rates are yours to set, in dollars per million tokens, because the price of a model is not
11//! something an SDK should hardcode: `:cost 0.20/1.00` in the REPL, or `JEV_PRICE=0.20/1.00` in
12//! the environment.
13
14use serde_json::{Value, json};
15use typesafe::Usage;
16
17use crate::mock;
18use crate::session::{Session, turns_to_json};
19
20/// Environment variable holding `<input>/<output>` dollars per million tokens.
21pub const PRICE_ENV: &str = "JEV_PRICE";
22
23/// A question shape this crate does not model: assume an answer the size of a noul's.
24const ASSUMED_ANSWER: &str = r#"{"name":{"type":"noul","noul":0.123}}"#;
25
26/// Dollars per million tokens, one rate for what goes up and one for what comes back.
27#[derive(Debug, Clone, Copy, PartialEq)]
28pub struct Rates {
29    /// Dollars per million input tokens.
30    pub input: f64,
31    /// Dollars per million output tokens.
32    pub output: f64,
33}
34
35/// What one question adds to a request and to the answers.
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct QuestionEstimate {
38    pub name: String,
39    /// `noul`, `choice`, `score`, or whatever a hand-built question object calls itself.
40    pub kind: String,
41    /// Tokens the question itself contributes to the request.
42    pub input_tokens: usize,
43    /// Tokens its answer is expected to contribute to the response.
44    pub output_tokens: usize,
45    /// True when the answer shape could not be derived and a noul-sized answer was assumed.
46    pub assumed: bool,
47}
48
49/// The token side of a call: where they go, and how many there are.
50#[derive(Debug, Clone, PartialEq, Eq)]
51pub struct Estimate {
52    pub model: String,
53    /// Tokens in the state being judged.
54    pub state_tokens: usize,
55    /// Tokens in the keys, braces and model name wrapped around the request.
56    pub envelope_tokens: usize,
57    /// Tokens in the braces wrapped around the answers.
58    pub answer_envelope_tokens: usize,
59    pub questions: Vec<QuestionEstimate>,
60    /// State, questions and envelope together.
61    pub input_tokens: usize,
62    /// Every expected answer, plus its envelope.
63    pub output_tokens: usize,
64}
65
66/// Dollars, split the way the rates are.
67#[derive(Debug, Clone, Copy, PartialEq)]
68pub struct Cost {
69    pub input: f64,
70    pub output: f64,
71    pub total: f64,
72}
73
74/// Tokens in a piece of text, estimated: a run of letters is a token per four characters, digits
75/// run denser, a run of punctuation pairs up the way `":"` and `"},` do in a real vocabulary, and
76/// a CJK character is a token on its own. Whitespace rides along with the token beside it. No
77/// tokenizer is shipped or downloaded to do this; the API reports the real counts in `usage` once
78/// a call has been made.
79pub fn estimate_tokens(text: &str) -> usize {
80    /// The run of like characters being counted; flushing it turns the run into tokens.
81    #[derive(Default)]
82    struct Runs {
83        letters: usize,
84        digits: usize,
85        marks: usize,
86    }
87    impl Runs {
88        fn flush(&mut self) -> usize {
89            let tokens =
90                self.letters.div_ceil(4) + self.digits.div_ceil(3) + self.marks.div_ceil(2);
91            *self = Runs::default();
92            tokens
93        }
94    }
95
96    let mut tokens = 0;
97    let mut runs = Runs::default();
98    for ch in text.chars() {
99        if ch.is_whitespace() {
100            // Whitespace rides along with the token beside it, so it only ends a run.
101            tokens += runs.flush();
102        } else if ch as u32 >= 0x2e80 {
103            // Anything above the CJK block is a character per token or worse; Latin is far denser.
104            tokens += runs.flush() + 1;
105        } else if ch.is_ascii_digit() {
106            if runs.letters > 0 || runs.marks > 0 {
107                tokens += runs.flush();
108            }
109            runs.digits += 1;
110        } else if ch.is_alphabetic() {
111            if runs.digits > 0 || runs.marks > 0 {
112                tokens += runs.flush();
113            }
114            runs.letters += 1;
115        } else {
116            if runs.letters > 0 || runs.digits > 0 {
117                tokens += runs.flush();
118            }
119            runs.marks += 1;
120        }
121    }
122    tokens + runs.flush()
123}
124
125/// Tokens in a JSON value, as it goes on the wire.
126pub fn estimate_json_tokens(value: &Value) -> usize {
127    estimate_tokens(&value.to_string())
128}
129
130/// Estimate one call: what the session sends, and what its answers come back as.
131pub fn estimate(session: &Session, model: &str) -> Estimate {
132    let state_tokens = estimate_json_tokens(&session.state);
133    let envelope_tokens =
134        estimate_json_tokens(&json!({"state": "", "model": model, "questions": {}}));
135    let answer_envelope_tokens = estimate_json_tokens(&json!({"model": model, "answers": {}}));
136
137    let questions: Vec<QuestionEstimate> = session
138        .questions
139        .iter()
140        .map(|(name, question)| {
141            let json = serde_json::to_value(question).unwrap_or(Value::Null);
142            let kind = json
143                .get("type")
144                .and_then(Value::as_str)
145                .unwrap_or("raw")
146                .to_owned();
147            let shape = mock::answer(&session.state, name, &json).map(|answer| {
148                let mut map = serde_json::Map::new();
149                map.insert(name.clone(), mock::answer_json(&answer));
150                Value::Object(map)
151            });
152            QuestionEstimate {
153                name: name.clone(),
154                kind,
155                input_tokens: estimate_tokens(&format!("{}:{}", json!(name), json)),
156                output_tokens: match &shape {
157                    Some(value) => estimate_json_tokens(value),
158                    None => estimate_tokens(ASSUMED_ANSWER),
159                },
160                assumed: shape.is_none(),
161            }
162        })
163        .collect();
164
165    let input_tokens =
166        state_tokens + envelope_tokens + questions.iter().map(|q| q.input_tokens).sum::<usize>();
167    let output_tokens =
168        answer_envelope_tokens + questions.iter().map(|q| q.output_tokens).sum::<usize>();
169    Estimate {
170        model: model.to_owned(),
171        state_tokens,
172        envelope_tokens,
173        answer_envelope_tokens,
174        questions,
175        input_tokens,
176        output_tokens,
177    }
178}
179
180/// What a conversation has cost, as opposed to what one call costs.
181///
182/// A thread is not cheap the way it looks: the state is sent whole every time, so asking again
183/// after each turn is a call per turn over a state that keeps growing, and the tokens add up
184/// faster than the transcript does. This is the number that surprises people, so it is worth
185/// printing next to the per-call one.
186#[derive(Debug, Clone, PartialEq, Eq)]
187pub struct Thread {
188    /// Turns in the state.
189    pub turns: usize,
190    /// One estimate per turn: what asking after that turn cost.
191    pub calls: Vec<Estimate>,
192    /// Every call's input tokens, added up.
193    pub input_tokens: usize,
194    /// Every call's output tokens, added up.
195    pub output_tokens: usize,
196}
197
198/// Estimate a call per turn, or `None` when the state is not a conversation.
199pub fn thread(session: &Session, model: &str) -> Option<Thread> {
200    let turns = session.turns()?;
201    let mut so_far = session.clone();
202    let calls: Vec<Estimate> = (1..=turns.len())
203        .map(|n| {
204            so_far.state = turns_to_json(&turns[..n]);
205            estimate(&so_far, model)
206        })
207        .collect();
208    Some(Thread {
209        turns: turns.len(),
210        input_tokens: calls.iter().map(|c| c.input_tokens).sum(),
211        output_tokens: calls.iter().map(|c| c.output_tokens).sum(),
212        calls,
213    })
214}
215
216/// Price a pair of token counts.
217pub fn price(input_tokens: u64, output_tokens: u64, rates: Rates) -> Cost {
218    let input = input_tokens as f64 / 1_000_000.0 * rates.input;
219    let output = output_tokens as f64 / 1_000_000.0 * rates.output;
220    Cost {
221        input,
222        output,
223        total: input + output,
224    }
225}
226
227/// Price an estimate.
228pub fn price_estimate(estimate: &Estimate, rates: Rates) -> Cost {
229    price(
230        estimate.input_tokens as u64,
231        estimate.output_tokens as u64,
232        rates,
233    )
234}
235
236/// Price what a call actually used. `None` when the API reported no counts, because a made-up
237/// number is worse than none.
238pub fn price_usage(usage: &Usage, rates: Rates) -> Option<Cost> {
239    Some(price(usage.input_tokens?, usage.output_tokens?, rates))
240}
241
242/// `0.20/1.00`, `0.20 1.00`, `$0.20, $1.00` — input first, output second, per million tokens.
243pub fn parse_rates(text: &str) -> Result<Rates, String> {
244    let parts: Vec<&str> = text
245        .split(|c: char| c.is_whitespace() || c == '/' || c == ',')
246        .map(|p| p.trim().trim_start_matches('$'))
247        .filter(|p| !p.is_empty())
248        .collect();
249    if parts.len() != 2 {
250        return Err(
251            "Two rates, input then output, in dollars per million tokens: :cost 0.20/1.00"
252                .to_owned(),
253        );
254    }
255    let mut values = [0.0f64; 2];
256    for (slot, part) in values.iter_mut().zip(parts) {
257        match part.parse::<f64>() {
258            Ok(n) if n.is_finite() && n >= 0.0 => *slot = n,
259            _ => return Err(format!("{text:?} is not a pair of dollar amounts.")),
260        }
261    }
262    Ok(Rates {
263        input: values[0],
264        output: values[1],
265    })
266}
267
268/// Rates from `JEV_PRICE`; `None` when it is unset or malformed.
269pub fn rates_from_env() -> Option<Rates> {
270    rates_from_str(&std::env::var(PRICE_ENV).ok()?)
271}
272
273/// Rates from a stored string; `None` when it is empty or malformed.
274pub fn rates_from_str(value: &str) -> Option<Rates> {
275    if value.trim().is_empty() {
276        return None;
277    }
278    parse_rates(value).ok()
279}
280
281/// `$0.20/$1.00 per Mtok`.
282pub fn format_rates(rates: Rates) -> String {
283    format!(
284        "${}/${} per Mtok",
285        amount(rates.input),
286        amount(rates.output)
287    )
288}
289
290/// The same pair as `JEV_PRICE` takes: `0.20/1.00`.
291pub fn rates_value(rates: Rates) -> String {
292    format!("{}/{}", amount(rates.input), amount(rates.output))
293}
294
295/// Dollars, with enough decimals to be readable at the size a single call costs.
296pub fn usd(amount: f64) -> String {
297    if amount == 0.0 {
298        return "$0".to_owned();
299    }
300    if amount < 0.000_001 {
301        return "<$0.000001".to_owned();
302    }
303    let digits: usize = if amount >= 1.0 {
304        2
305    } else if amount >= 0.01 {
306        4
307    } else {
308        6
309    };
310    format!("${amount:.digits$}")
311}
312
313/// Rates are dollars: two decimals unless the rate is finer than a cent.
314fn amount(n: f64) -> String {
315    if (n * 100.0).fract() == 0.0 {
316        format!("{n:.2}")
317    } else {
318        format!("{n}")
319    }
320}