gproxy_protocol/usage/
types.rs1use std::collections::BTreeMap;
9
10use rust_decimal::Decimal;
11
12use crate::connection::StreamFraming;
13
14#[derive(Debug, Clone, Default, PartialEq, Eq)]
16pub struct TokenUsage {
17 pub input_tokens: Option<u64>,
19 pub output_tokens: Option<u64>,
21 pub cached_input_tokens: Option<u64>,
22 pub cache_creation_5m_tokens: Option<u64>,
23 pub cache_creation_30m_tokens: Option<u64>,
24 pub cache_creation_1h_tokens: Option<u64>,
25 pub reasoning_tokens: Option<u64>,
26}
27
28#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
29pub enum UsageCompleteness {
30 Complete,
31 Partial,
32 #[default]
33 Unknown,
34}
35
36#[derive(Debug, Clone, Default, PartialEq, Eq)]
37pub struct NormalizedUsage {
38 pub tokens: TokenUsage,
39 pub metrics: BTreeMap<String, Decimal>,
42 pub dimensions: BTreeMap<String, String>,
44 pub actual_service_tier: Option<String>,
46 pub completeness: UsageCompleteness,
47 pub attempts: Vec<UsageAttempt>,
49 pub responses: Vec<ResponseUsage>,
53}
54
55#[derive(Debug, Clone, PartialEq, Eq)]
56pub struct ResponseUsage {
57 pub id: String,
58 pub usage: Box<NormalizedUsage>,
59}
60
61impl NormalizedUsage {
62 pub fn aggregate<'a>(values: impl IntoIterator<Item = &'a NormalizedUsage>) -> Self {
65 let mut total = Self {
66 completeness: UsageCompleteness::Complete,
67 ..Default::default()
68 };
69 let mut first = true;
70 fn add(target: &mut Option<u64>, value: Option<u64>) {
71 if let Some(v) = value {
72 *target = Some(target.unwrap_or(0).saturating_add(v));
73 }
74 }
75 for value in values {
76 add(&mut total.tokens.input_tokens, value.tokens.input_tokens);
77 add(&mut total.tokens.output_tokens, value.tokens.output_tokens);
78 add(
79 &mut total.tokens.cached_input_tokens,
80 value.tokens.cached_input_tokens,
81 );
82 add(
83 &mut total.tokens.reasoning_tokens,
84 value.tokens.reasoning_tokens,
85 );
86 add(
87 &mut total.tokens.cache_creation_5m_tokens,
88 value.tokens.cache_creation_5m_tokens,
89 );
90 add(
91 &mut total.tokens.cache_creation_30m_tokens,
92 value.tokens.cache_creation_30m_tokens,
93 );
94 add(
95 &mut total.tokens.cache_creation_1h_tokens,
96 value.tokens.cache_creation_1h_tokens,
97 );
98 for (key, count) in &value.metrics {
99 let sum = total.metrics.entry(key.clone()).or_default();
102 *sum = sum.saturating_add(*count);
103 }
104 if first {
105 total.dimensions = value.dimensions.clone();
106 total.actual_service_tier = value.actual_service_tier.clone();
107 first = false;
108 } else {
109 total
110 .dimensions
111 .retain(|key, v| value.dimensions.get(key) == Some(v));
112 if total.actual_service_tier != value.actual_service_tier {
113 total.actual_service_tier = None;
114 }
115 }
116 if value.completeness != UsageCompleteness::Complete {
117 total.completeness = UsageCompleteness::Partial;
118 }
119 }
120 total
121 }
122}
123
124#[derive(Debug, Clone, PartialEq, Eq)]
125pub struct UsageAttempt {
126 pub model: String,
127 pub usage: Box<NormalizedUsage>,
128 pub billable: Option<bool>,
130 pub started_at_ms: Option<i64>,
131}
132
133#[derive(Debug, Clone, Copy, PartialEq, Eq)]
135pub enum UsageTransport {
136 Http { framing: Option<StreamFraming> },
137 WebSocket,
138}
139
140#[derive(Debug, Clone, Copy, PartialEq, Eq)]
143pub enum UsageStreamEnd {
144 Complete,
145 Interrupted,
146}