omx(team): auto-checkpoint worker-4 [unknown]

This commit is contained in:
bellman
2026-05-15 10:21:19 +09:00
parent 3ff8743e79
commit a6ca5c489b
2 changed files with 35 additions and 20 deletions

View File

@@ -493,12 +493,7 @@ impl StreamState {
} }
if let Some(usage) = chunk.usage { if let Some(usage) = chunk.usage {
self.usage = Some(Usage { self.usage = Some(usage.normalized());
input_tokens: usage.prompt_tokens,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
output_tokens: usage.completion_tokens,
});
} }
for choice in chunk.choices { for choice in chunk.choices {
@@ -776,6 +771,29 @@ struct OpenAiUsage {
prompt_tokens: u32, prompt_tokens: u32,
#[serde(default)] #[serde(default)]
completion_tokens: u32, completion_tokens: u32,
#[serde(default)]
prompt_tokens_details: Option<OpenAiPromptTokensDetails>,
}
#[derive(Debug, Deserialize)]
struct OpenAiPromptTokensDetails {
#[serde(default)]
cached_tokens: u32,
}
impl OpenAiUsage {
fn normalized(&self) -> Usage {
let cached_tokens = self
.prompt_tokens_details
.as_ref()
.map_or(0, |details| details.cached_tokens);
Usage {
input_tokens: self.prompt_tokens.saturating_sub(cached_tokens),
cache_creation_input_tokens: 0,
cache_read_input_tokens: cached_tokens,
output_tokens: self.completion_tokens,
}
}
} }
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
@@ -1377,18 +1395,10 @@ fn normalize_response(
.finish_reason .finish_reason
.map(|value| normalize_finish_reason(&value)), .map(|value| normalize_finish_reason(&value)),
stop_sequence: None, stop_sequence: None,
usage: Usage { usage: response
input_tokens: response .usage
.usage .as_ref()
.as_ref() .map_or_else(Usage::default, OpenAiUsage::normalized),
.map_or(0, |usage| usage.prompt_tokens),
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
output_tokens: response
.usage
.as_ref()
.map_or(0, |usage| usage.completion_tokens),
},
request_id: None, request_id: None,
}) })
} }

View File

@@ -42,6 +42,9 @@ async fn send_message_uses_openai_compatible_endpoint_and_auth() {
.expect("request should succeed"); .expect("request should succeed");
assert_eq!(response.model, "grok-3"); assert_eq!(response.model, "grok-3");
assert_eq!(response.usage.input_tokens, 8);
assert_eq!(response.usage.cache_read_input_tokens, 3);
assert_eq!(response.usage.output_tokens, 5);
assert_eq!(response.total_tokens(), 16); assert_eq!(response.total_tokens(), 16);
assert_eq!( assert_eq!(
response.content, response.content,
@@ -284,7 +287,7 @@ async fn openai_streaming_requests_opt_into_usage_chunks() {
let sse = concat!( let sse = concat!(
"data: {\"id\":\"chatcmpl_openai_stream\",\"model\":\"gpt-5\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n", "data: {\"id\":\"chatcmpl_openai_stream\",\"model\":\"gpt-5\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n",
"data: {\"id\":\"chatcmpl_openai_stream\",\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", "data: {\"id\":\"chatcmpl_openai_stream\",\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
"data: {\"id\":\"chatcmpl_openai_stream\",\"choices\":[],\"usage\":{\"prompt_tokens\":9,\"completion_tokens\":4}}\n\n", "data: {\"id\":\"chatcmpl_openai_stream\",\"choices\":[],\"usage\":{\"prompt_tokens\":9,\"completion_tokens\":4,\"prompt_tokens_details\":{\"cached_tokens\":2}}}\n\n",
"data: [DONE]\n\n" "data: [DONE]\n\n"
); );
let server = spawn_server( let server = spawn_server(
@@ -339,8 +342,10 @@ async fn openai_streaming_requests_opt_into_usage_chunks() {
match &events[4] { match &events[4] {
StreamEvent::MessageDelta(MessageDeltaEvent { usage, .. }) => { StreamEvent::MessageDelta(MessageDeltaEvent { usage, .. }) => {
assert_eq!(usage.input_tokens, 9); assert_eq!(usage.input_tokens, 7);
assert_eq!(usage.cache_read_input_tokens, 2);
assert_eq!(usage.output_tokens, 4); assert_eq!(usage.output_tokens, 4);
assert_eq!(usage.total_tokens(), 13);
} }
other => panic!("expected message delta, got {other:?}"), other => panic!("expected message delta, got {other:?}"),
} }