[.NET] Return ChatCompletions instead of ChatResponseMessage for token usage.

dotnet/src/AutoGen.Mistral/DTOs/ChatCompletionResponse.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
 // ChatCompletionResponse.cs
 
 using System.Collections.Generic;

dotnet/src/AutoGen.Mistral/DTOs/Error.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
 // Error.cs
 
 using System.Text.Json.Serialization;

dotnet/src/AutoGen.Mistral/DTOs/Model.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
 // Model.cs
 
 using System;

dotnet/src/AutoGen.OpenAI/Agent/OpenAIChatAgent.cs

@@ -84,7 +84,7 @@ public async Task<IMessage> GenerateReplyAsync(
         var settings = this.CreateChatCompletionsOptions(options, messages);
         var reply = await this.openAIClient.GetChatCompletionsAsync(settings, cancellationToken);
 
-        return new MessageEnvelope<ChatResponseMessage>(reply.Value.Choices.First().Message, from: this.Name);
+        return new MessageEnvelope<ChatCompletions>(reply, from: this.Name);
     }
 
     public Task<IAsyncEnumerable<IStreamingMessage>> GenerateStreamingReplyAsync(
@@ -101,7 +101,7 @@ private async IAsyncEnumerable<IStreamingMessage> StreamingReplyAsync(
         [EnumeratorCancellation] CancellationToken cancellationToken = default)
     {
         var settings = this.CreateChatCompletionsOptions(options, messages);
-        var response = await this.openAIClient.GetChatCompletionsStreamingAsync(settings);
+        var response = await this.openAIClient.GetChatCompletionsStreamingAsync(settings, cancellationToken);
         await foreach (var update in response.WithCancellation(cancellationToken))
         {
             if (update.ChoiceIndex > 0)

dotnet/src/AutoGen.OpenAI/Middleware/OpenAIChatRequestMessageConnector.cs

@@ -98,6 +98,7 @@ public IMessage PostProcessMessage(IMessage message)
             Message => message,
             AggregateMessage<ToolCallMessage, ToolCallResultMessage> => message,
             IMessage<ChatResponseMessage> m => PostProcessMessage(m),
+            IMessage<ChatCompletions> m => PostProcessMessage(m),
             _ => throw new InvalidOperationException("The type of message is not supported. Must be one of TextMessage, ImageMessage, MultiModalMessage, ToolCallMessage, ToolCallResultMessage, Message, IMessage<ChatRequestMessage>, AggregateMessage<ToolCallMessage, ToolCallResultMessage>"),
         };
     }
@@ -129,15 +130,24 @@ public IMessage PostProcessMessage(IMessage message)
 
     private IMessage PostProcessMessage(IMessage<ChatResponseMessage> message)
     {
-        var chatResponseMessage = message.Content;
+        return PostProcessMessage(message.Content, message.From);
+    }
+
+    private IMessage PostProcessMessage(IMessage<ChatCompletions> message)
+    {
+        return PostProcessMessage(message.Content.Choices[0].Message, message.From);
+    }
+
+    private IMessage PostProcessMessage(ChatResponseMessage chatResponseMessage, string? from)
+    {
         if (chatResponseMessage.Content is string content)
         {
-            return new TextMessage(Role.Assistant, content, message.From);
+            return new TextMessage(Role.Assistant, content, from);
         }
 
         if (chatResponseMessage.FunctionCall is FunctionCall functionCall)
         {
-            return new ToolCallMessage(functionCall.Name, functionCall.Arguments, message.From);
+            return new ToolCallMessage(functionCall.Name, functionCall.Arguments, from);
         }
 
         if (chatResponseMessage.ToolCalls.Where(tc => tc is ChatCompletionsFunctionToolCall).Any())
@@ -148,7 +158,7 @@ private IMessage PostProcessMessage(IMessage<ChatResponseMessage> message)
 
             var toolCalls = functionToolCalls.Select(tc => new ToolCall(tc.Name, tc.Arguments));
 
-            return new ToolCallMessage(toolCalls, message.From);
+            return new ToolCallMessage(toolCalls, from);
         }
 
         throw new InvalidOperationException("Invalid ChatResponseMessage");

dotnet/test/AutoGen.Tests/OpenAIChatAgentTest.cs

@@ -41,9 +41,10 @@ public async Task BasicConversationTestAsync()
         var chatMessageContent = MessageEnvelope.Create(new ChatRequestUserMessage("Hello"));
         var reply = await openAIChatAgent.SendAsync(chatMessageContent);
 
-        reply.Should().BeOfType<MessageEnvelope<ChatResponseMessage>>();
-        reply.As<MessageEnvelope<ChatResponseMessage>>().From.Should().Be("assistant");
-        reply.As<MessageEnvelope<ChatResponseMessage>>().Content.Role.Should().Be(ChatRole.Assistant);
+        reply.Should().BeOfType<MessageEnvelope<ChatCompletions>>();
+        reply.As<MessageEnvelope<ChatCompletions>>().From.Should().Be("assistant");
+        reply.As<MessageEnvelope<ChatCompletions>>().Content.Choices.First().Message.Role.Should().Be(ChatRole.Assistant);
+        reply.As<MessageEnvelope<ChatCompletions>>().Content.Usage.TotalTokens.Should().BeGreaterThan(0);
 
         // test streaming
         var streamingReply = await openAIChatAgent.GenerateStreamingReplyAsync(new[] { chatMessageContent });