Skip to main content

gproxy_protocol/transform/generate/gemini_responses/
usage.rs

1use crate::{
2    transform::{Report, TransformError},
3    wire::{gemini as g, openai::responses::response as r},
4};
5
6#[derive(Debug, Default, Clone, Copy)]
7pub struct GeminiUsageFacts {
8    pub cache_write_tokens: Option<i64>,
9    pub cached_tokens: Option<i64>,
10}
11
12fn nonnegative(n: i64) -> Result<i64, TransformError> {
13    Ok(n)
14}
15
16fn count(v: Option<i64>, field: &str) -> Result<i64, TransformError> {
17    nonnegative(v.ok_or_else(|| TransformError::missing_metadata(field))?)
18}
19
20fn add(a: i64, b: i64) -> Result<i64, TransformError> {
21    nonnegative(a)?
22        .checked_add(nonnegative(b)?)
23        .ok_or_else(|| TransformError::invalid_result("usage", "token overflow"))
24}
25
26pub(super) fn to_responses(
27    input: g::UsageMetadata,
28    facts: GeminiUsageFacts,
29    report: &mut Report,
30) -> Result<r::ResponseUsage, TransformError> {
31    let prompt = add(
32        count(input.prompt_token_count, "prompt_token_count")?,
33        input.tool_use_prompt_token_count.unwrap_or(0),
34    )?;
35    let visible = count(input.candidates_token_count, "candidates_token_count")?;
36    let total = count(input.total_token_count, "total_token_count")?;
37    let output = total.checked_sub(prompt).ok_or_else(|| {
38        TransformError::invalid_result("usage", "inconsistent input/output total")
39    })?;
40    let thinking = output - visible;
41
42    let cached = count(
43        input.cached_content_token_count.or(facts.cached_tokens),
44        "cached_content_token_count actual fact",
45    )?;
46    let written = count(facts.cache_write_tokens, "cache_write_tokens actual fact")?;
47
48    if input.prompt_tokens_details.is_some()
49        || input.cache_tokens_details.is_some()
50        || input.candidates_tokens_details.is_some()
51        || input.tool_use_prompt_tokens_details.is_some()
52    {
53        report.omitted(
54            "usage.modality_details",
55            "Responses has no Gemini modality breakdown",
56        );
57    }
58    Ok(r::ResponseUsage::builder(
59        prompt,
60        r::ResponseInputTokensDetails::builder(written, cached).build(),
61        output,
62        r::ResponseOutputTokensDetails::builder(thinking).build(),
63        total,
64    )
65    .build())
66}
67
68pub(super) fn to_gemini(
69    input: r::ResponseUsage,
70    report: &mut Report,
71) -> Result<g::UsageMetadata, TransformError> {
72    let thinking = nonnegative(input.output_tokens_details.reasoning_tokens)?;
73    let visible = input.output_tokens.checked_sub(thinking).ok_or_else(|| {
74        TransformError::invalid_result("usage.reasoning", "reasoning exceeds output")
75    })?;
76
77    let mut out = g::UsageMetadata::builder().build();
78    out.prompt_token_count = Some(input.input_tokens);
79    out.candidates_token_count = Some(visible);
80    out.thoughts_token_count = Some(thinking);
81    out.total_token_count = Some(input.total_tokens);
82    out.cached_content_token_count = Some(input.input_tokens_details.cached_tokens);
83    report.omitted(
84        "usage.cache_write_tokens",
85        "Gemini has no distinct cache-write counter",
86    );
87    Ok(out)
88}