diff --git a/apis/src/anthropic/stream_events/mod.rs b/apis/src/anthropic/stream_events/mod.rs index 6d245202f1..cb9ec7372c 100644 --- a/apis/src/anthropic/stream_events/mod.rs +++ b/apis/src/anthropic/stream_events/mod.rs @@ -58,6 +58,12 @@ const FINISH_REASON_KEY: &str = "anthropic_stream.finish_reason"; /// Metadata key for accumulated output token count. const OUTPUT_TOKENS_KEY: &str = "anthropic_stream.output_tokens"; +/// Metadata key for accumulated input (prompt) token count. +const INPUT_TOKENS_KEY: &str = "anthropic_stream.input_tokens"; + +/// Metadata key for cached input token count. +const CACHE_READ_TOKENS_KEY: &str = "anthropic_stream.cache_read_tokens"; + /// Metadata key for the current content block index. const BLOCK_INDEX_KEY: &str = "anthropic_stream.block_index"; @@ -465,12 +471,20 @@ fn transform_chunk(ctx: &mut HttpFilterContext<'_>, chunk: &Value, output: &mut } } - if let Some(ot) = chunk - .get("usage") - .and_then(|u| u.get("completion_tokens")) - .and_then(Value::as_u64) - { - ctx.set_metadata(OUTPUT_TOKENS_KEY, ot.to_string()); + if let Some(usage) = chunk.get("usage") { + if let Some(ot) = usage.get("completion_tokens").and_then(Value::as_u64) { + ctx.set_metadata(OUTPUT_TOKENS_KEY, ot.to_string()); + } + if let Some(pt) = usage.get("prompt_tokens").and_then(Value::as_u64) { + ctx.set_metadata(INPUT_TOKENS_KEY, pt.to_string()); + } + if let Some(ct) = usage + .get("prompt_tokens_details") + .and_then(|d| d.get("cached_tokens")) + .and_then(Value::as_u64) + { + ctx.set_metadata(CACHE_READ_TOKENS_KEY, ct.to_string()); + } } } @@ -707,11 +721,7 @@ fn emit_message_delta(ctx: &HttpFilterContext<'_>, output: &mut Vec) { .get(FINISH_REASON_KEY) .map_or("end_turn", |v| map_stop_reason(v)); - let output_tokens: u64 = ctx - .filter_metadata - .get(OUTPUT_TOKENS_KEY) - .and_then(|v| v.parse().ok()) - .unwrap_or(0); + let usage = collect_delta_usage(ctx); emit_event( output, @@ -724,21 +734,39 @@ fn emit_message_delta(ctx: &HttpFilterContext<'_>, output: &mut Vec) { "stop_reason": stop_reason, "stop_sequence": null }, - "usage": message_delta_usage(output_tokens) + "usage": usage }), ); } +/// Collect token counts from metadata and build the terminal delta usage. +fn collect_delta_usage(ctx: &HttpFilterContext<'_>) -> MessageDeltaUsage { + let output_tokens: u64 = ctx + .filter_metadata + .get(OUTPUT_TOKENS_KEY) + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + + let prompt_tokens: Option = ctx.filter_metadata.get(INPUT_TOKENS_KEY).and_then(|v| v.parse().ok()); + + let cache_read: Option = ctx + .filter_metadata + .get(CACHE_READ_TOKENS_KEY) + .and_then(|v| v.parse().ok()); + + let input_tokens = prompt_tokens.map(|pt| match cache_read { + Some(cached) => pt.saturating_sub(cached), + None => pt, + }); + + MessageDeltaUsage::new(output_tokens, input_tokens, cache_read) +} + /// Build a schema-complete Anthropic `Message.usage` value. fn message_start_usage() -> MessageUsage { MessageUsage::new(0, 0, None) } -/// Build a schema-complete Anthropic `message_delta.usage` value. -fn message_delta_usage(output_tokens: u64) -> MessageDeltaUsage { - MessageDeltaUsage::new(output_tokens) -} - /// Map `OpenAI` finish reasons to Anthropic stop reasons. fn map_stop_reason(reason: &str) -> &str { match reason { @@ -987,7 +1015,7 @@ mod tests { fn message_delta_usage_matches_anthropic_schema() { let (filter, mut ctx) = make_filter_and_context(); - let chunk1 = "data: {\"id\":\"c1\",\"model\":\"gpt-4\",\"choices\":[{\"delta\":{\"content\":\"Hi\"},\"index\":0,\"finish_reason\":\"stop\"}],\"usage\":{\"completion_tokens\":7}}\n\n"; + let chunk1 = "data: {\"id\":\"c1\",\"model\":\"gpt-4\",\"choices\":[{\"delta\":{\"content\":\"Hi\"},\"index\":0,\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":15,\"completion_tokens\":7}}\n\n"; let mut body1 = Some(Bytes::from(chunk1)); drop(filter.on_response_body(&mut ctx, &mut body1, false).unwrap()); @@ -1011,13 +1039,63 @@ mod tests { &[ "cache_creation_input_tokens", "cache_read_input_tokens", - "input_tokens", "server_tool_use", ], "message_delta usage", ); assert_absent_fields(usage, &["output_tokens_details"], "message_delta usage"); assert_u64_field(usage, "output_tokens", 7, "message_delta usage"); + assert_u64_field(usage, "input_tokens", 15, "message_delta usage"); + } + + #[test] + fn message_delta_usage_with_cached_tokens() { + let (filter, mut ctx) = make_filter_and_context(); + + let chunk1 = "data: {\"id\":\"c1\",\"model\":\"gpt-4\",\"choices\":[{\"delta\":{\"content\":\"Hi\"},\"index\":0,\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":100,\"completion_tokens\":5,\"prompt_tokens_details\":{\"cached_tokens\":80}}}\n\n"; + let mut body1 = Some(Bytes::from(chunk1)); + drop(filter.on_response_body(&mut ctx, &mut body1, false).unwrap()); + + let done = "data: [DONE]\n\n"; + let mut body2 = Some(Bytes::from(done)); + drop(filter.on_response_body(&mut ctx, &mut body2, false).unwrap()); + + let out = String::from_utf8(body2.unwrap().to_vec()).unwrap(); + let event = event_data(&out, "message_delta"); + let usage = event.get("usage").unwrap(); + + assert_u64_field(usage, "output_tokens", 5, "message_delta usage"); + assert_u64_field( + usage, + "input_tokens", + 20, + "input_tokens should exclude cached (100 - 80)", + ); + assert_u64_field(usage, "cache_read_input_tokens", 80, "message_delta usage"); + } + + #[test] + fn message_delta_usage_without_usage_chunk() { + let (filter, mut ctx) = make_filter_and_context(); + + let chunk1 = "data: {\"id\":\"c1\",\"model\":\"gpt-4\",\"choices\":[{\"delta\":{\"content\":\"Hi\"},\"index\":0,\"finish_reason\":\"stop\"}]}\n\n"; + let mut body1 = Some(Bytes::from(chunk1)); + drop(filter.on_response_body(&mut ctx, &mut body1, false).unwrap()); + + let done = "data: [DONE]\n\n"; + let mut body2 = Some(Bytes::from(done)); + drop(filter.on_response_body(&mut ctx, &mut body2, false).unwrap()); + + let out = String::from_utf8(body2.unwrap().to_vec()).unwrap(); + let event = event_data(&out, "message_delta"); + let usage = event.get("usage").unwrap(); + + assert_u64_field(usage, "output_tokens", 0, "no usage chunk means zero output_tokens"); + assert_null_fields( + usage, + &["input_tokens", "cache_read_input_tokens"], + "no usage chunk means null input fields", + ); } #[test] diff --git a/apis/src/anthropic/to_openai/request.rs b/apis/src/anthropic/to_openai/request.rs index 465ac27868..8376e8b7aa 100644 --- a/apis/src/anthropic/to_openai/request.rs +++ b/apis/src/anthropic/to_openai/request.rs @@ -36,9 +36,7 @@ pub(crate) fn transform_request(body: &[u8]) -> Result, String> { chat.insert("max_completion_tokens".to_owned(), max_tokens.clone()); } - if let Some(stream) = obj.get("stream") { - chat.insert("stream".to_owned(), stream.clone()); - } + convert_stream(&mut chat, obj); map_parameters(&mut chat, obj); convert_tools(&mut chat, obj); @@ -498,6 +496,24 @@ fn non_empty_lines(lines: &[String]) -> Option { // Parameter Mapping // ----------------------------------------------------------------------------- +/// Copy `stream` and request streaming usage when enabled. +fn convert_stream(chat: &mut Map, obj: &Map) { + let Some(stream) = obj.get("stream") else { + return; + }; + chat.insert("stream".to_owned(), stream.clone()); + + if stream.as_bool() == Some(true) { + let mut opts = obj + .get("stream_options") + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + opts.insert("include_usage".to_owned(), Value::Bool(true)); + chat.insert("stream_options".to_owned(), Value::Object(opts)); + } +} + /// Map Anthropic parameters to Chat Completions-compatible equivalents. /// /// `top_k` has no standard Chat Completions equivalent but is preserved @@ -1237,4 +1253,42 @@ mod tests { assert_eq!(tools.len(), 1, "only non-filtered tools should remain"); assert_eq!(tools[0]["function"]["name"], "get_weather"); } + + #[test] + fn streaming_request_includes_usage_option() { + let body = br#"{"model":"claude-opus-4-8","max_tokens":1024,"stream":true,"messages":[{"role":"user","content":"Hi"}]}"#; + let result = transform_request(body).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + assert_eq!(parsed["stream"], true, "stream should be true"); + assert_eq!( + parsed["stream_options"]["include_usage"], true, + "stream_options.include_usage should be set" + ); + } + + #[test] + fn non_streaming_request_omits_stream_options() { + let body = br#"{"model":"claude-opus-4-8","max_tokens":1024,"messages":[{"role":"user","content":"Hi"}]}"#; + let result = transform_request(body).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + assert!( + parsed.get("stream_options").is_none(), + "stream_options should not be present without stream:true" + ); + } + + #[test] + fn stream_false_omits_stream_options() { + let body = br#"{"model":"claude-opus-4-8","max_tokens":1024,"stream":false,"messages":[{"role":"user","content":"Hi"}]}"#; + let result = transform_request(body).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + assert_eq!(parsed["stream"], false, "stream should be false"); + assert!( + parsed.get("stream_options").is_none(), + "stream_options should not be present when stream is false" + ); + } } diff --git a/apis/src/anthropic/wire.rs b/apis/src/anthropic/wire.rs index b431017cac..95ddfba9f6 100644 --- a/apis/src/anthropic/wire.rs +++ b/apis/src/anthropic/wire.rs @@ -144,12 +144,12 @@ pub(crate) struct MessageDeltaUsage { } impl MessageDeltaUsage { - /// Create terminal delta usage from the cumulative output token count. - pub(crate) fn new(output_tokens: u64) -> Self { + /// Create terminal delta usage with token counts. + pub(crate) fn new(output_tokens: u64, input_tokens: Option, cache_read_input_tokens: Option) -> Self { Self { cache_creation_input_tokens: None, - cache_read_input_tokens: None, - input_tokens: None, + cache_read_input_tokens, + input_tokens, output_tokens, server_tool_use: None, } diff --git a/tests/integration/fixtures/inference/recordings/openai/messages/basic-stream.json b/tests/integration/fixtures/inference/recordings/openai/messages/basic-stream.json index aeacf80fb1..768ff68019 100644 --- a/tests/integration/fixtures/inference/recordings/openai/messages/basic-stream.json +++ b/tests/integration/fixtures/inference/recordings/openai/messages/basic-stream.json @@ -3,15 +3,16 @@ "scenario_id": "messages/basic-stream", "protocol": "anthropic_messages", "provenance": { - "kind": "imported", + "kind": "live", "provider": "openai", "model": "gpt-4o", - "source_id": "tests/integration/messages/test_messages.py::test_messages_streaming_basic[txt=openai/gpt-4o]" + "source_id": null }, "normalization": { "version": 1, "linked_ids": { - "msg_18c8cd5a4f2dd888": "msg_recorded_0001" + "chatcmpl-EIu0ffPZypE9HdAyNsX3xyfsZuLMJ": "chatcmpl-recorded-0002", + "msg_18d0deb35bd64925": "msg_recorded_0001" } }, "turns": [ @@ -46,6 +47,12 @@ "headers": { "content-type": [ "text/event-stream" + ], + "request-id": [ + "req_55bee3e2fe904be0b4b993ce9415e87e" + ], + "x-request-id": [ + "req_55bee3e2fe904be0b4b993ce9415e87e" ] }, "body": { @@ -131,7 +138,7 @@ }, { "event": "message_delta", - "data": "{\"delta\":{\"container\":null,\"stop_details\":null,\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"type\":\"message_delta\",\"usage\":{\"cache_creation_input_tokens\":null,\"cache_read_input_tokens\":null,\"input_tokens\":null,\"output_tokens\":0,\"server_tool_use\":null}}", + "data": "{\"delta\":{\"container\":null,\"stop_details\":null,\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"type\":\"message_delta\",\"usage\":{\"cache_creation_input_tokens\":null,\"cache_read_input_tokens\":0,\"input_tokens\":13,\"output_tokens\":9,\"server_tool_use\":null}}", "id": null, "retry": null }, @@ -166,7 +173,10 @@ } ], "model": "gpt-4o", - "stream": true + "stream": true, + "stream_options": { + "include_usage": true + } } } }, @@ -174,7 +184,10 @@ "status": 200, "headers": { "content-type": [ - "text/event-stream" + "text/event-stream; charset=utf-8" + ], + "x-request-id": [ + "req_55bee3e2fe904be0b4b993ce9415e87e" ] }, "body": { @@ -182,67 +195,73 @@ "frames": [ { "event": null, - "data": "{\"choices\":[{\"delta\":{\"content\":\"\",\"function_call\":null,\"refusal\":null,\"role\":\"assistant\",\"tool_calls\":null},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"rec-dc11ec5aa35f\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "data": "{\"choices\":[{\"delta\":{\"content\":\"\",\"refusal\":null,\"role\":\"assistant\"},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "id": null, + "retry": null + }, + { + "event": null, + "data": "{\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", "id": null, "retry": null }, { "event": null, - "data": "{\"choices\":[{\"delta\":{\"content\":\"Hello\",\"function_call\":null,\"refusal\":null,\"role\":null,\"tool_calls\":null},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"rec-dc11ec5aa35f\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "data": "{\"choices\":[{\"delta\":{\"content\":\"!\"},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", "id": null, "retry": null }, { "event": null, - "data": "{\"choices\":[{\"delta\":{\"content\":\"!\",\"function_call\":null,\"refusal\":null,\"role\":null,\"tool_calls\":null},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"rec-dc11ec5aa35f\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "data": "{\"choices\":[{\"delta\":{\"content\":\" How\"},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", "id": null, "retry": null }, { "event": null, - "data": "{\"choices\":[{\"delta\":{\"content\":\" How\",\"function_call\":null,\"refusal\":null,\"role\":null,\"tool_calls\":null},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"rec-dc11ec5aa35f\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "data": "{\"choices\":[{\"delta\":{\"content\":\" can\"},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", "id": null, "retry": null }, { "event": null, - "data": "{\"choices\":[{\"delta\":{\"content\":\" can\",\"function_call\":null,\"refusal\":null,\"role\":null,\"tool_calls\":null},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"rec-dc11ec5aa35f\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "data": "{\"choices\":[{\"delta\":{\"content\":\" I\"},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", "id": null, "retry": null }, { "event": null, - "data": "{\"choices\":[{\"delta\":{\"content\":\" I\",\"function_call\":null,\"refusal\":null,\"role\":null,\"tool_calls\":null},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"rec-dc11ec5aa35f\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "data": "{\"choices\":[{\"delta\":{\"content\":\" assist\"},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", "id": null, "retry": null }, { "event": null, - "data": "{\"choices\":[{\"delta\":{\"content\":\" assist\",\"function_call\":null,\"refusal\":null,\"role\":null,\"tool_calls\":null},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"rec-dc11ec5aa35f\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "data": "{\"choices\":[{\"delta\":{\"content\":\" you\"},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", "id": null, "retry": null }, { "event": null, - "data": "{\"choices\":[{\"delta\":{\"content\":\" you\",\"function_call\":null,\"refusal\":null,\"role\":null,\"tool_calls\":null},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"rec-dc11ec5aa35f\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "data": "{\"choices\":[{\"delta\":{\"content\":\" today\"},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", "id": null, "retry": null }, { "event": null, - "data": "{\"choices\":[{\"delta\":{\"content\":\" today\",\"function_call\":null,\"refusal\":null,\"role\":null,\"tool_calls\":null},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"rec-dc11ec5aa35f\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "data": "{\"choices\":[{\"delta\":{\"content\":\"?\"},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", "id": null, "retry": null }, { "event": null, - "data": "{\"choices\":[{\"delta\":{\"content\":\"?\",\"function_call\":null,\"refusal\":null,\"role\":null,\"tool_calls\":null},\"finish_reason\":null,\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"rec-dc11ec5aa35f\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "data": "{\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\",\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", "id": null, "retry": null }, { "event": null, - "data": "{\"choices\":[{\"delta\":{\"content\":null,\"function_call\":null,\"refusal\":null,\"role\":null,\"tool_calls\":null},\"finish_reason\":\"stop\",\"index\":0,\"logprobs\":null}],\"created\":0,\"id\":\"rec-dc11ec5aa35f\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":null}", + "data": "{\"choices\":[],\"created\":0,\"id\":\"chatcmpl-recorded-0002\",\"model\":\"gpt-4o-2024-08-06\",\"object\":\"chat.completion.chunk\",\"service_tier\":\"default\",\"usage\":{\"completion_tokens\":9,\"completion_tokens_details\":{\"accepted_prediction_tokens\":0,\"audio_tokens\":0,\"reasoning_tokens\":0,\"rejected_prediction_tokens\":0},\"prompt_tokens\":13,\"prompt_tokens_details\":{\"audio_tokens\":0,\"cached_tokens\":0},\"total_tokens\":22}}", "id": null, "retry": null } diff --git a/tests/integration/fixtures/inference/recordings/vllm/messages/basic-stream.json b/tests/integration/fixtures/inference/recordings/vllm/messages/basic-stream.json index dc031510af..4fbc3e5bfe 100644 --- a/tests/integration/fixtures/inference/recordings/vllm/messages/basic-stream.json +++ b/tests/integration/fixtures/inference/recordings/vllm/messages/basic-stream.json @@ -167,7 +167,10 @@ } ], "model": "RedHatAI/Qwen3-Coder-Next-NVFP4", - "stream": true + "stream": true, + "stream_options": { + "include_usage": true + } } } }, diff --git a/tests/utils/src/inference_fixture/replay.rs b/tests/utils/src/inference_fixture/replay.rs index 4dae9d5e51..e5ad327ecc 100644 --- a/tests/utils/src/inference_fixture/replay.rs +++ b/tests/utils/src/inference_fixture/replay.rs @@ -3648,12 +3648,18 @@ mod tests { path: "/v1/chat/completions".to_owned(), headers: BTreeMap::from([("content-type".to_owned(), vec!["application/json".to_owned()])]), body: RecordedBody::Json { - value: json!({ - "model": model, - "max_completion_tokens": 64, - "stream": stream, - "messages": [{"role": "user", "content": prompt}], - }), + value: { + let mut obj = json!({ + "model": model, + "max_completion_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": prompt}], + }); + if stream { + obj["stream_options"] = json!({"include_usage": true}); + } + obj + }, }, } }