1use serde_json::{Value, json};
15use typesafe::Usage;
16
17use crate::mock;
18use crate::session::{Session, turns_to_json};
19
20pub const PRICE_ENV: &str = "JEV_PRICE";
22
23const ASSUMED_ANSWER: &str = r#"{"name":{"type":"noul","noul":0.123}}"#;
25
26#[derive(Debug, Clone, Copy, PartialEq)]
28pub struct Rates {
29 pub input: f64,
31 pub output: f64,
33}
34
35#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct QuestionEstimate {
38 pub name: String,
39 pub kind: String,
41 pub input_tokens: usize,
43 pub output_tokens: usize,
45 pub assumed: bool,
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
51pub struct Estimate {
52 pub model: String,
53 pub state_tokens: usize,
55 pub envelope_tokens: usize,
57 pub answer_envelope_tokens: usize,
59 pub questions: Vec<QuestionEstimate>,
60 pub input_tokens: usize,
62 pub output_tokens: usize,
64}
65
66#[derive(Debug, Clone, Copy, PartialEq)]
68pub struct Cost {
69 pub input: f64,
70 pub output: f64,
71 pub total: f64,
72}
73
74pub fn estimate_tokens(text: &str) -> usize {
80 #[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 tokens += runs.flush();
102 } else if ch as u32 >= 0x2e80 {
103 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
125pub fn estimate_json_tokens(value: &Value) -> usize {
127 estimate_tokens(&value.to_string())
128}
129
130pub 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#[derive(Debug, Clone, PartialEq, Eq)]
187pub struct Thread {
188 pub turns: usize,
190 pub calls: Vec<Estimate>,
192 pub input_tokens: usize,
194 pub output_tokens: usize,
196}
197
198pub 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
216pub 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
227pub 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
236pub fn price_usage(usage: &Usage, rates: Rates) -> Option<Cost> {
239 Some(price(usage.input_tokens?, usage.output_tokens?, rates))
240}
241
242pub 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
268pub fn rates_from_env() -> Option<Rates> {
270 rates_from_str(&std::env::var(PRICE_ENV).ok()?)
271}
272
273pub 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
281pub fn format_rates(rates: Rates) -> String {
283 format!(
284 "${}/${} per Mtok",
285 amount(rates.input),
286 amount(rates.output)
287 )
288}
289
290pub fn rates_value(rates: Rates) -> String {
292 format!("{}/{}", amount(rates.input), amount(rates.output))
293}
294
295pub 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
313fn amount(n: f64) -> String {
315 if (n * 100.0).fract() == 0.0 {
316 format!("{n:.2}")
317 } else {
318 format!("{n}")
319 }
320}