Mock Metadata Usage in a GetStreamingChatMessageContentsAsync response

[romainsemur]

Hi,

In this article about unit test (https://devblogs.microsoft.com/semantic-kernel/unit-testing-with-semantic-kernel/) , in the DoWorkWithPrompt() method, how to add a metadata with Usage values in the returns ?
.Returns(new List() { new(AuthorRole.Assistant, "AI response") }.ToAsyncEnumerable());

[sophialagerkranspandey]

Tagging @dmytrostruk

[dmytrostruk]

Hi @romainsemur ,
Here is an updated example with metadata including Usage in return:

[Fact]
public async Task DoWorkWithPrompt()
{
    // Arrange 
    var mockChatCompletion = new Mock<IChatCompletionService>();

    var mockResponse = new ChatMessageContent(
            AuthorRole.Assistant,
            "AI response",
            metadata: new Dictionary<string, object?>
            {
                { "Usage",
                    new Dictionary<string, int>
                    {
                        { "InputTokenCount", 18 },
                        { "OutputTokenCount", 242 },
                        { "TotalTokenCount", 260 }
                    }
                }
            });

    mockChatCompletion
        .Setup(x => x.GetChatMessageContentsAsync(
            It.IsAny<ChatHistory>(),
            It.IsAny<PromptExecutionSettings>(),
            It.IsAny<Kernel>(),
            It.IsAny<CancellationToken>()))
        .ReturnsAsync([mockResponse]);

    var kernelBuilder = Kernel.CreateBuilder();
    kernelBuilder.Services.AddSingleton(mockChatCompletion.Object);

    var kernel = kernelBuilder.Build();
    var service = new MyService(kernel);

    // Act 
    var result = await service.DoWorkAsync("Prompt to AI");

    // Assert 
    Assert.Equal("AI response", result.ToString());

    if (result.Metadata is not null &&
        result.Metadata.TryGetValue("Usage", out var metadataObj) &&
        metadataObj is Dictionary<string, int> metadata)
    {
        Assert.Equal(18, metadata["InputTokenCount"]);
        Assert.Equal(242, metadata["OutputTokenCount"]);
        Assert.Equal(260, metadata["TotalTokenCount"]);
    }
}

Please let us know if that works for you, thanks!

[romainsemur]

Hi @dmytrostruk , thanks for your reply.
However it's not working because metadata "Usage" is of type OpenAI.Chat.ChatTokenUsage.
As a reminder, I use GetStreamingChatMessageContentsAsync from IChatCompletionService (from Microsoft.SemanticKernel.Abstractions, Version=1.36.1.0)

[dmytrostruk]

@romainsemur Got it, thanks for the additional details. In this case, taking into account that OpenAI.Chat.ChatTokenUsage model doesn't have public constructor, you can use OpenAIChatModelFactory.ChatTokenUsage() method to initialize it on your side for testing purposes. Here is a full code example including streaming scenario:

[Fact]
public async Task DoWorkWithPrompt()
{
    var mockUsage = OpenAIChatModelFactory.ChatTokenUsage(
        outputTokenCount: 242,
        inputTokenCount: 18,
        totalTokenCount: 260);

    var mockResponse = new StreamingChatMessageContent(
        AuthorRole.Assistant,
        "AI response",
        metadata: new Dictionary<string, object?>
        {
            { "Usage", mockUsage }
        });

    // Arrange 
    var mockChatCompletion = new Mock<IChatCompletionService>();
    mockChatCompletion
        .Setup(x => x.GetStreamingChatMessageContentsAsync(
            It.IsAny<ChatHistory>(),
            It.IsAny<PromptExecutionSettings>(),
            It.IsAny<Kernel>(),
            It.IsAny<CancellationToken>()))
        .Returns(new List<StreamingChatMessageContent>() { mockResponse }.ToAsyncEnumerable());

    var kernelBuilder = Kernel.CreateBuilder();
    kernelBuilder.Services.AddSingleton(mockChatCompletion.Object);

    var kernel = kernelBuilder.Build();
    var service = new MyService(kernel);

    // Act 
    var result = new List<StreamingChatMessageContent>();
    await foreach (var item in service.DoWorkStreaming<StreamingChatMessageContent>("Prompt to AI"))
    {
        result.Add(item);
    }

    var streamingChatMessageContent = result[0];

    // Assert 
    Assert.Equal("AI response", streamingChatMessageContent.Content);

    if (streamingChatMessageContent.Metadata is not null &&
        streamingChatMessageContent.Metadata.TryGetValue("Usage", out var metadataObj) &&
        metadataObj is ChatTokenUsage usage)
    {
        Assert.Equal(18, usage.InputTokenCount);
        Assert.Equal(242, usage.OutputTokenCount);
        Assert.Equal(260, usage.TotalTokenCount);
    }
}

[romainsemur]

Awesome ! Thanks a lot @dmytrostruk , I didn't know that factory 👍