diff --git a/lib/components/fabro-llm/src/codec/openai_compatible/request.rs b/lib/components/fabro-llm/src/codec/openai_compatible/request.rs index 12be9e84b..f6bbf3b16 100644 --- a/lib/components/fabro-llm/src/codec/openai_compatible/request.rs +++ b/lib/components/fabro-llm/src/codec/openai_compatible/request.rs @@ -1,7 +1,7 @@ //! Request encoding: canonical `Request` → Chat Completions body. use super::translate; -use super::wire::{ApiRequest, ChatMessage}; +use super::wire::{ApiRequest, ChatMessage, StreamOptions}; use crate::codec::{CodecCtx, EncodedRequest, cache, merge_named_provider_options}; use crate::error::Error; @@ -10,8 +10,10 @@ use crate::error::Error; const KNOWN_OPTION_KEYS: &[&str] = &["auto_cache"]; /// Build the Chat Completions request for `ctx.request`. `stream` toggles the -/// `stream` body field. The body is assembled as a `serde_json::Value` so -/// `provider_options.` fields can be merged in before sending. +/// `stream` body field and the `stream_options.include_usage` opt-in that makes +/// providers emit the trailing usage chunk. The body is assembled as a +/// `serde_json::Value` so `provider_options.` fields can be +/// merged in before sending. /// /// Returns an error when the request contains a custom tool definition, which /// the Chat Completions tool envelope cannot represent. @@ -55,6 +57,9 @@ pub(super) fn encode(ctx: &CodecCtx<'_>, stream: bool) -> Result, #[serde(skip_serializing_if = "Option::is_none")] pub stream: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stream_options: Option, +} + +/// Streaming options. Chat Completions only emits the trailing usage chunk +/// when the request opts in, so without this a streamed response reports zero +/// tokens and costs are estimated at $0. +#[derive(serde::Serialize)] +pub(super) struct StreamOptions { + pub include_usage: bool, } #[derive(serde::Serialize)] diff --git a/lib/components/fabro-llm/tests/it/wire/openai_compatible.rs b/lib/components/fabro-llm/tests/it/wire/openai_compatible.rs index 79d61223c..90227f9c1 100644 --- a/lib/components/fabro-llm/tests/it/wire/openai_compatible.rs +++ b/lib/components/fabro-llm/tests/it/wire/openai_compatible.rs @@ -812,7 +812,8 @@ async fn stream_text_happy_path_capture() -> (WireCapture, Vec, #[serde(default)] pub stream: bool, + stream_options: Option, pub tools: Option>, pub tool_choice: Option, pub response_format: Option, pub stop: Option, } +#[derive(Clone, Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct ChatStreamOptions { + #[serde(default)] + include_usage: bool, +} + impl ChatCompletionsRequest { + pub fn include_stream_usage(&self) -> bool { + self.stream_options + .as_ref() + .is_some_and(|options| options.include_usage) + } + pub fn extract_user_text(&self) -> String { let pieces: Vec = self .messages diff --git a/test/twin/openai/src/sse.rs b/test/twin/openai/src/sse.rs index eb49c3f47..2d30a0f88 100644 --- a/test/twin/openai/src/sse.rs +++ b/test/twin/openai/src/sse.rs @@ -236,7 +236,11 @@ pub fn responses_sse_response(plan: &ResponsePlan, transport: TransportOptions) stream_response(events, transport) } -pub fn chat_sse_response(plan: &ResponsePlan, transport: TransportOptions) -> Response { +pub fn chat_sse_response( + plan: &ResponsePlan, + include_usage: bool, + transport: TransportOptions, +) -> Response { let mut events = Vec::new(); let content = plan.chat_content(); events.push(chat_chunk(&json!({ @@ -321,6 +325,16 @@ pub fn chat_sse_response(plan: &ResponsePlan, transport: TransportOptions) -> Re "finish_reason": if plan.tool_calls.is_empty() { "stop" } else { "tool_calls" }, }] }))); + if include_usage { + events.push(chat_chunk(&json!({ + "id": format!("chatcmpl_{}", plan.id), + "object": "chat.completion.chunk", + "created": plan.created, + "model": plan.model, + "choices": [], + "usage": plan.usage.chat_completions_json(), + }))); + } events.push("data: [DONE]\n\n".to_owned()); } diff --git a/test/twin/openai/tests/chat_completions_contract.rs b/test/twin/openai/tests/chat_completions_contract.rs index 6fe4c27fe..e17369c42 100644 --- a/test/twin/openai/tests/chat_completions_contract.rs +++ b/test/twin/openai/tests/chat_completions_contract.rs @@ -50,9 +50,48 @@ async fn chat_completions_stream_uses_same_canonical_plan() { assert_eq!(status, 200); assert!(joined.contains("\"content\":\"deterministic: stream same plan\"")); + assert!(!joined.contains("\"usage\"")); assert!(joined.contains("data: [DONE]")); } +#[tokio::test] +async fn chat_completions_stream_includes_usage_when_requested() { + let server = common::spawn_server().await.expect("server should start"); + + let (status, chunks) = server + .post_chat_stream(json!({ + "model": "gpt-test", + "messages": [{ "role": "user", "content": "stream with usage" }], + "stream": true, + "stream_options": { "include_usage": true } + })) + .await; + + assert_eq!(status, 200); + let transcript = + common::parse_sse_transcript(chunks.join("").as_bytes()).expect("valid SSE transcript"); + let usage_chunk = transcript + .events + .iter() + .filter(|event| event.data != "[DONE]") + .map(|event| { + serde_json::from_str::(&event.data).expect("valid JSON chunk") + }) + .find(|chunk| chunk.get("usage").is_some()) + .expect("trailing usage chunk"); + + assert_eq!(usage_chunk["choices"], json!([])); + assert_eq!( + usage_chunk["usage"], + json!({ + "prompt_tokens": 3, + "completion_tokens": 5, + "total_tokens": 8 + }) + ); + assert!(transcript.done); +} + #[tokio::test] async fn chat_completions_accepts_supported_openai_compatible_fields() { let server = common::spawn_server().await.expect("server should start");