gproxy_protocol/usage/
types.rs

1//! The metering vocabulary every reader produces and every consumer settles.
2//!
3//! These types used to live beside the channel usage traits. They moved here
4//! because the standard readers in this module produce them from the standard
5//! response shapes, and a crate that only knows dialects cannot depend on one
6//! that knows channels. `gproxy-channel` re-exports them at their old paths.
7
8use std::collections::BTreeMap;
9
10use rust_decimal::Decimal;
11
12use crate::connection::StreamFraming;
13
14/// Reported counts only. None is unknown; Some(0) is an explicit zero.
15#[derive(Debug, Clone, Default, PartialEq, Eq)]
16pub struct TokenUsage {
17    /// Ordinary input, excluding cache reads and cache creation.
18    pub input_tokens: Option<u64>,
19    /// Total output including reasoning; reasoning is a subset, not additive.
20    pub output_tokens: Option<u64>,
21    pub cached_input_tokens: Option<u64>,
22    /// Reported writes grouped by retention. Empty means unreported.
23    pub cache_creation: Vec<CacheCreation>,
24    pub reasoning_tokens: Option<u64>,
25}
26
27/// One reported cache-write quantity. Unknown retention is not a default TTL.
28#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
29pub struct CacheCreation {
30    pub tokens: u64,
31    pub ttl_seconds: Option<u64>,
32}
33
34impl TokenUsage {
35    pub fn cache_creation_tokens(&self) -> u64 {
36        self.cache_creation
37            .iter()
38            .fold(0, |sum, entry| sum.saturating_add(entry.tokens))
39    }
40
41    pub fn add_cache_creation(&mut self, tokens: Option<u64>, ttl_seconds: Option<u64>) {
42        let Some(tokens) = tokens else { return };
43        if let Some(entry) = self
44            .cache_creation
45            .iter_mut()
46            .find(|entry| entry.ttl_seconds == ttl_seconds)
47        {
48            entry.tokens = entry.tokens.saturating_add(tokens);
49        } else {
50            self.cache_creation.push(CacheCreation {
51                tokens,
52                ttl_seconds,
53            });
54        }
55    }
56}
57
58#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
59pub enum UsageCompleteness {
60    Complete,
61    Partial,
62    #[default]
63    Unknown,
64}
65
66#[derive(Debug, Clone, Default, PartialEq, Eq)]
67pub struct NormalizedUsage {
68    pub tokens: TokenUsage,
69    /// Other reported quantities, e.g. audio_seconds, image_outputs, web_searches.
70    /// Do not duplicate named token fields here; absence is not a measured zero.
71    pub metrics: BTreeMap<String, Decimal>,
72    /// Pricing qualifiers such as resolution, speed or inference_geo.
73    pub dimensions: BTreeMap<String, String>,
74    /// Actual upstream serving tier, which may differ from the requested tier.
75    pub actual_service_tier: Option<String>,
76    pub completeness: UsageCompleteness,
77    /// When present, per-attempt usage replaces aggregate usage for billing.
78    pub attempts: Vec<UsageAttempt>,
79    /// Independently identified upstream responses inside a long-lived call.
80    /// Their sum is the aggregate above; the settlement layer may deduplicate
81    /// these across connections before pricing. IDs are upstream identities.
82    pub responses: Vec<ResponseUsage>,
83}
84
85#[derive(Debug, Clone, PartialEq, Eq)]
86pub struct ResponseUsage {
87    pub id: String,
88    pub usage: Box<NormalizedUsage>,
89}
90
91impl NormalizedUsage {
92    /// Sum reported response quantities. Missing values remain unknown rather
93    /// than becoming invented zeros. Qualifiers survive only when all agree.
94    pub fn aggregate<'a>(values: impl IntoIterator<Item = &'a NormalizedUsage>) -> Self {
95        let mut total = Self {
96            completeness: UsageCompleteness::Complete,
97            ..Default::default()
98        };
99        let mut first = true;
100        fn add(target: &mut Option<u64>, value: Option<u64>) {
101            if let Some(v) = value {
102                *target = Some(target.unwrap_or(0).saturating_add(v));
103            }
104        }
105        for value in values {
106            add(&mut total.tokens.input_tokens, value.tokens.input_tokens);
107            add(&mut total.tokens.output_tokens, value.tokens.output_tokens);
108            add(
109                &mut total.tokens.cached_input_tokens,
110                value.tokens.cached_input_tokens,
111            );
112            add(
113                &mut total.tokens.reasoning_tokens,
114                value.tokens.reasoning_tokens,
115            );
116            for entry in &value.tokens.cache_creation {
117                total
118                    .tokens
119                    .add_cache_creation(Some(entry.tokens), entry.ttl_seconds);
120            }
121            for (key, count) in &value.metrics {
122                // Saturating: upstream-reported amounts, and a Decimal
123                // overflow panics.
124                let sum = total.metrics.entry(key.clone()).or_default();
125                *sum = sum.saturating_add(*count);
126            }
127            if first {
128                total.dimensions = value.dimensions.clone();
129                total.actual_service_tier = value.actual_service_tier.clone();
130                first = false;
131            } else {
132                total
133                    .dimensions
134                    .retain(|key, v| value.dimensions.get(key) == Some(v));
135                if total.actual_service_tier != value.actual_service_tier {
136                    total.actual_service_tier = None;
137                }
138            }
139            if value.completeness != UsageCompleteness::Complete {
140                total.completeness = UsageCompleteness::Partial;
141            }
142        }
143        total
144    }
145}
146
147#[derive(Debug, Clone, PartialEq, Eq)]
148pub struct UsageAttempt {
149    /// A delegated API may have a different billing operation from its parent session.
150    pub operation: Option<crate::Operation>,
151    pub model: String,
152    pub usage: Box<NormalizedUsage>,
153    /// None when the provider did not establish whether this attempt is charged.
154    pub billable: Option<bool>,
155    pub started_at_ms: Option<i64>,
156}
157
158/// How a streamed response reaches whoever watches it for usage.
159#[derive(Debug, Clone, Copy, PartialEq, Eq)]
160pub enum UsageTransport {
161    Http { framing: Option<StreamFraming> },
162    WebSocket,
163}
164
165/// How a watched stream ended. EOF alone does not establish complete
166/// provider usage; `Interrupted` means the stream stopped before its end.
167#[derive(Debug, Clone, Copy, PartialEq, Eq)]
168pub enum UsageStreamEnd {
169    Complete,
170    Interrupted,
171}
gproxy_protocol/usage/types.rs

gproxy_protocol/usage/
types.rs
