i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(usage): include cache token details in stream usage payloads

Browse Source
This commit is contained in:
ZheFox
2026-05-20 21:20:10 +08:00
parent 64ad0f694b
commit 3c6924238f
6 changed files with 310 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
crates/aether-ai-formats/src/formats/claude/messages/stream.rs
View File

@@ -753,13 +753,7 @@ impl ClaudeClientEmitter {
}),
);
let usage = usage.unwrap_or_default();
payload.insert(
"usage".to_string(),
json!({
"input_tokens": usage.input_tokens,
"output_tokens": usage.output_tokens,
}),
);
payload.insert("usage".to_string(), claude_usage_from_usage(&usage));
out.extend(encode_json_sse(
Some("message_delta"),
&Value::Object(payload),
@@ -1289,6 +1283,32 @@ mod tests {
assert!(sse.contains("\"usage\":{\"input_tokens\":0,\"output_tokens\":0}"));
}
#[test]
fn claude_client_emitter_includes_cache_usage_in_finish_events() {
let mut emitter = ClaudeClientEmitter::default();
let bytes = emitter
.emit(CanonicalStreamFrame {
id: "msg_cache".to_string(),
model: "claude-sonnet-4-5".to_string(),
event: CanonicalStreamEvent::Finish {
finish_reason: Some("stop".to_string()),
usage: Some(CanonicalUsage {
input_tokens: 10,
output_tokens: 2,
total_tokens: 12,
cache_creation_tokens: 5,
cache_read_tokens: 4,
..CanonicalUsage::default()
}),
},
})
.expect("finish should encode");
let sse = String::from_utf8(bytes).expect("sse should be utf8");
assert!(sse.contains("\"cache_creation_input_tokens\":5"));
assert!(sse.contains("\"cache_read_input_tokens\":4"));
}
#[test]
fn claude_client_emitter_emits_image_blocks_for_media_parts() {
let mut emitter = ClaudeClientEmitter::default();

25
crates/aether-ai-formats/src/formats/gemini/generate_content/stream.rs
View File

@@ -404,27 +404,10 @@ impl GeminiClientEmitter {
Value::Array(vec![Value::Object(candidate)]),
);
if let Some(usage) = usage {
let visible_output_tokens = usage.output_tokens.saturating_sub(usage.reasoning_tokens);
let mut usage_metadata = Map::new();
usage_metadata.insert(
"promptTokenCount".to_string(),
Value::from(usage.input_tokens),
response.insert(
"usageMetadata".to_string(),
gemini_usage_metadata_from_usage(&usage),
);
usage_metadata.insert(
"candidatesTokenCount".to_string(),
Value::from(visible_output_tokens),
);
usage_metadata.insert(
"totalTokenCount".to_string(),
Value::from(usage.total_tokens),
);
if usage.reasoning_tokens > 0 {
usage_metadata.insert(
"thoughtsTokenCount".to_string(),
Value::from(usage.reasoning_tokens),
);
}
response.insert("usageMetadata".to_string(), Value::Object(usage_metadata));
}
encode_json_sse(None, &Value::Object(response))
}
@@ -1000,6 +983,7 @@ mod tests {
output_tokens: 3,
reasoning_tokens: 1,
total_tokens: 4,
cache_read_tokens: 5,
..CanonicalUsage::default()
}),
},
@@ -1012,6 +996,7 @@ mod tests {
assert!(sse.contains("\"thoughtSignature\":\"sig_123\""));
assert!(sse.contains("\"thoughtsTokenCount\":1"));
assert!(sse.contains("\"candidatesTokenCount\":2"));
assert!(sse.contains("\"cachedContentTokenCount\":5"));
assert!(sse.contains("\"finishReason\":\"STOP\""));
}

32
crates/aether-ai-formats/src/formats/openai/chat/stream.rs
View File

@@ -1566,15 +1566,12 @@ impl OpenAIChatClientEmitter {
if let Some(usage) = usage {
out.extend(encode_json_sse(
None,
&build_openai_chat_usage_chunk(
&build_openai_chat_usage_chunk_from_usage(
self.response_id
.as_deref()
.unwrap_or("chatcmpl-local-stream"),
self.model.as_deref().unwrap_or("unknown"),
usage.input_tokens,
usage.output_tokens,
usage.total_tokens,
usage.reasoning_tokens,
&usage,
),
)?);
}
@@ -2202,20 +2199,6 @@ impl OpenAIResponsesClientEmitter {
}
ordered_output.sort_by_key(|(output_index, _)| *output_index);
let mut usage_payload = Map::new();
usage_payload.insert("input_tokens".to_string(), Value::from(usage.input_tokens));
usage_payload.insert(
"output_tokens".to_string(),
Value::from(usage.output_tokens),
);
usage_payload.insert("total_tokens".to_string(), Value::from(usage.total_tokens));
if usage.reasoning_tokens > 0 {
usage_payload.insert(
"output_tokens_details".to_string(),
json!({ "reasoning_tokens": usage.reasoning_tokens }),
);
}
json!({
"id": self.response_id(),
"object": "response",
@@ -2225,7 +2208,7 @@ impl OpenAIResponsesClientEmitter {
.into_iter()
.map(|(_, item)| item)
.collect::<Vec<_>>(),
"usage": usage_payload,
"usage": openai_responses_usage_from_usage(&usage),
})
}
@@ -3297,6 +3280,8 @@ mod tests {
input_tokens: 1,
output_tokens: 2,
total_tokens: 3,
cache_creation_tokens: 5,
cache_read_tokens: 4,
reasoning_tokens: 1,
..CanonicalUsage::default()
}),
@@ -3311,6 +3296,8 @@ mod tests {
assert!(sse.contains("\"prompt_tokens\":1"));
assert!(sse.contains("\"completion_tokens\":2"));
assert!(sse.contains("\"completion_tokens_details\":{\"reasoning_tokens\":1}"));
assert!(sse.contains("\"cached_creation_tokens\":5"));
assert!(sse.contains("\"cached_tokens\":4"));
assert!(sse.contains("\"total_tokens\":3"));
assert!(sse.contains("data: [DONE]\n\n"));
}
@@ -3418,6 +3405,8 @@ mod tests {
input_tokens: 1,
output_tokens: 2,
total_tokens: 3,
cache_creation_tokens: 5,
cache_read_tokens: 4,
reasoning_tokens: 1,
..CanonicalUsage::default()
}),
@@ -3430,6 +3419,9 @@ mod tests {
assert!(sse.contains("\"type\":\"reasoning\""));
assert!(sse.contains("\"text\":\"because\""));
assert!(sse.contains("\"output_tokens_details\":{\"reasoning_tokens\":1}"));
assert!(sse.contains("\"input_tokens_details\""));
assert!(sse.contains("\"cached_creation_tokens\":5"));
assert!(sse.contains("\"cached_tokens\":4"));
}
#[test]

52
crates/aether-ai-formats/src/formats/openai/image/stream.rs
View File

@@ -8,7 +8,8 @@ use crate::contracts::OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND;
use crate::formats::openai::responses::codex::CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT;
use crate::formats::shared::sse::{encode_done_sse, encode_json_sse};
use crate::formats::shared::stream_core::common::{
build_openai_chat_chunk, build_openai_chat_finish_chunk, build_openai_chat_usage_chunk,
build_openai_chat_chunk, build_openai_chat_finish_chunk,
build_openai_chat_usage_chunk_with_cache,
};
use crate::formats::shared::AiSurfaceFinalizeError;
@@ -501,18 +502,26 @@ impl OpenAiImageChatStreamState {
None,
&build_openai_chat_finish_chunk(&response_id, &model, Some("stop")),
)?);
if let Some((input_tokens, output_tokens, total_tokens, reasoning_tokens)) =
openai_image_chat_usage_counts(usage)
if let Some((
input_tokens,
output_tokens,
total_tokens,
reasoning_tokens,
cache_creation_tokens,
cache_read_tokens,
)) = openai_image_chat_usage_counts(usage)
{
output.extend(encode_json_sse(
None,
&build_openai_chat_usage_chunk(
&build_openai_chat_usage_chunk_with_cache(
&response_id,
&model,
input_tokens,
output_tokens,
total_tokens,
reasoning_tokens,
cache_creation_tokens,
cache_read_tokens,
),
)?);
}
@@ -900,7 +909,7 @@ fn image_chat_markdown(frame: &OpenAiImageChatFrame) -> String {
)
}
fn openai_image_chat_usage_counts(usage: Option<&Value>) -> Option<(u64, u64, u64, u64)> {
fn openai_image_chat_usage_counts(usage: Option<&Value>) -> Option<(u64, u64, u64, u64, u64, u64)> {
let usage = usage.and_then(Value::as_object)?;
let mut input_tokens = usage
.get("input_tokens")
@@ -912,6 +921,30 @@ fn openai_image_chat_usage_counts(usage: Option<&Value>) -> Option<(u64, u64, u6
.or_else(|| usage.get("completion_tokens"))
.and_then(Value::as_u64)
.unwrap_or(0);
let cache_creation_tokens = usage
.get("cache_creation_input_tokens")
.and_then(Value::as_u64)
.or_else(|| {
usage
.get("input_tokens_details")
.or_else(|| usage.get("prompt_tokens_details"))
.and_then(Value::as_object)
.and_then(|details| details.get("cached_creation_tokens"))
.and_then(Value::as_u64)
})
.unwrap_or(0);
let cache_read_tokens = usage
.get("cache_read_input_tokens")
.and_then(Value::as_u64)
.or_else(|| {
usage
.get("input_tokens_details")
.or_else(|| usage.get("prompt_tokens_details"))
.and_then(Value::as_object)
.and_then(|details| details.get("cached_tokens"))
.and_then(Value::as_u64)
})
.unwrap_or(0);
let total_tokens = usage
.get("total_tokens")
.and_then(Value::as_u64)
@@ -919,7 +952,14 @@ fn openai_image_chat_usage_counts(usage: Option<&Value>) -> Option<(u64, u64, u6
if input_tokens == 0 && total_tokens > output_tokens {
input_tokens = total_tokens.saturating_sub(output_tokens);
}
(total_tokens > 0).then_some((input_tokens, output_tokens, total_tokens, 0))
(total_tokens > 0).then_some((
input_tokens,
output_tokens,
total_tokens,
0,
cache_creation_tokens,
cache_read_tokens,
))
}
fn image_failure_error(event: &Value) -> Value {

196
crates/aether-ai-formats/src/formats/shared/stream_core/common.rs
View File

@@ -344,6 +344,156 @@ pub fn build_openai_chat_usage_chunk(
total_tokens: u64,
reasoning_tokens: u64,
) -> Value {
build_openai_chat_usage_chunk_with_cache(
id,
model,
prompt_tokens,
completion_tokens,
total_tokens,
reasoning_tokens,
0,
0,
)
}
pub fn build_openai_chat_usage_chunk_with_cache(
id: &str,
model: &str,
prompt_tokens: u64,
completion_tokens: u64,
total_tokens: u64,
reasoning_tokens: u64,
cache_creation_tokens: u64,
cache_read_tokens: u64,
) -> Value {
let usage = openai_chat_usage_payload(
prompt_tokens,
completion_tokens,
total_tokens,
reasoning_tokens,
cache_creation_tokens,
cache_read_tokens,
);
json!({
"id": id,
"object": "chat.completion.chunk",
"model": model,
"choices": [],
"usage": usage,
})
}
pub fn build_openai_chat_usage_chunk_from_usage(
id: &str,
model: &str,
usage: &CanonicalUsage,
) -> Value {
build_openai_chat_usage_chunk_with_cache(
id,
model,
usage.input_tokens,
usage.output_tokens,
usage.total_tokens,
usage.reasoning_tokens,
cache_creation_tokens_for_usage(usage),
usage.cache_read_tokens,
)
}
pub fn openai_responses_usage_from_usage(usage: &CanonicalUsage) -> Value {
let mut output = Map::new();
output.insert("input_tokens".to_string(), Value::from(usage.input_tokens));
output.insert(
"output_tokens".to_string(),
Value::from(usage.output_tokens),
);
output.insert("total_tokens".to_string(), Value::from(usage.total_tokens));
if usage.reasoning_tokens > 0 {
output.insert(
"output_tokens_details".to_string(),
json!({ "reasoning_tokens": usage.reasoning_tokens }),
);
}
insert_openai_token_details(
&mut output,
"input_tokens_details",
cache_creation_tokens_for_usage(usage),
usage.cache_read_tokens,
);
Value::Object(output)
}
pub fn claude_usage_from_usage(usage: &CanonicalUsage) -> Value {
let mut output = Map::new();
output.insert("input_tokens".to_string(), Value::from(usage.input_tokens));
output.insert(
"output_tokens".to_string(),
Value::from(usage.output_tokens),
);
if usage.cache_read_tokens > 0 {
output.insert(
"cache_read_input_tokens".to_string(),
Value::from(usage.cache_read_tokens),
);
}
let cache_creation_tokens = cache_creation_tokens_for_usage(usage);
if cache_creation_tokens > 0 {
output.insert(
"cache_creation_input_tokens".to_string(),
Value::from(cache_creation_tokens),
);
}
if usage.cache_creation_ephemeral_5m_tokens > 0 || usage.cache_creation_ephemeral_1h_tokens > 0
{
output.insert(
"cache_creation".to_string(),
json!({
"ephemeral_5m_input_tokens": usage.cache_creation_ephemeral_5m_tokens,
"ephemeral_1h_input_tokens": usage.cache_creation_ephemeral_1h_tokens,
}),
);
}
Value::Object(output)
}
pub fn gemini_usage_metadata_from_usage(usage: &CanonicalUsage) -> Value {
let visible_output_tokens = usage.output_tokens.saturating_sub(usage.reasoning_tokens);
let mut output = Map::new();
output.insert(
"promptTokenCount".to_string(),
Value::from(usage.input_tokens),
);
output.insert(
"candidatesTokenCount".to_string(),
Value::from(visible_output_tokens),
);
output.insert(
"totalTokenCount".to_string(),
Value::from(usage.total_tokens),
);
if usage.reasoning_tokens > 0 {
output.insert(
"thoughtsTokenCount".to_string(),
Value::from(usage.reasoning_tokens),
);
}
if usage.cache_read_tokens > 0 {
output.insert(
"cachedContentTokenCount".to_string(),
Value::from(usage.cache_read_tokens),
);
}
Value::Object(output)
}
fn openai_chat_usage_payload(
prompt_tokens: u64,
completion_tokens: u64,
total_tokens: u64,
reasoning_tokens: u64,
cache_creation_tokens: u64,
cache_read_tokens: u64,
) -> Map<String, Value> {
let mut usage = Map::new();
usage.insert("prompt_tokens".to_string(), Value::from(prompt_tokens));
usage.insert(
@@ -357,11 +507,43 @@ pub fn build_openai_chat_usage_chunk(
json!({ "reasoning_tokens": reasoning_tokens }),
);
}
json!({
"id": id,
"object": "chat.completion.chunk",
"model": model,
"choices": [],
"usage": usage,
})
insert_openai_token_details(
&mut usage,
"prompt_tokens_details",
cache_creation_tokens,
cache_read_tokens,
);
usage
}
fn insert_openai_token_details(
output: &mut Map<String, Value>,
details_key: &str,
cache_creation_tokens: u64,
cache_read_tokens: u64,
) {
if cache_creation_tokens == 0 && cache_read_tokens == 0 {
return;
}
let mut details = Map::new();
if cache_read_tokens > 0 {
details.insert("cached_tokens".to_string(), Value::from(cache_read_tokens));
}
if cache_creation_tokens > 0 {
details.insert(
"cached_creation_tokens".to_string(),
Value::from(cache_creation_tokens),
);
}
output.insert(details_key.to_string(), Value::Object(details));
}
fn cache_creation_tokens_for_usage(usage: &CanonicalUsage) -> u64 {
if usage.cache_creation_tokens > 0 {
usage.cache_creation_tokens
} else {
usage
.cache_creation_ephemeral_5m_tokens
.saturating_add(usage.cache_creation_ephemeral_1h_tokens)
}
}

37
crates/aether-ai-formats/src/formats/shared/sync_to_stream.rs
View File

@@ -14,7 +14,8 @@ use crate::formats::openai::chat::stream::{
};
use crate::formats::shared::sse::{encode_done_sse, encode_json_sse};
use crate::formats::shared::stream_core::common::{
build_openai_chat_chunk, build_openai_chat_finish_chunk, build_openai_chat_usage_chunk,
build_openai_chat_chunk, build_openai_chat_finish_chunk,
build_openai_chat_usage_chunk_with_cache,
};
use crate::formats::shared::stream_core::{
CanonicalStreamFrame, StreamingStandardFormatMatrix, StreamingStandardTerminalObserver,
@@ -171,20 +172,29 @@ fn maybe_bridge_openai_image_sync_json_to_chat_stream(
None,
&build_openai_chat_finish_chunk(&response_id, &model, Some("stop")),
)?);
if let Some((input_tokens, output_tokens, total_tokens, reasoning_tokens)) = summary
if let Some((
input_tokens,
output_tokens,
total_tokens,
reasoning_tokens,
cache_creation_tokens,
cache_read_tokens,
)) = summary
.standardized_usage
.as_ref()
.and_then(openai_chat_usage_counts)
{
sse_body.extend(encode_json_sse(
None,
&build_openai_chat_usage_chunk(
&build_openai_chat_usage_chunk_with_cache(
&response_id,
&model,
input_tokens,
output_tokens,
total_tokens,
reasoning_tokens,
cache_creation_tokens,
cache_read_tokens,
),
)?);
}
@@ -489,10 +499,12 @@ fn openai_image_standardized_usage(
(standardized_usage.signal_score() > 0).then_some(standardized_usage)
}
fn openai_chat_usage_counts(usage: &StandardizedUsage) -> Option<(u64, u64, u64, u64)> {
fn openai_chat_usage_counts(usage: &StandardizedUsage) -> Option<(u64, u64, u64, u64, u64, u64)> {
let input_tokens = usage.input_tokens.max(0) as u64;
let output_tokens = usage.output_tokens.max(0) as u64;
let reasoning_tokens = usage.reasoning_tokens.max(0) as u64;
let cache_creation_tokens = usage.cache_creation_tokens.max(0) as u64;
let cache_read_tokens = usage.cache_read_tokens.max(0) as u64;
let total_tokens = usage
.dimensions
.get("total_tokens")
@@ -502,7 +514,14 @@ fn openai_chat_usage_counts(usage: &StandardizedUsage) -> Option<(u64, u64, u64,
.saturating_add(output_tokens)
.saturating_add(reasoning_tokens)
});
(total_tokens > 0).then_some((input_tokens, output_tokens, total_tokens, reasoning_tokens))
(total_tokens > 0).then_some((
input_tokens,
output_tokens,
total_tokens,
reasoning_tokens,
cache_creation_tokens,
cache_read_tokens,
))
}
fn openai_image_bridge_response_id(
@@ -1207,7 +1226,11 @@ mod tests {
"usage": {
"total_tokens": 100,
"input_tokens": 50,
"output_tokens": 50
"output_tokens": 50,
"input_tokens_details": {
"cached_tokens": 20,
"cached_creation_tokens": 10
}
}
}),
"openai:image",
@@ -1222,6 +1245,8 @@ mod tests {
assert!(output.contains("![generated image](data:image/png;base64,aGVsbG8=)"));
assert!(output.contains("![generated image 2](data:image/png;base64,d29ybGQ=)"));
assert!(output.contains("\"finish_reason\":\"stop\""));
assert!(output.contains("\"cached_tokens\":20"));
assert!(output.contains("\"cached_creation_tokens\":10"));
assert!(output.contains("data: [DONE]"));
assert!(!output.contains("image_generation.completed"));