i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-11 19:59: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/openai/chat/stream.rs
+12 -20
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]
crates/aether-ai-formats/src/formats/openai/image/stream.rs
+46 -6
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 {