ai_agent/services/api/
empty_usage.rs

1//! Empty usage tracking - translated from emptyUsage.ts
2
3use serde::{Deserialize, Serialize};
4
5/// Token usage for an API call
6#[derive(Debug, Clone, Default, Serialize, Deserialize)]
7pub struct Usage {
8    #[serde(rename = "input_tokens")]
9    pub input_tokens: i64,
10    #[serde(rename = "output_tokens")]
11    pub output_tokens: i64,
12    #[serde(rename = "cache_creation_input_tokens")]
13    pub cache_creation_input_tokens: i64,
14    #[serde(rename = "cache_hit_input_tokens")]
15    pub cache_hit_input_tokens: i64,
16}
17
18impl Usage {
19    pub fn new() -> Self {
20        Self::default()
21    }
22
23    pub fn with_input_tokens(mut self, tokens: i64) -> Self {
24        self.input_tokens = tokens;
25        self
26    }
27
28    pub fn with_output_tokens(mut self, tokens: i64) -> Self {
29        self.output_tokens = tokens;
30        self
31    }
32
33    pub fn total(&self) -> i64 {
34        self.input_tokens + self.output_tokens
35    }
36}
37
38/// Create an empty usage struct (for errors, etc.)
39pub fn empty_usage() -> Usage {
40    Usage::default()
41}
42
43/// Create a usage struct with just input tokens
44pub fn input_only_usage(tokens: i64) -> Usage {
45    Usage::new().with_input_tokens(tokens)
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51
52    #[test]
53    fn test_usage_default() {
54        let usage = Usage::default();
55        assert_eq!(usage.input_tokens, 0);
56        assert_eq!(usage.output_tokens, 0);
57        assert_eq!(usage.total(), 0);
58    }
59
60    #[test]
61    fn test_usage_with_tokens() {
62        let usage = Usage::new().with_input_tokens(100).with_output_tokens(50);
63
64        assert_eq!(usage.input_tokens, 100);
65        assert_eq!(usage.output_tokens, 50);
66        assert_eq!(usage.total(), 150);
67    }
68}
ai_agent/services/api/empty_usage.rs

ai_agent/services/api/
empty_usage.rs
