From 678721bcf098061c36fbcc03df8bf7dbfdaa8daf Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Wed, 18 Sep 2024 11:14:53 +0200
Subject: [PATCH] Stream options.

---
 Cargo.toml           |  6 ++++--
 router/src/lib.rs    | 36 ++++++++++++++++++++++++++++++++++++
 router/src/server.rs | 21 +++++++++++++++++++--
 3 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index a50bba24..ffd45f16 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,7 +5,8 @@ members = [
   "backends/grpc-metadata",
   "backends/trtllm",
   "backends/client",
-  "launcher"
+  "launcher",
+  "router"
 ]
 default-members = [
   "benchmark",
@@ -13,7 +14,8 @@ default-members = [
   "backends/grpc-metadata",
   # "backends/trtllm",
   "backends/client",
-  "launcher"
+  "launcher",
+  "router"
 ]
 resolver = "2"
 
diff --git a/router/src/lib.rs b/router/src/lib.rs
index d8029c72..092cf7cb 100644
--- a/router/src/lib.rs
+++ b/router/src/lib.rs
@@ -684,6 +684,7 @@ pub(crate) struct ChatCompletionChunk {
     pub model: String,
     pub system_fingerprint: String,
     pub choices: Vec<ChatCompletionChoice>,
+    pub usage: Option<Usage>,
 }
 
 #[derive(Clone, Serialize, ToSchema)]
@@ -732,6 +733,7 @@ impl ChatCompletionChunk {
         created: u64,
         logprobs: Option<ChatCompletionLogprobs>,
         finish_reason: Option<String>,
+        usage: Option<Usage>,
     ) -> Self {
         let delta = match (delta, tool_calls) {
             (Some(delta), _) => ChatCompletionDelta::Chat(TextMessage {
@@ -766,6 +768,7 @@ impl ChatCompletionChunk {
                 logprobs,
                 finish_reason,
             }],
+            usage,
         }
     }
 }
@@ -880,6 +883,18 @@ pub(crate) struct ChatRequest {
     #[serde(default)]
     #[schema(nullable = true, default = "null", example = "null")]
     pub guideline: Option<String>,
+
+    /// Options for streaming response. Only set this when you set stream: true.
+    #[serde(default)]
+    #[schema(nullable = true, default = "null", example = "null")]
+    pub stream_options: Option<StreamOptions>,
+}
+
+#[derive(Clone, Deserialize, ToSchema, Serialize)]
+struct StreamOptions {
+    /// If set, an additional chunk will be streamed before the data: [DONE] message. The usage field on this chunk shows the token usage statistics for the entire request, and the choices field will always be an empty array. All other chunks will also include a usage field, but with a null value.
+    #[schema(example = "true")]
+    include_usage: bool,
 }
 
 pub fn default_tool_prompt() -> String {
@@ -1472,6 +1487,27 @@ mod tests {
         let textmsg: TextMessage = message.into();
         assert_eq!(textmsg.content, "Whats in this image?![](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/rabbit.png)");
     }
+
+    #[test]
+    fn test_chat_stream_options() {
+        let json = json!({
+            "model": "",
+            "stream_options": {"include_usage": true},
+            "messages": [{
+                "role": "user",
+                "content": "Hello"
+            }]
+        });
+        let request: ChatRequest = serde_json::from_str(json.to_string().as_str()).unwrap();
+
+        assert!(matches!(
+            request.stream_options,
+            Some(StreamOptions {
+                include_usage: true
+            })
+        ));
+    }
+
     #[test]
     fn openai_output() {
         let message = OutputMessage::ChatMessage(TextMessage {
diff --git a/router/src/server.rs b/router/src/server.rs
index 9cec2aaa..16ca1f10 100644
--- a/router/src/server.rs
+++ b/router/src/server.rs
@@ -590,7 +590,7 @@ async fn generate_stream_internal(
                                         let event = on_message_callback(stream_token);
                                         yield Ok(event);
                                     }
-                                    // Yield event for last token and compute timings
+                                    // Yield event for lat token and compute timings
                                     InferStreamResponse::End {
                                         token,
                                         generated_text,
@@ -1265,6 +1265,22 @@ async fn chat_completions(
                 (content, None)
             };
 
+            let (usage, finish_reason) = match stream_token.details {
+                Some(details) => {
+                    let completion_tokens = details.generated_tokens;
+                    let prompt_tokens = details.input_length;
+                    let total_tokens = prompt_tokens + completion_tokens;
+                    (
+                        Some(Usage {
+                            completion_tokens,
+                            prompt_tokens,
+                            total_tokens,
+                        }),
+                        Some(details.finish_reason.format(true)),
+                    )
+                }
+                None => (None, None),
+            };
             event
                 .json_data(CompletionType::ChatCompletionChunk(
                     ChatCompletionChunk::new(
@@ -1274,7 +1290,8 @@ async fn chat_completions(
                         tool_calls,
                         current_time,
                         logprobs,
-                        stream_token.details.map(|d| d.finish_reason.format(true)),
+                        finish_reason,
+                        usage,
                     ),
                 ))
                 .unwrap_or_else(|e| {