Skip to main content

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    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    /// Other reported quantities, e.g. audio_seconds, image_outputs, web_searches.
40    /// Do not duplicate named token fields here; absence is not a measured zero.
41    pub metrics: BTreeMap<String, Decimal>,
42    /// Pricing qualifiers such as resolution, speed or inference_geo.
43    pub dimensions: BTreeMap<String, String>,
44    /// Actual upstream serving tier, which may differ from the requested tier.
45    pub actual_service_tier: Option<String>,
46    pub completeness: UsageCompleteness,
47    /// When present, per-attempt usage replaces aggregate usage for billing.
48    pub attempts: Vec<UsageAttempt>,
49    /// Independently identified upstream responses inside a long-lived call.
50    /// Their sum is the aggregate above; the settlement layer may deduplicate
51    /// these across connections before pricing. IDs are upstream identities.
52    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    /// Sum reported response quantities. Missing values remain unknown rather
63    /// than becoming invented zeros. Qualifiers survive only when all agree.
64    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                // Saturating: upstream-reported amounts, and a Decimal
100                // overflow panics.
101                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    /// None when the provider did not establish whether this attempt is charged.
129    pub billable: Option<bool>,
130    pub started_at_ms: Option<i64>,
131}
132
133/// How a streamed response reaches whoever watches it for usage.
134#[derive(Debug, Clone, Copy, PartialEq, Eq)]
135pub enum UsageTransport {
136    Http { framing: Option<StreamFraming> },
137    WebSocket,
138}
139
140/// How a watched stream ended. EOF alone does not establish complete
141/// provider usage; `Interrupted` means the stream stopped before its end.
142#[derive(Debug, Clone, Copy, PartialEq, Eq)]
143pub enum UsageStreamEnd {
144    Complete,
145    Interrupted,
146}