typesafe_jev/
usage.rs

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

typesafe_jev/
usage.rs
