gproxy_protocol/usage/
output.rs

1//! Visible generated characters for fallback metering. Never count wire framing,
2//! media, signatures or response snapshots repeated after stream deltas.
3use serde_json::Value;
4
5use crate::{Dialect, Operation};
6
7fn text(value: &Value) -> u64 {
8    value.as_str().map_or(0, |text| text.chars().count() as u64)
9}
10
11fn items(value: &Value) -> impl Iterator<Item = &Value> {
12    value.as_array().into_iter().flatten()
13}
14
15fn arguments(value: &Value) -> u64 {
16    if value.is_null() || value.as_object().is_some_and(|object| object.is_empty()) {
17        0
18    } else if value.is_string() {
19        text(value)
20    } else {
21        value.to_string().chars().count() as u64
22    }
23}
24
25fn chat(message: &Value) -> u64 {
26    // Compatible providers can repeat the same reasoning under both names.
27    let reasoning = message["reasoning_content"]
28        .as_str()
29        .or_else(|| message["reasoning"].as_str())
30        .map_or(0, |text| text.chars().count() as u64);
31    text(&message["content"])
32        + text(&message["refusal"])
33        + reasoning
34        + text(&message["function_call"]["arguments"])
35        + items(&message["tool_calls"])
36            .map(|call| text(&call["function"]["arguments"]))
37            .sum::<u64>()
38}
39
40fn claude(block: &Value) -> u64 {
41    match block["type"].as_str() {
42        Some("text") => text(&block["text"]),
43        Some("thinking") => text(&block["thinking"]),
44        Some("tool_use" | "server_tool_use") => arguments(&block["input"]),
45        _ => 0,
46    }
47}
48
49fn gemini(value: &Value) -> u64 {
50    items(&value["candidates"])
51        .flat_map(|candidate| items(&candidate["content"]["parts"]))
52        .map(|part| text(&part["text"]) + arguments(&part["functionCall"]["args"]))
53        .sum()
54}
55
56pub(super) fn stream(dialect: Dialect, name: Option<&str>, value: &Value) -> u64 {
57    let kind = value["type"].as_str().or(name);
58    match dialect {
59        Dialect::OpenAiChat => items(&value["choices"])
60            .map(|choice| chat(&choice["delta"]))
61            .sum(),
62        Dialect::Gemini => gemini(value),
63        Dialect::Claude => match kind {
64            Some("content_block_start") => claude(&value["content_block"]),
65            Some("content_block_delta") => {
66                let delta = &value["delta"];
67                text(&delta["text"]) + text(&delta["thinking"]) + text(&delta["partial_json"])
68            }
69            _ => 0,
70        },
71        Dialect::OpenAi | Dialect::OpenAiResponsesWebSocket => match kind {
72            Some(
73                "response.output_text.delta"
74                | "response.reasoning_text.delta"
75                | "response.reasoning_summary_text.delta"
76                | "response.function_call_arguments.delta"
77                | "response.custom_tool_call_input.delta"
78                | "response.refusal.delta"
79                | "transcript.text.delta"
80                | "session.output_transcript.delta",
81            ) => text(&value["delta"]),
82            _ => 0,
83        },
84    }
85}
86
87/// Generated characters in a complete standard response. Unsupported bodies
88/// yield zero; binary/media sizes are never a substitute for output tokens.
89pub fn whole_output_characters(operation: Operation, dialect: Dialect, body: &[u8]) -> u64 {
90    if !matches!(
91        operation,
92        Operation::GenerateContent
93            | Operation::StreamGenerateContent
94            | Operation::CompactContent
95            | Operation::CreateDecision
96            | Operation::CreateTranscription
97            | Operation::CreateTranslation
98    ) {
99        return 0;
100    }
101    let Ok(value) = serde_json::from_slice::<Value>(body) else {
102        return 0;
103    };
104    if matches!(
105        operation,
106        Operation::CreateTranscription | Operation::CreateTranslation
107    ) {
108        return text(&value["text"]);
109    }
110    match dialect {
111        Dialect::OpenAiChat => items(&value["choices"])
112            .map(|choice| chat(&choice["message"]))
113            .sum(),
114        Dialect::Gemini => gemini(&value),
115        Dialect::Claude => items(&value["content"]).map(claude).sum(),
116        Dialect::OpenAi | Dialect::OpenAiResponsesWebSocket => items(&value["output"])
117            .map(|item| match item["type"].as_str() {
118                Some("message") => items(&item["content"])
119                    .map(|part| match part["type"].as_str() {
120                        Some("output_text") => text(&part["text"]),
121                        Some("refusal") => text(&part["refusal"]),
122                        _ => 0,
123                    })
124                    .sum(),
125                Some("reasoning") => {
126                    let content = item
127                        .get("content")
128                        .filter(|value| value.is_array())
129                        .unwrap_or(&item["summary"]);
130                    items(content).map(|part| text(&part["text"])).sum()
131                }
132                Some("function_call") => text(&item["arguments"]),
133                Some("custom_tool_call") => text(&item["input"]),
134                _ => 0,
135            })
136            .sum(),
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143    use crate::usage::{UsageReader, UsageTransport};
144
145    #[test]
146    fn fragmented_streams_count_content_without_metadata_or_terminal_replays() {
147        for (dialect, event, replay) in [
148            (
149                Dialect::OpenAiChat,
150                r#"{"choices":[{"delta":{"content":"你好","reasoning_content":"hmm","tool_calls":[{"function":{"arguments":"{}"}}]}}]}"#,
151                r#"{"usage":{"completion_tokens":100}}"#,
152            ),
153            (
154                Dialect::Claude,
155                r#"{"type":"content_block_delta","delta":{"text":"你好hmm{}"}}"#,
156                r#"{"type":"content_block_delta","delta":{"signature":"ignored"}}"#,
157            ),
158            (
159                Dialect::OpenAi,
160                r#"{"type":"response.output_text.delta","delta":"你好hmm{}"}"#,
161                r#"{"type":"response.completed","response":{"output":[{"type":"message","content":[{"type":"output_text","text":"你好hmm{}"}]}]}}"#,
162            ),
163            (
164                Dialect::Gemini,
165                r#"{"candidates":[{"content":{"parts":[{"text":"你好hmm{}","thoughtSignature":"ignored"}]}}]}"#,
166                r#"{"candidates":[{"content":{"parts":[{"inlineData":{"data":"ignored"}}]}}]}"#,
167            ),
168        ] {
169            for framing in [
170                crate::connection::StreamFraming::Sse,
171                crate::connection::StreamFraming::JsonArray,
172            ] {
173                if framing == crate::connection::StreamFraming::JsonArray
174                    && dialect != Dialect::Gemini
175                {
176                    continue;
177                }
178                let mut reader = UsageReader::new(
179                    Operation::StreamGenerateContent,
180                    dialect,
181                    UsageTransport::Http {
182                        framing: Some(framing),
183                    },
184                )
185                .unwrap();
186                reader.count_output_characters();
187                let wire = if framing == crate::connection::StreamFraming::Sse {
188                    format!(
189                        ": {}\n\ndata: {event}\n\ndata: {replay}\n\ndata: [DONE]\n\n",
190                        "keepalive".repeat(1024)
191                    )
192                } else {
193                    format!("[{event},{replay}]")
194                };
195                for chunk in wire.as_bytes().chunks(3) {
196                    reader.push(chunk);
197                }
198                assert_eq!(reader.output_characters(), 7, "{dialect:?} {framing:?}");
199            }
200        }
201    }
202
203    #[test]
204    fn whole_responses_count_text_thoughts_and_arguments_only() {
205        for (dialect, body) in [
206            (
207                Dialect::OpenAiChat,
208                r#"{"choices":[{"message":{"content":"你好","reasoning":"hmm","tool_calls":[{"function":{"arguments":"{}"}}],"audio":{"data":"ignored"}}}]}"#,
209            ),
210            (
211                Dialect::Claude,
212                r#"{"content":[{"type":"text","text":"你好"},{"type":"thinking","thinking":"hmm","signature":"ignored"},{"type":"tool_use","input":"{}"}]}"#,
213            ),
214            (
215                Dialect::OpenAi,
216                r#"{"output":[{"type":"message","content":[{"type":"output_text","text":"你好"}]},{"type":"reasoning","summary":[{"text":"hmm"}],"encrypted_content":"ignored"},{"type":"function_call","arguments":"{}"}]}"#,
217            ),
218            (
219                Dialect::Gemini,
220                r#"{"candidates":[{"content":{"parts":[{"text":"你好"},{"text":"hmm","thought":true},{"functionCall":{"args":"{}"}},{"inlineData":{"data":"ignored"}}]}}]}"#,
221            ),
222        ] {
223            assert_eq!(
224                whole_output_characters(Operation::GenerateContent, dialect, body.as_bytes()),
225                7
226            );
227        }
228        assert_eq!(
229            whole_output_characters(Operation::CreateSpeech, Dialect::OpenAi, b"audio bytes"),
230            0
231        );
232        assert_eq!(
233            whole_output_characters(
234                Operation::GenerateContent,
235                Dialect::OpenAi,
236                br#"{"error":{"message":"bad request"}}"#
237            ),
238            0
239        );
240    }
241}
gproxy_protocol/usage/output.rs

gproxy_protocol/usage/
output.rs
