Skip to main content

gproxy_protocol/transform/generate/claude_gemini/
usage.rs

1use crate::{
2    transform::{Report, TransformError},
3    wire::{claude::generate_content as c, gemini as g},
4};
5
6/// Measured counters absent from the source response. Supplied counters must
7/// agree with any corresponding upstream fields; omission is not zero.
8#[derive(Debug, Default, Clone, Copy)]
9pub struct ClaudeGeminiUsageFacts {
10    pub cache_creation_input_tokens: Option<i64>,
11    pub cache_read_input_tokens: Option<i64>,
12    pub thinking_tokens: Option<i64>,
13}
14
15fn count(v: i64) -> Result<i64, TransformError> {
16    Ok(v)
17}
18
19fn sum(a: i64, b: i64) -> Result<i64, TransformError> {
20    count(a)?
21        .checked_add(count(b)?)
22        .ok_or_else(|| TransformError::invalid_result("usage", "count overflow"))
23}
24
25fn actual(a: Option<i64>, b: Option<i64>, field: &str) -> Result<i64, TransformError> {
26    count(
27        a.or(b)
28            .ok_or_else(|| TransformError::missing_metadata(field))?,
29    )
30}
31
32pub(super) fn to_gemini(
33    input: c::Usage,
34    facts: ClaudeGeminiUsageFacts,
35    report: &mut Report,
36) -> Result<g::UsageMetadata, TransformError> {
37    let breakdown = input
38        .cache_creation
39        .flatten()
40        .map(|v| sum(v.ephemeral_1h_input_tokens, v.ephemeral_5m_input_tokens))
41        .map(crate::transform::optional)
42        .transpose()?
43        .flatten();
44    let written = actual(
45        input.cache_creation_input_tokens.flatten(),
46        breakdown,
47        "usage.cache_creation_input_tokens",
48    )
49    .or_else(|e| {
50        if input.cache_creation_input_tokens.flatten().is_none() && breakdown.is_none() {
51            actual(
52                None,
53                facts.cache_creation_input_tokens,
54                "usage.cache_creation_input_tokens",
55            )
56        } else {
57            Err(e)
58        }
59    })?;
60
61    let cached = actual(
62        input.cache_read_input_tokens.flatten(),
63        facts.cache_read_input_tokens,
64        "usage.cache_read_input_tokens",
65    )?;
66    let thinking = actual(
67        input
68            .output_tokens_details
69            .flatten()
70            .map(|v| v.thinking_tokens),
71        facts.thinking_tokens,
72        "usage.thinking_tokens",
73    )?;
74    let input_count = sum(sum(input.input_tokens, written)?, cached)?;
75    let output = count(input.output_tokens)?;
76
77    for (present, field) in [
78        (breakdown.is_some(), "cache_creation"),
79        (input.fallback_credit.is_some(), "fallback_credit"),
80        (input.inference_geo.is_some(), "inference_geo"),
81        (input.iterations.is_some(), "iterations"),
82        (input.server_tool_use.is_some(), "server_tool_use"),
83        (input.service_tier.is_some(), "service_tier"),
84        (input.speed.is_some(), "speed"),
85    ] {
86        if present {
87            report.omitted(
88                format!("usage.{field}"),
89                "Gemini has no equivalent Claude usage detail",
90            );
91        }
92    }
93    Ok(g::UsageMetadata::builder()
94        .prompt_token_count(input_count)
95        .cached_content_token_count(cached)
96        .candidates_token_count(output - thinking)
97        .thoughts_token_count(thinking)
98        .total_token_count(sum(input_count, output)?)
99        .build())
100}
101
102pub(super) fn to_claude(
103    input: g::UsageMetadata,
104    facts: ClaudeGeminiUsageFacts,
105    report: &mut Report,
106) -> Result<c::Usage, TransformError> {
107    let prompt = actual(input.prompt_token_count, None, "usage.prompt_token_count")?;
108    let prompt = sum(prompt, input.tool_use_prompt_token_count.unwrap_or(0))?;
109    let candidates = actual(
110        input.candidates_token_count,
111        None,
112        "usage.candidates_token_count",
113    )?;
114    let total = actual(input.total_token_count, None, "usage.total_token_count")?;
115    let output = total
116        .checked_sub(prompt)
117        .ok_or_else(|| TransformError::invalid_result("usage", "total smaller than components"))?;
118    let thinking = output - candidates;
119    actual(
120        input.thoughts_token_count,
121        Some(thinking),
122        "usage.thoughts_token_count",
123    )?;
124
125    let cached = actual(
126        input.cached_content_token_count,
127        facts.cache_read_input_tokens,
128        "usage.cached_content_token_count",
129    )?;
130    let written = actual(
131        None,
132        facts.cache_creation_input_tokens,
133        "usage.cache_creation_input_tokens",
134    )?;
135    let uncached = prompt
136        .checked_sub(sum(cached, written)?)
137        .ok_or_else(|| TransformError::invalid_result("usage.cache", "cache exceeds input"))?;
138    for (present, field) in [
139        (
140            input.prompt_tokens_details.is_some(),
141            "prompt_tokens_details",
142        ),
143        (input.cache_tokens_details.is_some(), "cache_tokens_details"),
144        (
145            input.candidates_tokens_details.is_some(),
146            "candidates_tokens_details",
147        ),
148        (
149            input.tool_use_prompt_tokens_details.is_some(),
150            "tool_use_prompt_tokens_details",
151        ),
152        (input.service_tier.is_some(), "service_tier"),
153    ] {
154        if present {
155            report.omitted(
156                format!("usage.{field}"),
157                "Claude lacks equivalent Gemini usage details",
158            );
159        }
160    }
161    let mut out = c::Usage::builder(uncached, output).build();
162    out.cache_creation_input_tokens = Some(Some(written));
163    out.cache_read_input_tokens = Some(Some(cached));
164    out.output_tokens_details = Some(Some(c::OutputTokensDetails::builder(thinking).build()));
165    Ok(out)
166}