diff --git a/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ChatClientIntegrationTests.cs b/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ChatClientIntegrationTests.cs index 0e4311a227a..edb6c5dd14c 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ChatClientIntegrationTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ChatClientIntegrationTests.cs @@ -229,16 +229,7 @@ public virtual async Task FunctionInvocation_AutomaticallyInvokeFunction_Paramet }); Assert.Contains(secretNumber.ToString(), response.Text); - - // If the underlying IChatClient provides usage data, function invocation should aggregate the - // usage data across all calls to produce a single Usage value on the final response - if (response.Usage is { } finalUsage) - { - var totalInputTokens = activities.Sum(a => (int?)a.GetTagItem("gen_ai.response.input_tokens")!); - var totalOutputTokens = activities.Sum(a => (int?)a.GetTagItem("gen_ai.response.output_tokens")!); - Assert.Equal(totalInputTokens, finalUsage.InputTokenCount); - Assert.Equal(totalOutputTokens, finalUsage.OutputTokenCount); - } + AssertUsageAgainstActivities(response, activities); } [ConditionalFact] @@ -306,16 +297,7 @@ public virtual async Task FunctionInvocation_OptionalParameter() }); Assert.Contains(secretNumber.ToString(), response.Text); - - // If the underlying IChatClient provides usage data, function invocation should aggregate the - // usage data across all calls to produce a single Usage value on the final response - if (response.Usage is { } finalUsage) - { - var totalInputTokens = activities.Sum(a => (int?)a.GetTagItem("gen_ai.response.input_tokens")!); - var totalOutputTokens = activities.Sum(a => (int?)a.GetTagItem("gen_ai.response.output_tokens")!); - Assert.Equal(totalInputTokens, finalUsage.InputTokenCount); - Assert.Equal(totalOutputTokens, finalUsage.OutputTokenCount); - } + AssertUsageAgainstActivities(response, activities); } [ConditionalFact] @@ -347,15 +329,21 @@ public virtual async Task FunctionInvocation_NestedParameters() }); Assert.Contains((secretNumber + 19).ToString(), response.Text); + AssertUsageAgainstActivities(response, activities); + } + private static void AssertUsageAgainstActivities(ChatResponse response, List activities) + { // If the underlying IChatClient provides usage data, function invocation should aggregate the - // usage data across all calls to produce a single Usage value on the final response + // usage data across all calls to produce a single Usage value on the final response. + // The FunctionInvokingChatClient then itself creates a span that will also be tagged with a sum + // across all consituent calls, which means our final answer will be double. if (response.Usage is { } finalUsage) { var totalInputTokens = activities.Sum(a => (int?)a.GetTagItem("gen_ai.response.input_tokens")!); var totalOutputTokens = activities.Sum(a => (int?)a.GetTagItem("gen_ai.response.output_tokens")!); - Assert.Equal(totalInputTokens, finalUsage.InputTokenCount); - Assert.Equal(totalOutputTokens, finalUsage.OutputTokenCount); + Assert.Equal(totalInputTokens, finalUsage.InputTokenCount * 2); + Assert.Equal(totalOutputTokens, finalUsage.OutputTokenCount * 2); } }