use serde_json::{Value, json};
use crate::ir::IrStreamEvent;
use super::u32_field;
pub(crate) fn from_chat(usage: &Value) -> IrStreamEvent {
let cache_read = nested_u32(usage, "prompt_tokens_details", "cached_tokens")
.filter(|&n| n > 0)
.or_else(|| u32_field(usage, "cached_tokens").filter(|&n| n > 0))
.or_else(|| u32_field(usage, "prompt_cache_hit_tokens"))
.unwrap_or(0);
let cache_write = nested_u32(usage, "prompt_tokens_details", "cache_write_tokens").unwrap_or(0);
let reasoning = nested_u32(usage, "completion_tokens_details", "reasoning_tokens").unwrap_or(0);
let audio_tokens = nested_u32(usage, "prompt_tokens_details", "audio_tokens").unwrap_or(0);
let completion_audio_tokens =
nested_u32(usage, "completion_tokens_details", "audio_tokens").unwrap_or(0);
IrStreamEvent::Usage {
prompt_tokens: u32_field(usage, "prompt_tokens")
.unwrap_or(0)
.saturating_sub(cache_read),
completion_tokens: u32_field(usage, "completion_tokens")
.unwrap_or(0)
.saturating_sub(reasoning),
cache_read_tokens: cache_read,
cache_write_tokens: cache_write,
reasoning_tokens: reasoning,
audio_tokens,
completion_audio_tokens,
}
}
pub(super) fn from_anthropic(usage: &Value) -> IrStreamEvent {
let reasoning = nested_u32(usage, "output_tokens_details", "thinking_tokens").unwrap_or(0);
IrStreamEvent::Usage {
prompt_tokens: u32_field(usage, "input_tokens").unwrap_or(0),
completion_tokens: u32_field(usage, "output_tokens")
.unwrap_or(0)
.saturating_sub(reasoning),
cache_read_tokens: u32_field(usage, "cache_read_input_tokens").unwrap_or(0),
cache_write_tokens: u32_field(usage, "cache_creation_input_tokens").unwrap_or(0),
reasoning_tokens: reasoning,
audio_tokens: 0,
completion_audio_tokens: 0,
}
}
pub(super) fn from_responses(usage: &Value) -> IrStreamEvent {
let cache_read = nested_u32(usage, "input_tokens_details", "cached_tokens").unwrap_or(0);
let cache_write = nested_u32(usage, "input_tokens_details", "cache_write_tokens").unwrap_or(0);
let reasoning = nested_u32(usage, "output_tokens_details", "reasoning_tokens").unwrap_or(0);
IrStreamEvent::Usage {
prompt_tokens: u32_field(usage, "input_tokens")
.unwrap_or(0)
.saturating_sub(cache_read),
completion_tokens: u32_field(usage, "output_tokens")
.unwrap_or(0)
.saturating_sub(reasoning),
cache_read_tokens: cache_read,
cache_write_tokens: cache_write,
reasoning_tokens: reasoning,
audio_tokens: 0,
completion_audio_tokens: 0,
}
}
pub(super) fn from_gemini(usage: &Value) -> IrStreamEvent {
let cache_read = u32_field(usage, "cachedContentTokenCount")
.or_else(|| u32_field(usage, "cached_content_token_count"))
.unwrap_or(0);
let reasoning = u32_field(usage, "thoughtsTokenCount")
.or_else(|| u32_field(usage, "thoughts_token_count"))
.unwrap_or(0);
let prompt = u32_field(usage, "promptTokenCount")
.or_else(|| u32_field(usage, "prompt_token_count"))
.unwrap_or(0);
let completion = u32_field(usage, "candidatesTokenCount")
.or_else(|| u32_field(usage, "candidates_token_count"))
.unwrap_or(0);
IrStreamEvent::Usage {
prompt_tokens: prompt.saturating_sub(cache_read),
completion_tokens: completion,
cache_read_tokens: cache_read,
cache_write_tokens: 0,
reasoning_tokens: reasoning,
audio_tokens: gemini_audio_tokens(usage, "promptTokensDetails", "prompt_tokens_details"),
completion_audio_tokens: gemini_audio_tokens(
usage,
"candidatesTokensDetails",
"candidates_tokens_details",
),
}
}
pub(super) fn encode_chat(
prompt_tokens: u32,
completion_tokens: u32,
cache_read_tokens: u32,
cache_write_tokens: u32,
reasoning_tokens: u32,
audio_tokens: u32,
completion_audio_tokens: u32,
) -> Value {
let prompt_wire = prompt_tokens.saturating_add(cache_read_tokens);
let completion_wire = completion_tokens.saturating_add(reasoning_tokens);
let mut usage = json!({
"prompt_tokens": prompt_wire,
"completion_tokens": completion_wire,
"total_tokens": prompt_wire.saturating_add(completion_wire),
});
if cache_read_tokens > 0 || cache_write_tokens > 0 || audio_tokens > 0 {
let mut details = json!({});
if cache_read_tokens > 0 {
details["cached_tokens"] = json!(cache_read_tokens);
}
if cache_write_tokens > 0 {
details["cache_write_tokens"] = json!(cache_write_tokens);
}
if audio_tokens > 0 {
details["audio_tokens"] = json!(audio_tokens);
}
usage["prompt_tokens_details"] = details;
}
if reasoning_tokens > 0 || completion_audio_tokens > 0 {
let mut details = json!({});
if reasoning_tokens > 0 {
details["reasoning_tokens"] = json!(reasoning_tokens);
}
if completion_audio_tokens > 0 {
details["audio_tokens"] = json!(completion_audio_tokens);
}
usage["completion_tokens_details"] = details;
}
json!({ "choices": [], "usage": usage })
}
pub(super) fn encode_anthropic(
prompt_tokens: u32,
completion_tokens: u32,
cache_read_tokens: u32,
cache_write_tokens: u32,
reasoning_tokens: u32,
) -> Value {
let mut usage = json!({
"input_tokens": prompt_tokens,
"output_tokens": completion_tokens.saturating_add(reasoning_tokens),
});
if cache_read_tokens > 0 {
usage["cache_read_input_tokens"] = json!(cache_read_tokens);
}
if cache_write_tokens > 0 {
usage["cache_creation_input_tokens"] = json!(cache_write_tokens);
}
if reasoning_tokens > 0 {
usage["output_tokens_details"] = json!({ "thinking_tokens": reasoning_tokens });
}
json!({
"type": "message_delta",
"delta": {},
"usage": usage,
})
}
pub(super) fn encode_responses(
prompt_tokens: u32,
completion_tokens: u32,
cache_read_tokens: u32,
cache_write_tokens: u32,
reasoning_tokens: u32,
) -> Value {
let input_tokens = prompt_tokens.saturating_add(cache_read_tokens);
let output_tokens = completion_tokens.saturating_add(reasoning_tokens);
let mut usage = json!({
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": input_tokens.saturating_add(output_tokens),
});
if cache_read_tokens > 0 || cache_write_tokens > 0 {
let mut details = json!({});
if cache_read_tokens > 0 {
details["cached_tokens"] = json!(cache_read_tokens);
}
if cache_write_tokens > 0 {
details["cache_write_tokens"] = json!(cache_write_tokens);
}
usage["input_tokens_details"] = details;
}
if reasoning_tokens > 0 {
usage["output_tokens_details"] = json!({ "reasoning_tokens": reasoning_tokens });
}
json!({
"type": "response.completed",
"response": { "status": "completed", "usage": usage },
})
}
pub(super) fn encode_gemini(
prompt_tokens: u32,
completion_tokens: u32,
cache_read_tokens: u32,
reasoning_tokens: u32,
audio_tokens: u32,
completion_audio_tokens: u32,
) -> Value {
let prompt_wire = prompt_tokens.saturating_add(cache_read_tokens);
let mut usage = json!({
"promptTokenCount": prompt_wire,
"candidatesTokenCount": completion_tokens,
"totalTokenCount": prompt_wire
.saturating_add(completion_tokens)
.saturating_add(reasoning_tokens),
});
if cache_read_tokens > 0 {
usage["cachedContentTokenCount"] = json!(cache_read_tokens);
}
if reasoning_tokens > 0 {
usage["thoughtsTokenCount"] = json!(reasoning_tokens);
}
if audio_tokens > 0 {
usage["promptTokensDetails"] = json!([{
"modality": "AUDIO",
"tokenCount": audio_tokens
}]);
}
if completion_audio_tokens > 0 {
usage["candidatesTokensDetails"] = json!([{
"modality": "AUDIO",
"tokenCount": completion_audio_tokens
}]);
}
json!({ "usageMetadata": usage })
}
fn gemini_audio_tokens(usage: &Value, camel: &str, snake: &str) -> u32 {
let details = usage
.get(camel)
.or_else(|| usage.get(snake))
.and_then(Value::as_array);
let Some(details) = details else {
return 0;
};
let mut total = 0u32;
for item in details {
let modality = item.get("modality").and_then(Value::as_str).unwrap_or("");
if !modality.eq_ignore_ascii_case("AUDIO") {
continue;
}
if let Some(n) = u32_field(item, "tokenCount").or_else(|| u32_field(item, "token_count")) {
total = total.saturating_add(n);
}
}
total
}
fn nested_u32(value: &Value, object: &str, key: &str) -> Option<u32> {
u32_field(value.get(object)?, key)
}