Coalesce TextReasoningContent

Author: stephentoub

src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatCompletion/ChatResponseExtensions.cs

@@ -180,53 +180,60 @@ static async Task<ChatResponse> ToChatResponseAsync(
         }
     }
 
-    /// <summary>Coalesces sequential <see cref="TextContent"/> content elements.</summary>
+    /// <summary>Coalesces sequential <see cref="AIContent"/> content elements.</summary>
     internal static void CoalesceTextContent(List<AIContent> contents)
     {
-        StringBuilder? coalescedText = null;
+        Coalesce<TextContent>(contents, static text => new(text));
+        Coalesce<TextReasoningContent>(contents, static text => new(text));
 
-        // Iterate through all of the items in the list looking for contiguous items that can be coalesced.
-        int start = 0;
-        while (start < contents.Count - 1)
+        // This implementation relies on TContent's ToString returning its exact text.
+        static void Coalesce<TContent>(List<AIContent> contents, Func<string, TContent> fromText)
+            where TContent : AIContent
         {
-            // We need at least two TextContents in a row to be able to coalesce.
-            if (contents[start] is not TextContent firstText)
-            {
-                start++;
-                continue;
-            }
-
-            if (contents[start + 1] is not TextContent secondText)
-            {
-                start += 2;
-                continue;
-            }
+            StringBuilder? coalescedText = null;
 
-            // Append the text from those nodes and continue appending subsequent TextContents until we run out.
-            // We null out nodes as their text is appended so that we can later remove them all in one O(N) operation.
-            coalescedText ??= new();
-            _ = coalescedText.Clear().Append(firstText.Text).Append(secondText.Text);
-            contents[start + 1] = null!;
-            int i = start + 2;
-            for (; i < contents.Count && contents[i] is TextContent next; i++)
+            // Iterate through all of the items in the list looking for contiguous items that can be coalesced.
+            int start = 0;
+            while (start < contents.Count - 1)
             {
-                _ = coalescedText.Append(next.Text);
-                contents[i] = null!;
+                // We need at least two TextContents in a row to be able to coalesce.
+                if (contents[start] is not TContent firstText)
+                {
+                    start++;
+                    continue;
+                }
+
+                if (contents[start + 1] is not TContent secondText)
+                {
+                    start += 2;
+                    continue;
+                }
+
+                // Append the text from those nodes and continue appending subsequent TextContents until we run out.
+                // We null out nodes as their text is appended so that we can later remove them all in one O(N) operation.
+                coalescedText ??= new();
+                _ = coalescedText.Clear().Append(firstText).Append(secondText);
+                contents[start + 1] = null!;
+                int i = start + 2;
+                for (; i < contents.Count && contents[i] is TContent next; i++)
+                {
+                    _ = coalescedText.Append(next);
+                    contents[i] = null!;
+                }
+
+                // Store the replacement node. We inherit the properties of the first text node. We don't
+                // currently propagate additional properties from the subsequent nodes. If we ever need to,
+                // we can add that here.
+                var newContent = fromText(coalescedText.ToString());
+                contents[start] = newContent;
+                newContent.AdditionalProperties = firstText.AdditionalProperties?.Clone();
+
+                start = i;
             }
 
-            // Store the replacement node.
-            contents[start] = new TextContent(coalescedText.ToString())
-            {
-                // We inherit the properties of the first text node. We don't currently propagate additional
-                // properties from the subsequent nodes. If we ever need to, we can add that here.
-                AdditionalProperties = firstText.AdditionalProperties?.Clone(),
-            };
-
-            start = i;
+            // Remove all of the null slots left over from the coalescing process.
+            _ = contents.RemoveAll(u => u is null);
         }
-
-        // Remove all of the null slots left over from the coalescing process.
-        _ = contents.RemoveAll(u => u is null);
     }
 
     /// <summary>Finalizes the <paramref name="response"/> object.</summary>

test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatCompletion/ChatResponseUpdateExtensionsTests.cs

@@ -145,6 +145,44 @@ void AddGap()
         }
     }
 
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public async Task ToChatResponse_CoalescesTextContentAndTextReasoningContentSeparately(bool useAsync)
+    {
+        ChatResponseUpdate[] updates =
+        {
+            new(null, "A"),
+            new(null, "B"),
+            new(null, "C"),
+            new() { Contents = [new TextReasoningContent("D")] },
+            new() { Contents = [new TextReasoningContent("E")] },
+            new() { Contents = [new TextReasoningContent("F")] },
+            new(null, "G"),
+            new(null, "H"),
+            new() { Contents = [new TextReasoningContent("I")] },
+            new() { Contents = [new TextReasoningContent("J")] },
+            new(null, "K"),
+            new() { Contents = [new TextReasoningContent("L")] },
+            new(null, "M"),
+            new(null, "N"),
+            new() { Contents = [new TextReasoningContent("O")] },
+            new() { Contents = [new TextReasoningContent("P")] },
+        };
+
+        ChatResponse response = useAsync ? await YieldAsync(updates).ToChatResponseAsync() : updates.ToChatResponse();
+        ChatMessage message = Assert.Single(response.Messages);
+        Assert.Equal(8, message.Contents.Count);
+        Assert.Equal("ABC", Assert.IsType<TextContent>(message.Contents[0]).Text);
+        Assert.Equal("DEF", Assert.IsType<TextReasoningContent>(message.Contents[1]).Text);
+        Assert.Equal("GH", Assert.IsType<TextContent>(message.Contents[2]).Text);
+        Assert.Equal("IJ", Assert.IsType<TextReasoningContent>(message.Contents[3]).Text);
+        Assert.Equal("K", Assert.IsType<TextContent>(message.Contents[4]).Text);
+        Assert.Equal("L", Assert.IsType<TextReasoningContent>(message.Contents[5]).Text);
+        Assert.Equal("MN", Assert.IsType<TextContent>(message.Contents[6]).Text);
+        Assert.Equal("OP", Assert.IsType<TextReasoningContent>(message.Contents[7]).Text);
+    }
+
     [Fact]
     public async Task ToChatResponse_UsageContentExtractedFromContents()
     {