From acd880e607af6952f8e62a68add7e38117eb4d5e Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Fri, 5 Sep 2025 23:28:14 -0400 Subject: [PATCH] Factor ChatMessage.Role into ChatResponseUpdate coalescing logic If the role changes, consider it a new message. --- .../ChatCompletion/ChatResponseExtensions.cs | 8 ++- .../ChatResponseUpdateExtensionsTests.cs | 60 ++++++++++++++++++- 2 files changed, 65 insertions(+), 3 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatCompletion/ChatResponseExtensions.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatCompletion/ChatResponseExtensions.cs index e41368f115d..b9909857be2 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatCompletion/ChatResponseExtensions.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatCompletion/ChatResponseExtensions.cs @@ -303,7 +303,7 @@ private static void FinalizeResponse(ChatResponse response) private static void ProcessUpdate(ChatResponseUpdate update, ChatResponse response) { // If there is no message created yet, or if the last update we saw had a different - // message ID than the newest update, create a new message. + // message ID or role than the newest update, create a new message. ChatMessage message; var isNewMessage = false; if (response.Messages.Count == 0) @@ -316,6 +316,12 @@ private static void ProcessUpdate(ChatResponseUpdate update, ChatResponse respon { isNewMessage = true; } + else if (update.Role is { } updateRole + && response.Messages[response.Messages.Count - 1].Role is { } lastRole + && updateRole != lastRole) + { + isNewMessage = true; + } if (isNewMessage) { diff --git a/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatCompletion/ChatResponseUpdateExtensionsTests.cs b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatCompletion/ChatResponseUpdateExtensionsTests.cs index 45a82542da8..5cfb7ace025 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatCompletion/ChatResponseUpdateExtensionsTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatCompletion/ChatResponseUpdateExtensionsTests.cs @@ -29,7 +29,7 @@ public async Task ToChatResponse_SuccessfullyCreatesResponse(bool useAsync) ChatResponseUpdate[] updates = [ new(ChatRole.Assistant, "Hello") { ResponseId = "someResponse", MessageId = "12345", CreatedAt = new DateTimeOffset(1, 2, 3, 4, 5, 6, TimeSpan.Zero), ModelId = "model123" }, - new(new("human"), ", ") { AuthorName = "Someone", AdditionalProperties = new() { ["a"] = "b" } }, + new(ChatRole.Assistant, ", ") { AuthorName = "Someone", AdditionalProperties = new() { ["a"] = "b" } }, new(null, "world!") { CreatedAt = new DateTimeOffset(2, 2, 3, 4, 5, 6, TimeSpan.Zero), ConversationId = "123", AdditionalProperties = new() { ["c"] = "d" } }, new() { Contents = [new UsageContent(new() { InputTokenCount = 1, OutputTokenCount = 2 })] }, @@ -53,7 +53,7 @@ public async Task ToChatResponse_SuccessfullyCreatesResponse(bool useAsync) ChatMessage message = response.Messages.Single(); Assert.Equal("12345", message.MessageId); - Assert.Equal(new ChatRole("human"), message.Role); + Assert.Equal(ChatRole.Assistant, message.Role); Assert.Equal("Someone", message.AuthorName); Assert.Null(message.AdditionalProperties); @@ -65,6 +65,62 @@ public async Task ToChatResponse_SuccessfullyCreatesResponse(bool useAsync) Assert.Equal("Hello, world!", response.Text); } + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ToChatResponse_RoleOrIdChangeDictatesMessageChange(bool useAsync) + { + ChatResponseUpdate[] updates = + [ + new(null, "!") { MessageId = "1" }, + new(ChatRole.Assistant, "a") { MessageId = "1" }, + new(ChatRole.Assistant, "b") { MessageId = "2" }, + new(ChatRole.User, "c") { MessageId = "2" }, + new(ChatRole.User, "d") { MessageId = "2" }, + new(ChatRole.Assistant, "e") { MessageId = "3" }, + new(ChatRole.Tool, "f") { MessageId = "4" }, + new(ChatRole.Tool, "g") { MessageId = "4" }, + new(ChatRole.Tool, "h") { MessageId = "5" }, + new(new("human"), "i") { MessageId = "6" }, + new(new("human"), "j") { MessageId = "7" }, + new(new("human"), "k") { MessageId = "7" }, + new(null, "l") { MessageId = "7" }, + new(null, "m") { MessageId = "8" }, + ]; + + ChatResponse response = useAsync ? + updates.ToChatResponse() : + await YieldAsync(updates).ToChatResponseAsync(); + Assert.Equal(9, response.Messages.Count); + + Assert.Equal("!a", response.Messages[0].Text); + Assert.Equal(ChatRole.Assistant, response.Messages[0].Role); + + Assert.Equal("b", response.Messages[1].Text); + Assert.Equal(ChatRole.Assistant, response.Messages[1].Role); + + Assert.Equal("cd", response.Messages[2].Text); + Assert.Equal(ChatRole.User, response.Messages[2].Role); + + Assert.Equal("e", response.Messages[3].Text); + Assert.Equal(ChatRole.Assistant, response.Messages[3].Role); + + Assert.Equal("fg", response.Messages[4].Text); + Assert.Equal(ChatRole.Tool, response.Messages[4].Role); + + Assert.Equal("h", response.Messages[5].Text); + Assert.Equal(ChatRole.Tool, response.Messages[5].Role); + + Assert.Equal("i", response.Messages[6].Text); + Assert.Equal(new ChatRole("human"), response.Messages[6].Role); + + Assert.Equal("jkl", response.Messages[7].Text); + Assert.Equal(new ChatRole("human"), response.Messages[7].Role); + + Assert.Equal("m", response.Messages[8].Text); + Assert.Equal(ChatRole.Assistant, response.Messages[8].Role); + } + [Theory] [InlineData(false)] [InlineData(true)]