From 9f1a0ded7af5ebbffecbe7c075c55b18c78cc4e6 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Fri, 6 Jun 2025 21:49:44 +0100 Subject: [PATCH] Add Usage Metadata for ChatClientChatCOmpletionService + UT --- .../samples/AgentUtilities/BaseAgentsTest.cs | 6 +- .../AI/ChatClient/ChatMessageExtensions.cs | 2 +- .../ChatResponseUpdateExtensions.cs | 10 + ...entChatCompletionServiceConversionTests.cs | 455 ++++++++++++++++++ .../Functions/KernelFunctionTests.cs | 3 +- 5 files changed, 473 insertions(+), 3 deletions(-) create mode 100644 dotnet/src/SemanticKernel.UnitTests/AI/ChatCompletion/ChatClientChatCompletionServiceConversionTests.cs diff --git a/dotnet/src/InternalUtilities/samples/AgentUtilities/BaseAgentsTest.cs b/dotnet/src/InternalUtilities/samples/AgentUtilities/BaseAgentsTest.cs index 0b40625eaa16..f6bdd7669ac3 100644 --- a/dotnet/src/InternalUtilities/samples/AgentUtilities/BaseAgentsTest.cs +++ b/dotnet/src/InternalUtilities/samples/AgentUtilities/BaseAgentsTest.cs @@ -9,8 +9,8 @@ using Microsoft.SemanticKernel.ChatCompletion; using OpenAI.Assistants; using OpenAI.Files; - using ChatTokenUsage = OpenAI.Chat.ChatTokenUsage; +using UsageDetails = Microsoft.Extensions.AI.UsageDetails; /// /// Base class for samples that demonstrate the usage of host agents @@ -125,6 +125,10 @@ protected void WriteAgentChatMessage(ChatMessageContent message) { WriteUsage(chatUsage.TotalTokenCount, chatUsage.InputTokenCount, chatUsage.OutputTokenCount); } + else if (usage is UsageDetails usageDetails) + { + WriteUsage(usageDetails.TotalTokenCount ?? 0, usageDetails.InputTokenCount ?? 0, usageDetails.OutputTokenCount ?? 0); + } } string FormatAuthor() => message.AuthorName is not null ? $" - {message.AuthorName ?? " * "}" : string.Empty; diff --git a/dotnet/src/SemanticKernel.Abstractions/AI/ChatClient/ChatMessageExtensions.cs b/dotnet/src/SemanticKernel.Abstractions/AI/ChatClient/ChatMessageExtensions.cs index 1501cb71d988..561838ceb660 100644 --- a/dotnet/src/SemanticKernel.Abstractions/AI/ChatClient/ChatMessageExtensions.cs +++ b/dotnet/src/SemanticKernel.Abstractions/AI/ChatClient/ChatMessageExtensions.cs @@ -15,7 +15,7 @@ internal static ChatMessageContent ToChatMessageContent(this ChatMessage message ModelId = response?.ModelId, AuthorName = message.AuthorName, InnerContent = response?.RawRepresentation ?? message.RawRepresentation, - Metadata = message.AdditionalProperties, + Metadata = new AdditionalPropertiesDictionary(message.AdditionalProperties ?? []) { ["Usage"] = response?.Usage }, Role = new AuthorRole(message.Role.Value), }; diff --git a/dotnet/src/SemanticKernel.Abstractions/AI/ChatClient/ChatResponseUpdateExtensions.cs b/dotnet/src/SemanticKernel.Abstractions/AI/ChatClient/ChatResponseUpdateExtensions.cs index 8ec9698484b0..614de18c43cb 100644 --- a/dotnet/src/SemanticKernel.Abstractions/AI/ChatClient/ChatResponseUpdateExtensions.cs +++ b/dotnet/src/SemanticKernel.Abstractions/AI/ChatClient/ChatResponseUpdateExtensions.cs @@ -1,5 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. +using System.Collections.Generic; using System.Text.Json; using Microsoft.SemanticKernel; using Microsoft.SemanticKernel.ChatCompletion; @@ -34,9 +35,18 @@ item is Microsoft.Extensions.AI.FunctionCallContent fcc ? if (resultContent is not null) { + resultContent.InnerContent = item.RawRepresentation; resultContent.ModelId = update.ModelId; content.Items.Add(resultContent); } + + if (item is Microsoft.Extensions.AI.UsageContent uc) + { + content.Metadata = new Dictionary(update.AdditionalProperties ?? []) + { + ["Usage"] = uc + }; + } } return content; diff --git a/dotnet/src/SemanticKernel.UnitTests/AI/ChatCompletion/ChatClientChatCompletionServiceConversionTests.cs b/dotnet/src/SemanticKernel.UnitTests/AI/ChatCompletion/ChatClientChatCompletionServiceConversionTests.cs new file mode 100644 index 000000000000..703d01a95d99 --- /dev/null +++ b/dotnet/src/SemanticKernel.UnitTests/AI/ChatCompletion/ChatClientChatCompletionServiceConversionTests.cs @@ -0,0 +1,455 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Microsoft.SemanticKernel; +using Microsoft.SemanticKernel.ChatCompletion; +using Xunit; + +namespace SemanticKernel.UnitTests.AI.ChatCompletion; + +/// +/// Unit tests for ChatClientChatCompletionService conversion logic. +/// Tests verify that metadata and usage content are properly preserved when converting +/// from IChatClient abstractions to Semantic Kernel types. +/// +public sealed class ChatClientChatCompletionServiceConversionTests +{ + [Fact] + public async Task GetChatMessageContentsAsyncWithUsageDetailsPreservesUsageInMetadata() + { + // Arrange + using var chatClient = new TestChatClient + { + CompleteAsyncDelegate = (messages, options, cancellationToken) => + { + return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Test response")]) + { + Usage = new UsageDetails { InputTokenCount = 10, OutputTokenCount = 20, TotalTokenCount = 30 }, + ModelId = "test-model", + RawRepresentation = "raw-response" + }); + } + }; + + var service = chatClient.AsChatCompletionService(); + var chatHistory = new ChatHistory(); + chatHistory.AddUserMessage("Test message"); + + // Act + var result = await service.GetChatMessageContentsAsync(chatHistory); + + // Assert + Assert.Single(result); + var message = result[0]; + Assert.NotNull(message.Metadata); + Assert.True(message.Metadata.ContainsKey("Usage")); + var usageDetails = Assert.IsType(message.Metadata["Usage"]); + Assert.Equal(10, usageDetails.InputTokenCount); + Assert.Equal(20, usageDetails.OutputTokenCount); + Assert.Equal(30, usageDetails.TotalTokenCount); + Assert.Equal("test-model", message.ModelId); + Assert.Equal("raw-response", message.InnerContent); + } + + [Fact] + public async Task GetChatMessageContentsAsyncWithoutUsageHasNullUsageInMetadata() + { + // Arrange + using var chatClient = new TestChatClient + { + CompleteAsyncDelegate = (messages, options, cancellationToken) => + { + return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Test response")]) + { + Usage = null, + ModelId = "test-model" + }); + } + }; + + var service = chatClient.AsChatCompletionService(); + var chatHistory = new ChatHistory(); + chatHistory.AddUserMessage("Test message"); + + // Act + var result = await service.GetChatMessageContentsAsync(chatHistory); + + // Assert + Assert.Single(result); + var message = result[0]; + Assert.NotNull(message.Metadata); + Assert.True(message.Metadata.ContainsKey("Usage")); + Assert.Null(message.Metadata["Usage"]); + } + + [Fact] + public async Task GetStreamingChatMessageContentsAsyncWithUsageContentPreservesUsageInMetadata() + { + // Arrange + var expectedUsage = new UsageContent(new UsageDetails { InputTokenCount = 5, OutputTokenCount = 10, TotalTokenCount = 15 }); + using var chatClient = new TestChatClient + { + CompleteStreamingAsyncDelegate = (messages, options, cancellationToken) => + { + return new[] + { + new ChatResponseUpdate(ChatRole.Assistant, "Hello"), + new ChatResponseUpdate(ChatRole.Assistant, " World") { Contents = [expectedUsage] }, + new ChatResponseUpdate(ChatRole.Assistant, "!") { ModelId = "test-model" } + }.ToAsyncEnumerable(); + } + }; + + var service = chatClient.AsChatCompletionService(); + var chatHistory = new ChatHistory(); + chatHistory.AddUserMessage("Test message"); + + // Act + var results = new List(); + await foreach (var update in service.GetStreamingChatMessageContentsAsync(chatHistory)) + { + results.Add(update); + } + + // Assert + Assert.Equal(3, results.Count); + + // Check the update with usage content + var usageUpdate = results[1]; + Assert.NotNull(usageUpdate.Metadata); + Assert.True(usageUpdate.Metadata.ContainsKey("Usage")); + Assert.Equal(expectedUsage, usageUpdate.Metadata["Usage"]); + + // Check model ID is preserved + var modelUpdate = results[2]; + Assert.Equal("test-model", modelUpdate.ModelId); + } + + [Fact] + public async Task GetStreamingChatMessageContentsAsyncWithInnerContentPreservesInnerContent() + { + // Arrange + using var chatClient = new TestChatClient + { + CompleteStreamingAsyncDelegate = (messages, options, cancellationToken) => + { + return new[] + { + new ChatResponseUpdate(ChatRole.Assistant, "Test") + { + RawRepresentation = "raw-stream-data", + ModelId = "test-model" + } + }.ToAsyncEnumerable(); + } + }; + + var service = chatClient.AsChatCompletionService(); + var chatHistory = new ChatHistory(); + chatHistory.AddUserMessage("Test message"); + + // Act + var results = new List(); + await foreach (var update in service.GetStreamingChatMessageContentsAsync(chatHistory)) + { + results.Add(update); + } + + // Assert + Assert.Single(results); + var message = results[0]; + Assert.Equal("raw-stream-data", message.InnerContent); + Assert.Equal("test-model", message.ModelId); + } + + [Fact] + public async Task GetChatMessageContentsAsyncWithAdditionalPropertiesPreservesMetadata() + { + // Arrange + var additionalProps = new AdditionalPropertiesDictionary + { + ["custom-key"] = "custom-value", + ["another-key"] = 42 + }; + + using var chatClient = new TestChatClient + { + CompleteAsyncDelegate = (messages, options, cancellationToken) => + { + var message = new ChatMessage(ChatRole.Assistant, "Test response") + { + AdditionalProperties = additionalProps + }; + return Task.FromResult(new ChatResponse([message]) + { + Usage = new UsageDetails { InputTokenCount = 5, OutputTokenCount = 15, TotalTokenCount = 20 } + }); + } + }; + + var service = chatClient.AsChatCompletionService(); + var chatHistory = new ChatHistory(); + chatHistory.AddUserMessage("Test message"); + + // Act + var result = await service.GetChatMessageContentsAsync(chatHistory); + + // Assert + Assert.Single(result); + var message = result[0]; + Assert.NotNull(message.Metadata); + Assert.True(message.Metadata.ContainsKey("Usage")); + var usageDetails = Assert.IsType(message.Metadata["Usage"]); + Assert.Equal(5, usageDetails.InputTokenCount); + Assert.Equal(15, usageDetails.OutputTokenCount); + Assert.Equal(20, usageDetails.TotalTokenCount); + Assert.True(message.Metadata.ContainsKey("custom-key")); + Assert.Equal("custom-value", message.Metadata["custom-key"]); + Assert.True(message.Metadata.ContainsKey("another-key")); + Assert.Equal(42, message.Metadata["another-key"]); + } + + [Fact] + public async Task GetStreamingChatMessageContentsAsyncWithAdditionalPropertiesPreservesMetadata() + { + // Arrange + var additionalProps = new AdditionalPropertiesDictionary + { + ["custom-key"] = "custom-value", + ["stream-id"] = "stream-123" + }; + + using var chatClient = new TestChatClient + { + CompleteStreamingAsyncDelegate = (messages, options, cancellationToken) => + { + return new[] + { + new ChatResponseUpdate(ChatRole.Assistant, "Test") + { + AdditionalProperties = additionalProps, + RawRepresentation = "raw-stream-data" + } + }.ToAsyncEnumerable(); + } + }; + + var service = chatClient.AsChatCompletionService(); + var chatHistory = new ChatHistory(); + chatHistory.AddUserMessage("Test message"); + + // Act + var results = new List(); + await foreach (var update in service.GetStreamingChatMessageContentsAsync(chatHistory)) + { + results.Add(update); + } + + // Assert + Assert.Single(results); + var message = results[0]; + Assert.NotNull(message.Metadata); + Assert.True(message.Metadata.ContainsKey("custom-key")); + Assert.Equal("custom-value", message.Metadata["custom-key"]); + Assert.True(message.Metadata.ContainsKey("stream-id")); + Assert.Equal("stream-123", message.Metadata["stream-id"]); + Assert.Equal("raw-stream-data", message.InnerContent); + } + + [Fact] + public async Task GetChatMessageContentsAsyncWithFunctionCallContentPreservesInnerContentAndMetadata() + { + // Arrange + var functionCall = new Microsoft.Extensions.AI.FunctionCallContent("call-456", "WeatherFunction", + new Dictionary { ["location"] = "Seattle", ["units"] = "metric" }) + { + RawRepresentation = "function-call-raw" + }; + + using var chatClient = new TestChatClient + { + CompleteAsyncDelegate = (messages, options, cancellationToken) => + { + var message = new ChatMessage(ChatRole.Assistant, [functionCall]) + { + RawRepresentation = "message-raw-content" + }; + return Task.FromResult(new ChatResponse([message]) + { + Usage = new UsageDetails { InputTokenCount = 15, OutputTokenCount = 25, TotalTokenCount = 40 }, + ModelId = "function-model", + RawRepresentation = "response-raw" + }); + } + }; + + var service = chatClient.AsChatCompletionService(); + var chatHistory = new ChatHistory(); + chatHistory.AddUserMessage("What's the weather?"); + + // Act + var result = await service.GetChatMessageContentsAsync(chatHistory); + + // Assert + Assert.Single(result); + var message = result[0]; + Assert.Equal("response-raw", message.InnerContent); + Assert.Equal("function-model", message.ModelId); + Assert.NotNull(message.Metadata); + Assert.True(message.Metadata.ContainsKey("Usage")); + var usageDetails = Assert.IsType(message.Metadata["Usage"]); + Assert.Equal(15, usageDetails.InputTokenCount); + Assert.Equal(25, usageDetails.OutputTokenCount); + Assert.Equal(40, usageDetails.TotalTokenCount); + + Assert.Single(message.Items); + var functionCallContent = Assert.IsType(message.Items[0]); + Assert.Equal("call-456", functionCallContent.Id); + Assert.Equal("WeatherFunction", functionCallContent.FunctionName); + Assert.Equal("function-call-raw", functionCallContent.InnerContent); + Assert.Equal("function-model", functionCallContent.ModelId); + } + + [Fact] + public async Task GetStreamingChatMessageContentsAsyncWithFunctionCallsPreservesInnerContent() + { + // Arrange + var functionCall = new Microsoft.Extensions.AI.FunctionCallContent("call-123", "TestFunction", + new Dictionary { ["param"] = "value" }) + { + RawRepresentation = "function-raw-data" + }; + + using var chatClient = new TestChatClient + { + CompleteStreamingAsyncDelegate = (messages, options, cancellationToken) => + { + return new[] + { + new ChatResponseUpdate(ChatRole.Assistant, [functionCall]) + { + ModelId = "test-model", + RawRepresentation = "update-raw-data" + } + }.ToAsyncEnumerable(); + } + }; + + var service = chatClient.AsChatCompletionService(); + var chatHistory = new ChatHistory(); + chatHistory.AddUserMessage("Test message"); + + // Act + var results = new List(); + await foreach (var update in service.GetStreamingChatMessageContentsAsync(chatHistory)) + { + results.Add(update); + } + + // Assert + Assert.Single(results); + var message = results[0]; + Assert.Equal("update-raw-data", message.InnerContent); + Assert.Equal("test-model", message.ModelId); + Assert.Single(message.Items); + + var functionCallItem = Assert.IsType(message.Items[0]); + Assert.Equal("call-123", functionCallItem.CallId); + Assert.Equal("TestFunction", functionCallItem.Name); + Assert.Equal("function-raw-data", functionCallItem.InnerContent); + Assert.Equal("test-model", functionCallItem.ModelId); + } + + [Fact] + public async Task GetStreamingChatMessageContentsAsyncWithTextAndUsageContentCreatesCorrectStreamingContent() + { + // Arrange + var expectedUsage = new UsageContent(new UsageDetails { InputTokenCount = 8, OutputTokenCount = 12, TotalTokenCount = 20 }); + var textContent = new Microsoft.Extensions.AI.TextContent("Hello World"); + + using var chatClient = new TestChatClient + { + CompleteStreamingAsyncDelegate = (messages, options, cancellationToken) => + { + return new[] + { + new ChatResponseUpdate(ChatRole.Assistant, [textContent, expectedUsage]) + { + ModelId = "test-model", + RawRepresentation = "combined-content-raw" + } + }.ToAsyncEnumerable(); + } + }; + + var service = chatClient.AsChatCompletionService(); + var chatHistory = new ChatHistory(); + chatHistory.AddUserMessage("Test message"); + + // Act + var results = new List(); + await foreach (var update in service.GetStreamingChatMessageContentsAsync(chatHistory)) + { + results.Add(update); + } + + // Assert + Assert.Single(results); + var message = results[0]; + + // Should have only text content as streaming item, usage goes to metadata + Assert.Single(message.Items); + + // Check text content + var streamingTextContent = Assert.IsType(message.Items[0]); + Assert.Equal("Hello World", streamingTextContent.Text); + Assert.Equal("test-model", streamingTextContent.ModelId); + + // Check overall message metadata - usage content should be in metadata + Assert.NotNull(message.Metadata); + Assert.True(message.Metadata.ContainsKey("Usage")); + Assert.Equal(expectedUsage, message.Metadata["Usage"]); + Assert.Equal("combined-content-raw", message.InnerContent); + Assert.Equal("test-model", message.ModelId); + } + + /// + /// Test implementation of IChatClient for unit testing. + /// + private sealed class TestChatClient : IChatClient + { + public Func, ChatOptions?, CancellationToken, Task>? CompleteAsyncDelegate { get; set; } + public Func, ChatOptions?, CancellationToken, IAsyncEnumerable>? CompleteStreamingAsyncDelegate { get; set; } + + public ChatClientMetadata Metadata { get; set; } = new("TestChatClient", null, "test-model"); + + public Task GetResponseAsync(IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + { + return this.CompleteAsyncDelegate?.Invoke(messages, options, cancellationToken) + ?? throw new NotImplementedException("CompleteAsyncDelegate not set"); + } + + public IAsyncEnumerable GetStreamingResponseAsync(IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + { + return this.CompleteStreamingAsyncDelegate?.Invoke(messages, options, cancellationToken) + ?? throw new NotImplementedException("CompleteStreamingAsyncDelegate not set"); + } + + public TService? GetService(object? key = null) where TService : class + { + return typeof(TService) == typeof(ChatClientMetadata) ? (TService)(object)this.Metadata : null; + } + + public object? GetService(Type serviceType, object? serviceKey = null) + { + return serviceType == typeof(ChatClientMetadata) ? this.Metadata : null; + } + + public void Dispose() { } + } +} diff --git a/dotnet/src/SemanticKernel.UnitTests/Functions/KernelFunctionTests.cs b/dotnet/src/SemanticKernel.UnitTests/Functions/KernelFunctionTests.cs index 42e8f9b61ded..102fb4d38b5f 100644 --- a/dotnet/src/SemanticKernel.UnitTests/Functions/KernelFunctionTests.cs +++ b/dotnet/src/SemanticKernel.UnitTests/Functions/KernelFunctionTests.cs @@ -3,10 +3,11 @@ using System; using System.Threading.Tasks; using Microsoft.Extensions.Logging; +using Microsoft.SemanticKernel; using Moq; using Xunit; -namespace Microsoft.SemanticKernel.UnitTests.Functions; +namespace SemanticKernel.UnitTests.Functions; /// /// Tests for cloning with a instance.