typesafe_jev/
usage.rs

1use std::sync::atomic::{AtomicU64, Ordering};
2
3use serde_json::Value;
4
5use crate::USD_PER_INPUT_MTOK;
6
7/// Thread-safe request and token counters.
8///
9/// A [`Client`](crate::Client) fills its own `Usage` from the `usage` object of every successful
10/// reply and counts each retry. The counters are atomic, so the totals can be read while requests
11/// are in flight; they are then a snapshot, not a transaction.
12///
13/// [`Usage::record`] and [`Usage::record_retry`] are public so that a caller who talks to another
14/// provider through the same code path can keep the same accounting.
15#[derive(Debug, Default)]
16pub struct Usage {
17    requests: AtomicU64,
18    retries: AtomicU64,
19    input_tokens: AtomicU64,
20    output_tokens: AtomicU64,
21}
22
23impl Usage {
24    /// One completed request and the tokens it used.
25    pub fn record(&self, input_tokens: u64, output_tokens: u64) {
26        self.requests.fetch_add(1, Ordering::Relaxed);
27        self.input_tokens.fetch_add(input_tokens, Ordering::Relaxed);
28        self.output_tokens.fetch_add(output_tokens, Ordering::Relaxed);
29    }
30
31    /// One attempt that had to be repeated. A request that is retried twice before it succeeds
32    /// counts as one request and two retries.
33    pub fn record_retry(&self) {
34        self.retries.fetch_add(1, Ordering::Relaxed);
35    }
36
37    /// One completed request, from the reply's `usage` object (`input_tokens` and
38    /// `output_tokens`; a missing or malformed field counts as zero).
39    pub(crate) fn record_reported(&self, usage: Option<&Value>) {
40        let field = |k: &str| usage.and_then(|u| u.get(k)).and_then(Value::as_u64).unwrap_or(0);
41        self.record(field("input_tokens"), field("output_tokens"));
42    }
43
44    /// Requests that were answered.
45    pub fn requests(&self) -> u64 {
46        self.requests.load(Ordering::Relaxed)
47    }
48
49    /// Attempts that were repeated after a transient failure.
50    pub fn retries(&self) -> u64 {
51        self.retries.load(Ordering::Relaxed)
52    }
53
54    /// Input tokens, as reported by the API.
55    pub fn input_tokens(&self) -> u64 {
56        self.input_tokens.load(Ordering::Relaxed)
57    }
58
59    /// Output tokens, as reported by the API. Jev reports them, but does not charge for them.
60    pub fn output_tokens(&self) -> u64 {
61        self.output_tokens.load(Ordering::Relaxed)
62    }
63
64    /// The input tokens at Jev's list price, [`USD_PER_INPUT_MTOK`]. An estimate: the price is a
65    /// constant in this crate, not something the API reports.
66    pub fn cost_usd(&self) -> f64 {
67        self.input_tokens() as f64 / 1_000_000.0 * USD_PER_INPUT_MTOK
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use serde_json::json;
75
76    #[test]
77    fn counts_requests_retries_and_tokens() {
78        let usage = Usage::default();
79        assert_eq!((usage.requests(), usage.retries(), usage.input_tokens(), usage.output_tokens()), (0, 0, 0, 0));
80        usage.record(100, 10);
81        usage.record_retry();
82        usage.record_reported(Some(&json!({"input_tokens": 50, "output_tokens": 5})));
83        usage.record_reported(Some(&json!({"input_tokens": "not a number"})));
84        usage.record_reported(None);
85        assert_eq!((usage.requests(), usage.retries(), usage.input_tokens(), usage.output_tokens()), (4, 1, 150, 15));
86        assert!((usage.cost_usd() - 150.0 / 1_000_000.0 * USD_PER_INPUT_MTOK).abs() < 1e-15);
87    }
88}
typesafe_jev/usage.rs

typesafe_jev/
usage.rs
