Skip to main content

anthropic_tools/common/
usage.rs

1//! Token usage information from API responses.
2//!
3//! This module provides the [`Usage`] struct for tracking token consumption:
4//!
5//! - Input tokens consumed
6//! - Output tokens generated
7//! - Cache creation and read tokens (for prompt caching)
8//!
9//! # Example
10//!
11//! ```rust
12//! use anthropic_tools::common::usage::Usage;
13//!
14//! let usage = Usage::new(100, 50);
15//! assert_eq!(usage.input_tokens, 100);
16//! assert_eq!(usage.output_tokens, 50);
17//! assert_eq!(usage.total_tokens(), 150);
18//! ```
19//!
20//! # With Prompt Caching
21//!
22//! When using prompt caching, the usage will include cache-related fields:
23//!
24//! ```rust
25//! use anthropic_tools::common::usage::Usage;
26//!
27//! let mut usage = Usage::new(100, 50);
28//! // These would be set by the API response
29//! // usage.cache_creation_input_tokens = Some(50);
30//! // usage.cache_read_input_tokens = Some(25);
31//! assert_eq!(usage.cached_tokens(), 0); // No cache tokens in this example
32//! ```
33
34use serde::{Deserialize, Serialize};
35
36/// Token usage information from Anthropic API response
37#[derive(Debug, Clone, Default, Serialize, Deserialize)]
38pub struct Usage {
39    /// The number of input tokens used
40    pub input_tokens: usize,
41
42    /// The number of output tokens generated
43    pub output_tokens: usize,
44
45    /// The number of input tokens used to create the cache entry
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub cache_creation_input_tokens: Option<usize>,
48
49    /// The number of input tokens read from the cache
50    #[serde(skip_serializing_if = "Option::is_none")]
51    pub cache_read_input_tokens: Option<usize>,
52}
53
54impl Usage {
55    /// Create a new Usage instance
56    pub fn new(input_tokens: usize, output_tokens: usize) -> Self {
57        Self {
58            input_tokens,
59            output_tokens,
60            cache_creation_input_tokens: None,
61            cache_read_input_tokens: None,
62        }
63    }
64
65    /// Get total tokens (input + output)
66    pub fn total_tokens(&self) -> usize {
67        self.input_tokens + self.output_tokens
68    }
69
70    /// Get total cached tokens
71    pub fn cached_tokens(&self) -> usize {
72        self.cache_creation_input_tokens.unwrap_or(0) + self.cache_read_input_tokens.unwrap_or(0)
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn test_usage_new() {
82        let usage = Usage::new(100, 50);
83        assert_eq!(usage.input_tokens, 100);
84        assert_eq!(usage.output_tokens, 50);
85        assert_eq!(usage.total_tokens(), 150);
86    }
87
88    #[test]
89    fn test_usage_deserialize() {
90        let json = r#"{
91            "input_tokens": 100,
92            "output_tokens": 50,
93            "cache_creation_input_tokens": 10,
94            "cache_read_input_tokens": 20
95        }"#;
96
97        let usage: Usage = serde_json::from_str(json).unwrap();
98        assert_eq!(usage.input_tokens, 100);
99        assert_eq!(usage.output_tokens, 50);
100        assert_eq!(usage.cache_creation_input_tokens, Some(10));
101        assert_eq!(usage.cache_read_input_tokens, Some(20));
102        assert_eq!(usage.cached_tokens(), 30);
103    }
104
105    #[test]
106    fn test_usage_serialize() {
107        let usage = Usage::new(100, 50);
108        let json = serde_json::to_string(&usage).unwrap();
109        assert!(json.contains("\"input_tokens\":100"));
110        assert!(json.contains("\"output_tokens\":50"));
111        // cache fields should be omitted when None
112        assert!(!json.contains("cache_creation_input_tokens"));
113        assert!(!json.contains("cache_read_input_tokens"));
114    }
115}