diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs index 1366422b8ea..70fddc68718 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs @@ -252,6 +252,13 @@ public override async Task CompleteAsync(IList chat } } + // If the original chat completion included usage data, + // add that into the message so it's available in the history. + if (KeepFunctionCallingMessages && response.Usage is { } usage) + { + response.Message.Contents = [.. response.Message.Contents, new UsageContent(usage)]; + } + // Add the responses from the function calls into the history. var modeAndMessages = await ProcessFunctionCallsAsync(chatMessages, options, functionCallContents, iteration, cancellationToken).ConfigureAwait(false); if (modeAndMessages.MessagesAdded is not null) diff --git a/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ChatClientIntegrationTests.cs b/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ChatClientIntegrationTests.cs index cf113c878f6..ab8e7613edb 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ChatClientIntegrationTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ChatClientIntegrationTests.cs @@ -163,13 +163,25 @@ public virtual async Task FunctionInvocation_AutomaticallyInvokeFunction_Paramet int secretNumber = 42; - var response = await chatClient.CompleteAsync("What is the current secret number?", new() + List messages = + [ + new(ChatRole.User, "What is the current secret number?") + ]; + + var response = await chatClient.CompleteAsync(messages, new() { Tools = [AIFunctionFactory.Create(() => secretNumber, "GetSecretNumber")] }); Assert.Single(response.Choices); Assert.Contains(secretNumber.ToString(), response.Message.Text); + + if (response.Usage is { } finalUsage) + { + UsageContent? intermediate = messages.SelectMany(m => m.Contents).OfType().FirstOrDefault(); + Assert.NotNull(intermediate); + Assert.True(finalUsage.TotalTokenCount > intermediate.Details.TotalTokenCount); + } } [ConditionalFact]