Skip to content

Commit

Permalink
Fix a few FunctionInvocationChatClient streaming issues (#5680)
Browse files Browse the repository at this point in the history
- The non-streaming path explicitly throws if the response contains multiple choices. The streaming path wasn't doing the same and was instead silently producing bad results.
- The streaming path was yielding function call content _and_ adding them to the chat history. It should only have been doing the latter.

This fixes both issues. We also had close to zero test coverage in our FunctionInvocationChatClient tests for streaming, only for non-streaming. This also fixes that.
  • Loading branch information
stephentoub authored Nov 20, 2024
1 parent 7ebb34d commit 1a4a54f
Show file tree
Hide file tree
Showing 2 changed files with 371 additions and 156 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -216,7 +216,7 @@ public override async Task<ChatCompletion> CompleteAsync(IList<ChatMessage> chat
// doesn't realize this and is wasting their budget requesting extra choices we'd never use.
if (response.Choices.Count > 1)
{
throw new InvalidOperationException($"Automatic function call invocation only accepts a single choice, but {response.Choices.Count} choices were received.");
ThrowForMultipleChoices();
}

// Extract any function call contents on the first choice. If there are none, we're done.
Expand Down Expand Up @@ -301,22 +301,47 @@ public override async IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteSt
_ = Throw.IfNull(chatMessages);

HashSet<ChatMessage>? messagesToRemove = null;
List<FunctionCallContent> functionCallContents = [];
int? choice;
try
{
for (int iteration = 0; ; iteration++)
{
List<FunctionCallContent>? functionCallContents = null;
await foreach (var chunk in base.CompleteStreamingAsync(chatMessages, options, cancellationToken).ConfigureAwait(false))
choice = null;
functionCallContents.Clear();
await foreach (var update in base.CompleteStreamingAsync(chatMessages, options, cancellationToken).ConfigureAwait(false))
{
// We're going to emit all StreamingChatMessage items upstream, even ones that represent
// function calls, because a given StreamingChatMessage can contain other content too.
yield return chunk;
// function calls, because a given StreamingChatMessage can contain other content, too.
// And if we yield the function calls, and the consumer adds all the content into a message
// that's then added into history, they'll end up with function call contents that aren't
// directly paired with function result contents, which may cause issues for some models
// when the history is later sent again.

// Find all the FCCs. We need to track these separately in order to be able to process them later.
int preFccCount = functionCallContents.Count;
functionCallContents.AddRange(update.Contents.OfType<FunctionCallContent>());

// If there were any, remove them from the update. We do this before yielding the update so
// that we're not modifying an instance already provided back to the caller.
int addedFccs = functionCallContents.Count - preFccCount;
if (addedFccs > preFccCount)
{
update.Contents = addedFccs == update.Contents.Count ?
[] : update.Contents.Where(c => c is not FunctionCallContent).ToList();
}

foreach (var item in chunk.Contents.OfType<FunctionCallContent>())
// Only one choice is allowed with automatic function calling.
if (choice is null)
{
choice = update.ChoiceIndex;
}
else if (choice != update.ChoiceIndex)
{
functionCallContents ??= [];
functionCallContents.Add(item);
ThrowForMultipleChoices();
}

yield return update;
}

// If there are no tools to call, or for any other reason we should stop, return the response.
Expand Down Expand Up @@ -373,6 +398,16 @@ public override async IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteSt
}
}

/// <summary>Throws an exception when multiple choices are received.</summary>
private static void ThrowForMultipleChoices()
{
// If there's more than one choice, we don't know which one to add to chat history, or which
// of their function calls to process. This should not happen except if the developer has
// explicitly requested multiple choices. We fail aggressively to avoid cases where a developer
// doesn't realize this and is wasting their budget requesting extra choices we'd never use.
throw new InvalidOperationException("Automatic function call invocation only accepts a single choice, but multiple choices were received.");
}

/// <summary>
/// Removes all of the messages in <paramref name="messagesToRemove"/> from <paramref name="messages"/>
/// and all of the content in <paramref name="contentToRemove"/> from the messages in <paramref name="messages"/>.
Expand Down
Loading

0 comments on commit 1a4a54f

Please sign in to comment.