i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-13 12:39:50 +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
crates/aether-ai-formats/src/formats/shared/stream_core/common.rs
+189 -7
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)
}
}
crates/aether-ai-formats/src/formats/shared/sync_to_stream.rs
+31 -6
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"));