1use serde::{Deserialize, Serialize};
2use std::collections::BTreeMap;
3
4#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
21#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
22pub struct Usage {
23 pub input_tokens: u32,
24 #[serde(default)]
25 pub cache_read_tokens: u32,
26 #[serde(default)]
27 pub cache_write_tokens: u32,
28 pub output_tokens: u32,
29 #[serde(default)]
30 pub reasoning_tokens: u32,
31 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
32 pub server_tool_calls: BTreeMap<String, u32>,
33}
34
35impl Usage {
36 pub fn prompt_tokens(&self) -> u32 {
39 self.input_tokens + self.cache_read_tokens + self.cache_write_tokens
40 }
41
42 pub fn completion_tokens(&self) -> u32 {
45 self.output_tokens + self.reasoning_tokens
46 }
47
48 pub fn total_tokens(&self) -> u32 {
50 self.prompt_tokens() + self.completion_tokens()
51 }
52}
53
54mod peek {
55 use crate::types::{AnthropicUsage, GeminiUsage, OpenAiUsage};
56 use serde::Deserialize;
57
58 #[derive(Deserialize)]
59 pub struct OpenAi {
60 pub usage: Option<OpenAiUsage>,
61 }
62
63 #[derive(Deserialize)]
64 pub struct Anthropic {
65 pub usage: Option<AnthropicUsage>,
66 }
67
68 #[derive(Deserialize)]
69 #[serde(rename_all = "camelCase")]
70 pub struct Gemini {
71 pub usage_metadata: Option<GeminiUsage>,
72 }
73}
74
75impl From<&[u8]> for Usage {
76 fn from(raw: &[u8]) -> Self {
77 if let Ok(peek::OpenAi { usage: Some(u) }) = crate::json::from_slice(raw)
78 && (u.prompt_tokens > 0 || u.completion_tokens > 0)
79 {
80 return Usage::from(&u);
81 }
82
83 if let Ok(peek::Anthropic { usage: Some(u) }) = crate::json::from_slice(raw)
84 && (u.input_tokens > 0 || u.output_tokens > 0)
85 {
86 return Usage::from(&u);
87 }
88
89 if let Ok(peek::Gemini {
90 usage_metadata: Some(u),
91 }) = crate::json::from_slice(raw)
92 && u.total_token_count > 0
93 {
94 return Usage::from(&u);
95 }
96
97 Usage::default()
98 }
99}