bamboo_domain/
token_usage.rs

1//! Shared token usage value object used across domain, application, and infrastructure layers.
2
3use serde::{Deserialize, Serialize};
4
5/// Token consumption statistics for a single LLM call or aggregated period.
6///
7/// This is a stable, cross-layer value object. Every crate that needs to
8/// represent "how many tokens were used" should use this type (or re-export it)
9/// instead of defining a local duplicate.
10#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
11pub struct TokenUsage {
12    /// Tokens in the LLM request (prompt / input).
13    pub prompt_tokens: u64,
14    /// Tokens in the LLM response (completion / output).
15    pub completion_tokens: u64,
16    /// Total tokens consumed (normally prompt + completion).
17    pub total_tokens: u64,
18}
19
20impl TokenUsage {
21    /// Accumulate another usage snapshot into this one.
22    pub fn add_assign(&mut self, other: TokenUsage) {
23        self.prompt_tokens = self.prompt_tokens.saturating_add(other.prompt_tokens);
24        self.completion_tokens = self
25            .completion_tokens
26            .saturating_add(other.completion_tokens);
27        self.total_tokens = self.total_tokens.saturating_add(other.total_tokens);
28    }
29
30    /// Recompute `total_tokens` from the two component fields.
31    pub fn recompute_total(&mut self) {
32        self.total_tokens = self.prompt_tokens.saturating_add(self.completion_tokens);
33    }
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39
40    #[test]
41    fn default_is_zero() {
42        let usage = TokenUsage::default();
43        assert_eq!(usage.prompt_tokens, 0);
44        assert_eq!(usage.completion_tokens, 0);
45        assert_eq!(usage.total_tokens, 0);
46    }
47
48    #[test]
49    fn add_assign_accumulates() {
50        let mut usage1 = TokenUsage {
51            prompt_tokens: 100,
52            completion_tokens: 50,
53            total_tokens: 150,
54        };
55        let usage2 = TokenUsage {
56            prompt_tokens: 200,
57            completion_tokens: 100,
58            total_tokens: 300,
59        };
60        usage1.add_assign(usage2);
61        assert_eq!(usage1.prompt_tokens, 300);
62        assert_eq!(usage1.completion_tokens, 150);
63        assert_eq!(usage1.total_tokens, 450);
64    }
65
66    #[test]
67    fn recompute_total_uses_saturating_add() {
68        let mut usage = TokenUsage {
69            prompt_tokens: u64::MAX - 5,
70            completion_tokens: u64::MAX - 9,
71            total_tokens: 0,
72        };
73        usage.recompute_total();
74        assert_eq!(usage.total_tokens, u64::MAX);
75    }
76}
bamboo_domain/token_usage.rs

bamboo_domain/
token_usage.rs
