Skip to content

Commit 0e845f5

Browse files
stephentoubstephentoub
committed
Address feedback
1 parent 8868320 commit 0e845f5

File tree

3 files changed

+236
-203
lines changed

3 files changed

+236
-203
lines changed

‎src/Libraries/Microsoft.Extensions.AI/ChatCompletion/CachingChatClient.cs‎

Lines changed: 104 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,13 @@
33

44
using System.Collections.Generic;
55
using System.Runtime.CompilerServices;
6+
using System.Text;
67
using System.Threading;
78
using System.Threading.Tasks;
89
using Microsoft.Shared.Diagnostics;
910

11+
#pragma warning disable S127 // "for" loop stop conditions should be invariant
12+
1013
namespace Microsoft.Extensions.AI;
1114

1215
/// <summary>
@@ -83,7 +86,7 @@ public override async IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteSt
8386
// If the caching client is configured to coalesce streaming updates, do so now within the capturedItems list.
8487
if (CoalesceStreamingUpdates)
8588
{
86-
capturedItems.CoalesceTextContent();
89+
CoalesceTextContent(capturedItems);
8790
}
8891

8992
// Write the captured items to the cache.
@@ -137,4 +140,104 @@ public override async IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteSt
137140
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
138141
/// <returns>A <see cref="Task"/> representing the completion of the operation.</returns>
139142
protected abstract Task WriteCacheStreamingAsync(string key, IReadOnlyList<StreamingChatCompletionUpdate> value, CancellationToken cancellationToken);
143+
144+
/// <summary>Coalesces sequential <see cref="TextContent"/> updates.</summary>
145+
private static void CoalesceTextContent(List<StreamingChatCompletionUpdate> updates)
146+
{
147+
StringBuilder coalescedText = new();
148+
149+
// Iterate through all of the items in the list looking for contiguous items that can be coalesced.
150+
for (int startInclusive = 0; startInclusive < updates.Count; startInclusive++)
151+
{
152+
// If an item isn't generally coalescable, skip it.
153+
StreamingChatCompletionUpdate update = updates[startInclusive];
154+
if (update.ChoiceIndex != 0 ||
155+
update.Contents.Count != 1 ||
156+
update.Contents[0] is not TextContent textContent)
157+
{
158+
continue;
159+
}
160+
161+
// We found a coalescable item. Look for more contiguous items that are also coalescable with it.
162+
int endExclusive = startInclusive + 1;
163+
for (; endExclusive < updates.Count; endExclusive++)
164+
{
165+
StreamingChatCompletionUpdate next = updates[endExclusive];
166+
if (next.ChoiceIndex != 0 ||
167+
next.Contents.Count != 1 ||
168+
next.Contents[0] is not TextContent ||
169+
170+
// changing role or author would be really strange, but check anyway
171+
(update.Role is not null && next.Role is not null && update.Role != next.Role) ||
172+
(update.AuthorName is not null && next.AuthorName is not null && update.AuthorName != next.AuthorName))
173+
{
174+
break;
175+
}
176+
}
177+
178+
// If we couldn't find anything to coalesce, there's nothing to do.
179+
if (endExclusive - startInclusive <= 1)
180+
{
181+
continue;
182+
}
183+
184+
// We found a coalescable run of items. Create a new node to represent the run. We create a new one
185+
// rather than reappropriating one of the existing ones so as not to mutate an item already yielded.
186+
_ = coalescedText.Clear().Append(updates[startInclusive].Text);
187+
188+
TextContent coalescedContent = new(null) // will patch the text after examining all items in the run
189+
{
190+
AdditionalProperties = textContent.AdditionalProperties?.Clone(),
191+
};
192+
193+
StreamingChatCompletionUpdate coalesced = new()
194+
{
195+
AdditionalProperties = update.AdditionalProperties?.Clone(),
196+
AuthorName = update.AuthorName,
197+
CompletionId = update.CompletionId,
198+
Contents = [coalescedContent],
199+
CreatedAt = update.CreatedAt,
200+
FinishReason = update.FinishReason,
201+
ModelId = update.ModelId,
202+
Role = update.Role,
203+
204+
// Explicitly don't include RawRepresentation. It's not applicable if one update ends up being used
205+
// to represent multiple, and it won't be serialized anyway.
206+
};
207+
208+
// Replace the starting node with the coalesced node.
209+
updates[startInclusive] = coalesced;
210+
211+
// Now iterate through all the rest of the updates in the run, updating the coalesced node with relevant properties,
212+
// and nulling out the nodes along the way. We do this rather than removing the entry in order to avoid an O(N^2) operation.
213+
// We'll remove all the null entries at the end of the loop, using RemoveAll to do so, which can remove all of
214+
// the nulls in a single O(N) pass.
215+
for (int i = startInclusive + 1; i < endExclusive; i++)
216+
{
217+
// Grab the next item.
218+
StreamingChatCompletionUpdate next = updates[i];
219+
updates[i] = null!;
220+
221+
var nextContent = (TextContent)next.Contents[0];
222+
_ = coalescedText.Append(nextContent.Text);
223+
224+
coalesced.AuthorName ??= next.AuthorName;
225+
coalesced.CompletionId ??= next.CompletionId;
226+
coalesced.CreatedAt ??= next.CreatedAt;
227+
coalesced.FinishReason ??= next.FinishReason;
228+
coalesced.ModelId ??= next.ModelId;
229+
coalesced.Role ??= next.Role;
230+
}
231+
232+
// Complete the coalescing by patching the text of the coalesced node.
233+
coalesced.Text = coalescedText.ToString();
234+
235+
// Jump to the last update in the run, so that when we loop around and bump ahead,
236+
// we're at the next update just after the run.
237+
startInclusive = endExclusive - 1;
238+
}
239+
240+
// Remove all of the null slots left over from the coalescing process.
241+
_ = updates.RemoveAll(u => u is null);
242+
}
140243
}

‎src/Libraries/Microsoft.Extensions.AI/ChatCompletion/StreamingChatCompletionUpdateExtensions.cs‎

Lines changed: 39 additions & 148 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,9 @@
66
using System.Runtime.InteropServices;
77
#endif
88
using System.Text;
9-
using System.Threading.Tasks;
109
using Microsoft.Shared.Diagnostics;
1110

11+
#pragma warning disable S109 // Magic numbers should not be used
1212
#pragma warning disable S127 // "for" loop stop conditions should be invariant
1313

1414
namespace Microsoft.Extensions.AI;
@@ -18,79 +18,17 @@ namespace Microsoft.Extensions.AI;
1818
/// </summary>
1919
public static class StreamingChatCompletionUpdateExtensions
2020
{
21-
/// <summary>
22-
/// Augments an <see cref="IAsyncEnumerable{StreamingChatCompletionUpdate}"/> produced by
23-
/// <see cref="IChatClient.CompleteStreamingAsync"/> to add a merged message to the list of messages
24-
/// upon completion of iterating through the enumerable.
25-
/// </summary>
26-
/// <param name="updates">The list of updates to coalesce into a message added to <paramref name="chatMessages"/>.</param>
27-
/// <param name="chatMessages">The list of messages to which the chat completion message should be added.</param>
21+
/// <summary>Combines <see cref="StreamingChatCompletionUpdate"/> instances into a single <see cref="ChatCompletion"/>.</summary>
22+
/// <param name="updates">The updates to be combined.</param>
2823
/// <param name="coalesceContent">
29-
/// <see langword="true"/> to attempt to coalesce contiguous streaming update <see cref="AIContent"/> items, where applicable,
24+
/// <see langword="true"/> to attempt to coalesce contiguous <see cref="AIContent"/> items, where applicable,
3025
/// into a single <see cref="AIContent"/>, in order to reduce the number of individual content items that are included in
31-
/// manufactured <see cref="ChatMessage"/> instances. When <see langword="false"/>, the content items are kept unaltered.
32-
/// A value of <see langword="true"/> does not impact the updates yielded from the enumerable, only the final message added
33-
/// to <paramref name="chatMessages"/>.
26+
/// the manufactured <see cref="ChatMessage"/> instances. When <see langword="false"/>, the original content items are used.
27+
/// The default is <see langword="true"/>.
3428
/// </param>
35-
/// <returns>
36-
/// An enumerable that yields the same updates as <paramref name="updates"/>. It must be enumerated to completion
37-
/// or disposed to ensure that the final message is added to <paramref name="chatMessages"/>.
38-
/// </returns>
39-
/// <remarks>
40-
/// <para>
41-
/// If any updates were yielded, the method will add a message to <paramref name="chatMessages"/> upon
42-
/// completion of iteration, regardless of whether completion occurs due to exhausting the input enumerable,
43-
/// due to an exception ocurring, or due to early disposal of the enumerator. Each iteration through the
44-
/// returned enumerable may possibly add a message, so it should typically only be enumerated once.
45-
/// </para>
46-
/// <para>
47-
/// At most one <see cref="ChatMessage"/> will be added to <paramref name="chatMessages"/>. If the supplied updates
48-
/// represent multiple choices, only one choice will be added.
49-
/// </para>
50-
/// </remarks>
51-
public static IAsyncEnumerable<StreamingChatCompletionUpdate> WithMessageAddedAsync(
52-
this IAsyncEnumerable<StreamingChatCompletionUpdate> updates,
53-
IList<ChatMessage> chatMessages,
54-
bool coalesceContent = true)
55-
{
56-
_ = Throw.IfNull(updates);
57-
_ = Throw.IfNull(chatMessages);
58-
59-
return ImplAsync(updates, chatMessages, coalesceContent);
60-
61-
static async IAsyncEnumerable<StreamingChatCompletionUpdate> ImplAsync(
62-
IAsyncEnumerable<StreamingChatCompletionUpdate> updates,
63-
IList<ChatMessage> chatMessages,
64-
bool coalesceContent = true)
65-
{
66-
List<StreamingChatCompletionUpdate> updatesList = [];
67-
try
68-
{
69-
await foreach (var update in updates.ConfigureAwait(false))
70-
{
71-
updatesList.Add(update);
72-
yield return update;
73-
}
74-
}
75-
finally
76-
{
77-
if (coalesceContent)
78-
{
79-
CoalesceTextContent(updatesList);
80-
}
81-
82-
if (ToChatCompletion(updatesList) is { Choices.Count: > 0 } completion)
83-
{
84-
chatMessages.Add(completion.Message);
85-
}
86-
}
87-
}
88-
}
89-
90-
/// <summary>Combines <see cref="StreamingChatCompletionUpdate"/> instances into a single <see cref="ChatCompletion"/>.</summary>
91-
/// <param name="updates">The updates to be combined.</param>
9229
/// <returns>The combined <see cref="ChatCompletion"/>.</returns>
93-
public static ChatCompletion ToChatCompletion(this IEnumerable<StreamingChatCompletionUpdate> updates)
30+
public static ChatCompletion ToChatCompletion(
31+
this IEnumerable<StreamingChatCompletionUpdate> updates, bool coalesceContent = true)
9432
{
9533
_ = Throw.IfNull(updates);
9634

@@ -146,6 +84,11 @@ public static ChatCompletion ToChatCompletion(this IEnumerable<StreamingChatComp
14684
entry.Value.Role = ChatRole.Assistant;
14785
}
14886

87+
if (coalesceContent)
88+
{
89+
CoalesceTextContent((List<AIContent>)entry.Value.Contents);
90+
}
91+
14992
completion.Choices.Add(entry.Value);
15093

15194
if (completion.Usage is null)
@@ -164,103 +107,51 @@ public static ChatCompletion ToChatCompletion(this IEnumerable<StreamingChatComp
164107
return completion;
165108
}
166109

167-
/// <summary>Coalesces sequential <see cref="TextContent"/> updates.</summary>
168-
internal static void CoalesceTextContent(this List<StreamingChatCompletionUpdate> updates)
110+
/// <summary>Coalesces sequential <see cref="TextContent"/> content elements.</summary>
111+
private static void CoalesceTextContent(List<AIContent> contents)
169112
{
170-
StringBuilder coalescedText = new();
113+
StringBuilder? coalescedText = null;
171114

172115
// Iterate through all of the items in the list looking for contiguous items that can be coalesced.
173-
for (int startInclusive = 0; startInclusive < updates.Count; startInclusive++)
116+
int start = 0;
117+
while (start < contents.Count - 1)
174118
{
175-
// If an item isn't generally coalescable, skip it.
176-
StreamingChatCompletionUpdate update = updates[startInclusive];
177-
if (update.ChoiceIndex != 0 ||
178-
update.Contents.Count != 1 ||
179-
update.Contents[0] is not TextContent textContent)
119+
// We need at least two TextContents in a row to be able to coalesce.
120+
if (contents[start] is not TextContent firstText)
180121
{
122+
start++;
181123
continue;
182124
}
183125

184-
// We found a coalescable item. Look for more contiguous items that are also coalescable with it.
185-
int endExclusive = startInclusive + 1;
186-
for (; endExclusive < updates.Count; endExclusive++)
187-
{
188-
StreamingChatCompletionUpdate next = updates[endExclusive];
189-
if (next.ChoiceIndex != 0 ||
190-
next.Contents.Count != 1 ||
191-
next.Contents[0] is not TextContent ||
192-
193-
// changing role or author would be really strange, but check anyway
194-
(update.Role is not null && next.Role is not null && update.Role != next.Role) ||
195-
(update.AuthorName is not null && next.AuthorName is not null && update.AuthorName != next.AuthorName))
196-
{
197-
break;
198-
}
199-
}
200-
201-
// If we couldn't find anything to coalesce, there's nothing to do.
202-
if (endExclusive - startInclusive <= 1)
126+
if (contents[start + 1] is not TextContent secondText)
203127
{
128+
start += 2;
204129
continue;
205130
}
206131

207-
// We found a coalescable run of items. Create a new node to represent the run. We create a new one
208-
// rather than reappropriating one of the existing ones so as not to mutate an item already yielded.
209-
_ = coalescedText.Clear().Append(updates[startInclusive].Text);
210-
211-
TextContent coalescedContent = new(null) // will patch the text after examining all items in the run
132+
// Append the text from those nodes and continue appending subsequent TextContents until we run out.
133+
// We null out nodes as their text is appended so that we can later remove them all in one O(N) operation.
134+
_ = (coalescedText ??= new()).Clear().Append(firstText.Text).Append(secondText.Text);
135+
contents[start + 1] = null!;
136+
int i = start + 2;
137+
for (; i < contents.Count && contents[i] is TextContent next; i++)
212138
{
213-
AdditionalProperties = textContent.AdditionalProperties?.Clone(),
214-
};
139+
_ = coalescedText.Append(next.Text);
140+
contents[i] = null!;
141+
}
215142

216-
StreamingChatCompletionUpdate coalesced = new()
143+
// Store the replacement node.
144+
contents[start] = new TextContent(coalescedText.ToString())
217145
{
218-
AdditionalProperties = update.AdditionalProperties?.Clone(),
219-
AuthorName = update.AuthorName,
220-
CompletionId = update.CompletionId,
221-
Contents = [coalescedContent],
222-
CreatedAt = update.CreatedAt,
223-
FinishReason = update.FinishReason,
224-
ModelId = update.ModelId,
225-
Role = update.Role,
226-
227-
// Explicitly don't include RawRepresentation. It's not applicable if one update ends up being used
228-
// to represent multiple, and it won't be serialized anyway.
146+
// We inherit the properties of the first text node. We don't currently propagate additional
147+
// properties from the subsequent nodes. If we ever need to, we can add that here.
148+
AdditionalProperties = firstText.AdditionalProperties?.Clone(),
229149
};
230150

231-
// Replace the starting node with the coalesced node.
232-
updates[startInclusive] = coalesced;
233-
234-
// Now iterate through all the rest of the updates in the run, updating the coalesced node with relevant properties,
235-
// and nulling out the nodes along the way. We do this rather than removing the entry in order to avoid an O(N^2) operation.
236-
// We'll remove all the null entries at the end of the loop, using RemoveAll to do so, which can remove all of
237-
// the nulls in a single O(N) pass.
238-
for (int i = startInclusive + 1; i < endExclusive; i++)
239-
{
240-
// Grab the next item.
241-
StreamingChatCompletionUpdate next = updates[i];
242-
updates[i] = null!;
243-
244-
var nextContent = (TextContent)next.Contents[0];
245-
_ = coalescedText.Append(nextContent.Text);
246-
247-
coalesced.AuthorName ??= next.AuthorName;
248-
coalesced.CompletionId ??= next.CompletionId;
249-
coalesced.CreatedAt ??= next.CreatedAt;
250-
coalesced.FinishReason ??= next.FinishReason;
251-
coalesced.ModelId ??= next.ModelId;
252-
coalesced.Role ??= next.Role;
253-
}
254-
255-
// Complete the coalescing by patching the text of the coalesced node.
256-
coalesced.Text = coalescedText.ToString();
257-
258-
// Jump to the last update in the run, so that when we loop around and bump ahead,
259-
// we're at the next update just after the run.
260-
startInclusive = endExclusive - 1;
151+
start = i;
261152
}
262153

263154
// Remove all of the null slots left over from the coalescing process.
264-
_ = updates.RemoveAll(u => u is null);
155+
_ = contents.RemoveAll(u => u is null);
265156
}
266157
}

0 commit comments

Comments
 (0)