diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatCompletion/ChatResponseExtensions.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatCompletion/ChatResponseExtensions.cs index 34e882a77aa..62bab814d13 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatCompletion/ChatResponseExtensions.cs +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatCompletion/ChatResponseExtensions.cs @@ -536,6 +536,34 @@ private static void ProcessUpdate(ChatResponseUpdate update, ChatResponse respon message.MessageId = update.MessageId; } + // AdditionalProperties are scoped to the message if the update has a MessageId, + // otherwise they're scoped to the response. + if (update.AdditionalProperties is not null) + { + if (update.MessageId is { Length: > 0 }) + { + if (message.AdditionalProperties is null) + { + message.AdditionalProperties = new(update.AdditionalProperties); + } + else + { + message.AdditionalProperties.SetAll(update.AdditionalProperties); + } + } + else + { + if (response.AdditionalProperties is null) + { + response.AdditionalProperties = new(update.AdditionalProperties); + } + else + { + response.AdditionalProperties.SetAll(update.AdditionalProperties); + } + } + } + foreach (var content in update.Contents) { switch (content) @@ -579,18 +607,6 @@ private static void ProcessUpdate(ChatResponseUpdate update, ChatResponse respon { response.ModelId = update.ModelId; } - - if (update.AdditionalProperties is not null) - { - if (response.AdditionalProperties is null) - { - response.AdditionalProperties = new(update.AdditionalProperties); - } - else - { - response.AdditionalProperties.SetAll(update.AdditionalProperties); - } - } } /// Gets whether both strings are not null/empty and not the same as each other. 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 a921f75d580..8cfd0639dd9 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatCompletion/ChatResponseUpdateExtensionsTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatCompletion/ChatResponseUpdateExtensionsTests.cs @@ -445,6 +445,91 @@ await YieldAsync(updates).ToChatResponseAsync() : Assert.Equal(ChatRole.Assistant, response.Messages[2].Role); } + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ToChatResponse_AdditionalPropertiesGoToMessages(bool useAsync) + { + ChatResponseUpdate[] updates = + [ + + // First message with AdditionalProperties (MessageId makes properties go to message) + new(ChatRole.Assistant, "First message") { MessageId = "msg1", AdditionalProperties = new() { ["key1"] = "value1" } }, + new(null, " part 2") { MessageId = "msg1", AdditionalProperties = new() { ["key2"] = "value2" } }, + + // Second message with different AdditionalProperties (same keys, different values) + new(ChatRole.User, "Second message") { MessageId = "msg2", AdditionalProperties = new() { ["key1"] = "different_value1" } }, + new(null, " part 2") { MessageId = "msg2", AdditionalProperties = new() { ["key3"] = "value3" } }, + + // Third message with no AdditionalProperties + new(ChatRole.Assistant, "Third message") { MessageId = "msg3" }, + ]; + + ChatResponse response = useAsync ? + await YieldAsync(updates).ToChatResponseAsync() : + updates.ToChatResponse(); + + Assert.Equal(3, response.Messages.Count); + Assert.Null(response.AdditionalProperties); + + // First message should have its own AdditionalProperties + var msg1 = response.Messages[0]; + Assert.Equal("First message part 2", msg1.Text); + Assert.NotNull(msg1.AdditionalProperties); + Assert.Equal(2, msg1.AdditionalProperties.Count); + Assert.Equal("value1", msg1.AdditionalProperties["key1"]); + Assert.Equal("value2", msg1.AdditionalProperties["key2"]); + + // Second message should have its own AdditionalProperties (with different value for key1) + var msg2 = response.Messages[1]; + Assert.Equal("Second message part 2", msg2.Text); + Assert.NotNull(msg2.AdditionalProperties); + Assert.Equal(2, msg2.AdditionalProperties.Count); + Assert.Equal("different_value1", msg2.AdditionalProperties["key1"]); + Assert.Equal("value3", msg2.AdditionalProperties["key3"]); + + // Third message should have no AdditionalProperties + var msg3 = response.Messages[2]; + Assert.Equal("Third message", msg3.Text); + Assert.Null(msg3.AdditionalProperties); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ToChatResponse_AdditionalPropertiesRoutingBasedOnMessageId(bool useAsync) + { + // This test explicitly verifies that: + // - Updates WITH MessageId route AdditionalProperties to the message + // - Updates WITHOUT MessageId route AdditionalProperties to the response + ChatResponseUpdate[] updates = + [ + + // Update with MessageId - properties should go to message + new(ChatRole.Assistant, "Hello") { MessageId = "msg1", AdditionalProperties = new() { ["messageKey"] = "messageValue" } }, + + // Update without MessageId - properties should go to response + new() { AdditionalProperties = new() { ["responseKey"] = "responseValue" } }, + ]; + + ChatResponse response = useAsync ? + await YieldAsync(updates).ToChatResponseAsync() : + updates.ToChatResponse(); + + // Verify message-scoped properties (update had MessageId) + var message = Assert.Single(response.Messages); + Assert.NotNull(message.AdditionalProperties); + Assert.Single(message.AdditionalProperties); + Assert.Equal("messageValue", message.AdditionalProperties["messageKey"]); + Assert.False(message.AdditionalProperties.ContainsKey("responseKey")); + + // Verify response-scoped properties (update had no MessageId) + Assert.NotNull(response.AdditionalProperties); + Assert.Single(response.AdditionalProperties); + Assert.Equal("responseValue", response.AdditionalProperties["responseKey"]); + Assert.False(response.AdditionalProperties.ContainsKey("messageKey")); + } + [Theory] [InlineData(false)] [InlineData(true)]