llm_sdk/
usage_ext.rs

1use crate::{LanguageModelPricing, ModelUsage};
2
3impl ModelUsage {
4    #[must_use]
5    pub fn calculate_cost(&self, pricing: &LanguageModelPricing) -> f64 {
6        let input_text_tokens = self
7            .input_tokens_details
8            .as_ref()
9            .and_then(|details| details.text_tokens)
10            .unwrap_or(self.input_tokens);
11        let input_audio_tokens = self
12            .input_tokens_details
13            .as_ref()
14            .and_then(|details| details.audio_tokens)
15            .unwrap_or(0);
16        let input_image_tokens = self
17            .input_tokens_details
18            .as_ref()
19            .and_then(|details| details.image_tokens)
20            .unwrap_or(0);
21        let input_cached_text_tokens = self
22            .input_tokens_details
23            .as_ref()
24            .and_then(|details| details.cached_text_tokens)
25            .unwrap_or(0);
26        let input_cached_audio_tokens = self
27            .input_tokens_details
28            .as_ref()
29            .and_then(|details| details.cached_audio_tokens)
30            .unwrap_or(0);
31        let input_cached_image_tokens = self
32            .input_tokens_details
33            .as_ref()
34            .and_then(|details| details.cached_image_tokens)
35            .unwrap_or(0);
36
37        let output_text_tokens = self
38            .output_tokens_details
39            .as_ref()
40            .and_then(|details| details.text_tokens)
41            .unwrap_or(self.output_tokens);
42        let output_audio_tokens = self
43            .output_tokens_details
44            .as_ref()
45            .and_then(|details| details.audio_tokens)
46            .unwrap_or(0);
47        let output_image_tokens = self
48            .output_tokens_details
49            .as_ref()
50            .and_then(|details| details.image_tokens)
51            .unwrap_or(0);
52
53        f64::from(input_text_tokens) * pricing.input_cost_per_text_token.unwrap_or(0.0)
54            + f64::from(input_audio_tokens) * pricing.input_cost_per_audio_token.unwrap_or(0.0)
55            + f64::from(input_image_tokens) * pricing.input_cost_per_image_token.unwrap_or(0.0)
56            + f64::from(input_cached_text_tokens) * pricing.input_cost_per_text_token.unwrap_or(0.0)
57            + f64::from(input_cached_audio_tokens)
58                * pricing.input_cost_per_audio_token.unwrap_or(0.0)
59            + f64::from(input_cached_image_tokens)
60                * pricing.input_cost_per_image_token.unwrap_or(0.0)
61            + f64::from(output_text_tokens) * pricing.output_cost_per_text_token.unwrap_or(0.0)
62            + f64::from(output_audio_tokens) * pricing.output_cost_per_audio_token.unwrap_or(0.0)
63            + f64::from(output_image_tokens) * pricing.output_cost_per_image_token.unwrap_or(0.0)
64    }
65
66    pub fn add(&mut self, other: &Self) {
67        self.input_tokens += other.input_tokens;
68        self.output_tokens += other.output_tokens;
69
70        if let Some(other_input_details) = &other.input_tokens_details {
71            let self_input_details = self.input_tokens_details.get_or_insert_default();
72            if let Some(text_tokens) = other_input_details.text_tokens {
73                self_input_details.text_tokens =
74                    Some(self_input_details.text_tokens.unwrap_or(0) + text_tokens);
75            }
76            if let Some(audio_tokens) = other_input_details.audio_tokens {
77                self_input_details.audio_tokens =
78                    Some(self_input_details.audio_tokens.unwrap_or(0) + audio_tokens);
79            }
80            if let Some(image_tokens) = other_input_details.image_tokens {
81                self_input_details.image_tokens =
82                    Some(self_input_details.image_tokens.unwrap_or(0) + image_tokens);
83            }
84            if let Some(cached_text_tokens) = other_input_details.cached_text_tokens {
85                self_input_details.cached_text_tokens =
86                    Some(self_input_details.cached_text_tokens.unwrap_or(0) + cached_text_tokens);
87            }
88            if let Some(cached_audio_tokens) = other_input_details.cached_audio_tokens {
89                self_input_details.cached_audio_tokens =
90                    Some(self_input_details.cached_audio_tokens.unwrap_or(0) + cached_audio_tokens);
91            }
92            if let Some(cached_image_tokens) = other_input_details.cached_image_tokens {
93                self_input_details.cached_image_tokens =
94                    Some(self_input_details.cached_image_tokens.unwrap_or(0) + cached_image_tokens);
95            }
96        }
97
98        if let Some(other_output_details) = &other.output_tokens_details {
99            let self_output_details = self.output_tokens_details.get_or_insert_default();
100            if let Some(text_tokens) = other_output_details.text_tokens {
101                self_output_details.text_tokens =
102                    Some(self_output_details.text_tokens.unwrap_or(0) + text_tokens);
103            }
104            if let Some(audio_tokens) = other_output_details.audio_tokens {
105                self_output_details.audio_tokens =
106                    Some(self_output_details.audio_tokens.unwrap_or(0) + audio_tokens);
107            }
108            if let Some(image_tokens) = other_output_details.image_tokens {
109                self_output_details.image_tokens =
110                    Some(self_output_details.image_tokens.unwrap_or(0) + image_tokens);
111            }
112        }
113    }
114}