.Net: Adding usage metadata to OpenAI Streaming Chunks (#9022)

### Motivation and Context Currently we have metadata for non-streaming results, and to be consistent, the same should be available for streaming for the final chunk. - Resolves #6826
microsoft · Sep 30, 2024 · 5b0a1af · 5b0a1af
1 parent 2724f34
commit 5b0a1af
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 27 deletions.
diff --git a/dotnet/samples/Concepts/ChatCompletion/OpenAI_ChatCompletionStreaming.cs b/dotnet/samples/Concepts/ChatCompletion/OpenAI_ChatCompletionStreaming.cs
@@ -16,7 +16,7 @@ public class OpenAI_ChatCompletionStreaming(ITestOutputHelper output) : BaseTest
     /// This example demonstrates chat completion streaming using OpenAI.
     /// </summary>
     [Fact]
-    public Task StreamServicePromptAsync()
+    public async Task StreamServicePromptAsync()
     {
         Assert.NotNull(TestConfiguration.OpenAI.ChatModelId);
         Assert.NotNull(TestConfiguration.OpenAI.ApiKey);
@@ -25,7 +25,25 @@ public Task StreamServicePromptAsync()
 
         OpenAIChatCompletionService chatCompletionService = new(TestConfiguration.OpenAI.ChatModelId, TestConfiguration.OpenAI.ApiKey);
 
-        return this.StartStreamingChatAsync(chatCompletionService);
+        Console.WriteLine("Chat content:");
+        Console.WriteLine("------------------------");
+
+        var chatHistory = new ChatHistory("You are a librarian, expert about books");
+        OutputLastMessage(chatHistory);
+
+        // First user message
+        chatHistory.AddUserMessage("Hi, I'm looking for book suggestions");
+        OutputLastMessage(chatHistory);
+
+        // First assistant message
+        await StreamMessageOutputAsync(chatCompletionService, chatHistory, AuthorRole.Assistant);
+
+        // Second user message
+        chatHistory.AddUserMessage("I love history and philosophy, I'd like to learn something new about Greece, any suggestion?");
+        OutputLastMessage(chatHistory);
+
+        // Second assistant message
+        await StreamMessageOutputAsync(chatCompletionService, chatHistory, AuthorRole.Assistant);
     }
 
     /// <summary>
@@ -196,30 +214,7 @@ public async Task StreamFunctionCallContentAsync()
         }
     }
 
-    private async Task StartStreamingChatAsync(IChatCompletionService chatCompletionService)
-    {
-        Console.WriteLine("Chat content:");
-        Console.WriteLine("------------------------");
-
-        var chatHistory = new ChatHistory("You are a librarian, expert about books");
-        OutputLastMessage(chatHistory);
-
-        // First user message
-        chatHistory.AddUserMessage("Hi, I'm looking for book suggestions");
-        OutputLastMessage(chatHistory);
-
-        // First assistant message
-        await StreamMessageOutputAsync(chatCompletionService, chatHistory, AuthorRole.Assistant);
-
-        // Second user message
-        chatHistory.AddUserMessage("I love history and philosophy, I'd like to learn something new about Greece, any suggestion?");
-        OutputLastMessage(chatHistory);
-
-        // Second assistant message
-        await StreamMessageOutputAsync(chatCompletionService, chatHistory, AuthorRole.Assistant);
-    }
-
-    private async Task StreamMessageOutputAsync(IChatCompletionService chatCompletionService, ChatHistory chatHistory, AuthorRole authorRole)
+    private async Task StreamMessageOutputAsync(OpenAIChatCompletionService chatCompletionService, ChatHistory chatHistory, AuthorRole authorRole)
     {
         bool roleWritten = false;
         string fullMessage = string.Empty;
@@ -237,6 +232,13 @@ private async Task StreamMessageOutputAsync(IChatCompletionService chatCompletio
                 fullMessage += chatUpdate.Content;
                 Console.Write(chatUpdate.Content);
             }
+
+            // The last message in the chunk has the usage metadata.
+            // https://platform.openai.com/docs/api-reference/chat/create#chat-create-stream_options
+            if (chatUpdate.Metadata?["Usage"] is not null)
+            {
+                Console.WriteLine(chatUpdate.Metadata["Usage"]?.AsJson());
+            }
         }
 
         Console.WriteLine("\n------------------------");
@@ -259,6 +261,13 @@ private async Task<string> StreamMessageOutputFromKernelAsync(Kernel kernel, str
                 fullMessage += chatUpdate.Content;
                 Console.Write(chatUpdate.Content);
             }
+
+            // The last message in the chunk has the usage metadata.
+            // https://platform.openai.com/docs/api-reference/chat/create#chat-create-stream_options
+            if (chatUpdate.Metadata?["Usage"] is not null)
+            {
+                Console.WriteLine(chatUpdate.Metadata["Usage"]?.AsJson());
+            }
         }
         Console.WriteLine("\n------------------------");
         return fullMessage;
@@ -342,7 +351,8 @@ private void OutputInnerContent(OpenAI.Chat.StreamingChatCompletionUpdate stream
             }
         }
 
-        /// The last message in the chunk is a <see cref="ChatDoneResponseStream"/> type with additional metadata.
+        // The last message in the chunk has the usage metadata.
+        // https://platform.openai.com/docs/api-reference/chat/create#chat-create-stream_options
         if (streamChunk.Usage is not null)
         {
             Console.WriteLine($"Usage input tokens: {streamChunk.Usage.InputTokens}");

diff --git a/...t/src/Connectors/Connectors.OpenAI.UnitTests/Services/OpenAIChatCompletionServiceTests.cs b/...t/src/Connectors/Connectors.OpenAI.UnitTests/Services/OpenAIChatCompletionServiceTests.cs
@@ -323,6 +323,10 @@ public async Task GetStreamingChatMessageContentsWorksCorrectlyAsync()
 
         await enumerator.MoveNextAsync();
         Assert.Equal("Stop", enumerator.Current.Metadata?["FinishReason"]);
+
+        await enumerator.MoveNextAsync();
+        Assert.NotNull(enumerator.Current.Metadata?["Usage"]);
+        Assert.Equal("{\"OutputTokens\":8,\"InputTokens\":13,\"TotalTokens\":21}", JsonSerializer.Serialize(enumerator.Current.Metadata?["Usage"]));
     }
 
     [Fact]

diff --git a/...nnectors/Connectors.OpenAI.UnitTests/TestData/chat_completion_streaming_test_response.txt b/...nnectors/Connectors.OpenAI.UnitTests/TestData/chat_completion_streaming_test_response.txt
@@ -2,4 +2,6 @@ data: {"id":"chatcmpl-96fqQVHGjG9Yzs4ZMB1K6nfy2oEoo","object":"chat.completion.c
 
 data: {"id":"chatcmpl-96fqQVHGjG9Yzs4ZMB1K6nfy2oEoo","object":"chat.completion.chunk","created":1711377846,"model":"gpt-4-0125-preview","system_fingerprint":"fp_a7daf7c51e","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}]}
 
+data: {"id":"chatcmpl-96fqQVHGjG9Yzs4ZMB1K6nfy2oEoo","object":"chat.completion.chunk","created":1711377846,"model":"gpt-4-0125-preview","system_fingerprint":"fp_a7daf7c51e","choices":[],"usage":{"prompt_tokens":13,"completion_tokens":8,"total_tokens":21,"completion_tokens_details":{"reasoning_tokens":0}}}
+
 data: [DONE]
diff --git a/dotnet/src/Connectors/Connectors.OpenAI/Core/ClientCore.ChatCompletion.cs b/dotnet/src/Connectors/Connectors.OpenAI/Core/ClientCore.ChatCompletion.cs
@@ -118,6 +118,7 @@ protected record ToolCallingConfig(IList<ChatTool>? Tools, ChatToolChoice? Choic
             { nameof(completionUpdate.CreatedAt), completionUpdate.CreatedAt },
             { nameof(completionUpdate.SystemFingerprint), completionUpdate.SystemFingerprint },
             { nameof(completionUpdate.RefusalUpdate), completionUpdate.RefusalUpdate },
+            { nameof(completionUpdate.Usage), completionUpdate.Usage },
 
             // Serialization of this struct behaves as an empty object {}, need to cast to string to avoid it.
             { nameof(completionUpdate.FinishReason), completionUpdate.FinishReason?.ToString() },