gproxy_protocol/transform/count_tokens/
response.rs

1use crate::{
2    transform::{Converted, Report, TransformError},
3    wire::{claude::count_tokens as c, gemini::count_tokens as g, openai::count_tokens as o},
4};
5
6/// Facts from the original Claude counting request and any context-edit pass.
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum ClaudeCountContext {
9    /// The counted input is identical to the original input: no edits occurred.
10    Unedited,
11    /// An actual count obtained before applying context management.
12    OriginalInputTokens(i64),
13    /// Edits were requested, but their original token count is unavailable.
14    Unknown,
15}
16
17impl ClaudeCountContext {
18    pub fn for_request(request: &c::CountTokensRequestBody) -> Self {
19        if request
20            .context_management
21            .as_ref()
22            .and_then(|config| config.edits.as_ref())
23            .is_none_or(Vec::is_empty)
24        {
25            Self::Unedited
26        } else {
27            Self::Unknown
28        }
29    }
30
31    fn original(self, counted: i64) -> Result<i64, TransformError> {
32        match self {
33            Self::Unedited => Ok(counted),
34            Self::OriginalInputTokens(value) => {
35                nonnegative(value, "context_management.original_input_tokens")
36            }
37            Self::Unknown => Err(TransformError::missing_metadata(
38                "context_management.original_input_tokens",
39            )),
40        }
41    }
42}
43
44pub fn openai_to_claude_response(
45    input: o::CountTokensResponseBody,
46    context: ClaudeCountContext,
47) -> Result<Converted<c::CountTokensResponseBody>, TransformError> {
48    claude(input.input_tokens, context, Report::default())
49}
50
51pub fn gemini_to_claude_response(
52    input: g::CountTokensResponseBody,
53    context: ClaudeCountContext,
54) -> Result<Converted<c::CountTokensResponseBody>, TransformError> {
55    let (count, report) = gemini_count(&input)?;
56    claude(count, context, report)
57}
58
59pub fn claude_to_openai_response(
60    input: c::CountTokensResponseBody,
61) -> Result<Converted<o::CountTokensResponseBody>, TransformError> {
62    let report = claude_report(&input)?;
63    openai(input.input_tokens, report)
64}
65
66pub fn gemini_to_openai_response(
67    input: g::CountTokensResponseBody,
68) -> Result<Converted<o::CountTokensResponseBody>, TransformError> {
69    let (count, report) = gemini_count(&input)?;
70    openai(count, report)
71}
72
73pub fn openai_to_gemini_response(
74    input: o::CountTokensResponseBody,
75) -> Result<Converted<g::CountTokensResponseBody>, TransformError> {
76    gemini(input.input_tokens, Report::default())
77}
78
79pub fn claude_to_gemini_response(
80    input: c::CountTokensResponseBody,
81) -> Result<Converted<g::CountTokensResponseBody>, TransformError> {
82    let report = claude_report(&input)?;
83    gemini(input.input_tokens, report)
84}
85
86fn claude(
87    count: i64,
88    context: ClaudeCountContext,
89    report: Report,
90) -> Result<Converted<c::CountTokensResponseBody>, TransformError> {
91    let count = nonnegative(count, "input_tokens")?;
92    let original = context.original(count)?;
93    Ok(Converted {
94        value: c::CountTokensResponseBody::builder(
95            c::CountTokensContextManagementResponse::builder(original).build(),
96            count,
97        )
98        .build(),
99        report,
100    })
101}
102
103fn openai(
104    count: i64,
105    report: Report,
106) -> Result<Converted<o::CountTokensResponseBody>, TransformError> {
107    Ok(Converted {
108        value: o::CountTokensResponseBody::builder(
109            nonnegative(count, "input_tokens")?,
110            o::CountTokensObject::ResponseInputTokens,
111        )
112        .build(),
113        report,
114    })
115}
116
117fn gemini(
118    count: i64,
119    report: Report,
120) -> Result<Converted<g::CountTokensResponseBody>, TransformError> {
121    Ok(Converted {
122        value: g::CountTokensResponseBody::builder()
123            .total_tokens(nonnegative(count, "input_tokens")?)
124            .build(),
125        report,
126    })
127}
128
129fn claude_report(input: &c::CountTokensResponseBody) -> Result<Report, TransformError> {
130    nonnegative(
131        input.context_management.original_input_tokens,
132        "context_management.original_input_tokens",
133    )?;
134    let mut report = Report::default();
135    report.omitted(
136        "context_management.original_input_tokens",
137        "target reports only the counted input, without the count before context edits",
138    );
139    Ok(report)
140}
141
142fn gemini_count(input: &g::CountTokensResponseBody) -> Result<(i64, Report), TransformError> {
143    let count = nonnegative(
144        input
145            .total_tokens
146            .ok_or_else(|| TransformError::missing_metadata("totalTokens"))?,
147        "totalTokens",
148    )?;
149    let mut report = Report::default();
150    if let Some(cached) = input.cached_content_token_count {
151        nonnegative(cached, "cachedContentTokenCount")?;
152        if cached > count {
153            return Err(TransformError::invalid_result(
154                "cachedContentTokenCount",
155                "cached tokens exceed total input tokens",
156            ));
157        }
158        report.omitted(
159            "cachedContentTokenCount",
160            "target count result has no cached-token detail",
161        );
162    }
163    for (field, values) in [
164        ("promptTokensDetails", &input.prompt_tokens_details),
165        ("cacheTokensDetails", &input.cache_tokens_details),
166    ] {
167        if let Some(values) = values {
168            for (index, value) in values.iter().enumerate() {
169                if let Some(count) = value.token_count {
170                    nonnegative(count, &format!("{field}[{index}].tokenCount"))?;
171                }
172            }
173            if !values.is_empty() {
174                report.omitted(field, "target count result has no modality-token detail");
175            }
176        }
177    }
178    Ok((count, report))
179}
180
181fn nonnegative(count: i64, field: &str) -> Result<i64, TransformError> {
182    if count < 0 {
183        return Err(TransformError::invalid_result(
184            field,
185            "token count cannot be negative",
186        ));
187    }
188    Ok(count)
189}
gproxy_protocol/transform/count_tokens/response.rs

gproxy_protocol/transform/count_tokens/
response.rs
