Issues with cancellation tokens in IChatCompletionService.

[MenaceFOP3]

Hi everyone,

so quite a simple one possibly, but that one has me banging my head against the wall.

I have this code:

using System.Collections.Concurrent;
using System.Runtime.CompilerServices;

using Interfaces;

using Microsoft.SemanticKernel;
using Microsoft.SemanticKernel.ChatCompletion;
using Microsoft.SemanticKernel.Connectors.OpenAI;

public class LLMChatCompletionsService(
    Kernel kernel) : ILLMChatCompletionsService
{
    private static readonly ConcurrentDictionary<string, ChatHistory> __chatHistories = new();
    private IChatCompletionService _chatService;

    private OpenAIPromptExecutionSettings _executionSettings;

    public OpenAIPromptExecutionSettings ExecutionSettings => _executionSettings ??= GetExecutionSettings();

    public IChatCompletionService ChatService =>
        _chatService ??= kernel.GetRequiredService<IChatCompletionService>();

    public ChatHistory this[string chatHistoryId]
    {
        get => __chatHistories.GetOrAdd(chatHistoryId, _ => []);
        set => __chatHistories[chatHistoryId] = value;
    }

    public ConcurrentDictionary<string, ChatHistory> ChatHistories => __chatHistories;

    public async Task<string> GetChatResponseAsync(string chatHistoryId, string userMessage,
        CancellationToken cancellationToken)
    {
        return await GetChatResponseAsync(GetOrCreateChatHistory(chatHistoryId), userMessage, cancellationToken);
    }

    public async Task<string> GetChatResponseAsync(ChatHistory chatHistory, string userMessage,
        CancellationToken cancellationToken)
    {
        chatHistory.AddUserMessage(userMessage);

        ChatMessageContent llmResponse =
            await ChatService.GetChatMessageContentAsync(chatHistory, ExecutionSettings, kernel, cancellationToken);

        chatHistory.Add(llmResponse);

        return chatHistory.Last().Content;
    }

    public async IAsyncEnumerable<string> GetChatStreamResponseAsync(string chatHistoryId, string userMessage,
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        ChatHistory chatHistory = GetOrCreateChatHistory(chatHistoryId);

        await foreach (string chunk in GetChatStreamResponseAsync(chatHistory, userMessage, cancellationToken))
        {
            yield return chunk;
        }
    }

    public async IAsyncEnumerable<string> GetChatStreamResponseAsync(ChatHistory chatHistory, string userMessage,
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        chatHistory.AddUserMessage(userMessage);

        await foreach (StreamingChatMessageContent streamingResponse in ChatService
                           .GetStreamingChatMessageContentsAsync(chatHistory)
                           .ConfigureAwait(false)
                           .WithCancellation(cancellationToken))
        {
            if (cancellationToken.IsCancellationRequested)
            {
                yield break;
            }

            // Yield each chunk as it arrives.
            yield return streamingResponse.Content;
        }
    }

    public ChatHistory CreateNewChatHistory(string chatHistoryId, string systemPrompt)
    {
        ChatHistory newHistory = __chatHistories[chatHistoryId] = new(systemPrompt);

        return newHistory;
    }

    public ChatHistory GetOrCreateChatHistory(string chatHistoryId)
    {
        return __chatHistories.GetOrAdd(chatHistoryId, []);
    }

    private static OpenAIPromptExecutionSettings GetExecutionSettings()
    {
        return new()
        {
            //// The maximum number of tokens to generate in the response.
            MaxTokens = 20000,

            //// The temperature of the model. Must be between 0 and 1.
            //// Higher values mean the model will take more risks.
            Temperature = 0.5f

            //// The maximum number of tokens to generate in the response.
            //TopP = 1.0f,
            //// The maximum number of tokens to generate in the response.
            //FrequencyPenalty = 0.0f,
            //// The maximum number of tokens to generate in the response.
            //PresencePenalty = 0.0f
        };
    }
}

So nothing crazy, and it works ok. That is until I issue a cancellation token. at which point I am getting the TaskCanceledException, as expected.

The actual issue is that the exception happens on the line:

ChatService
      .GetStreamingChatMessageContentsAsync(chatHistory)
      .ConfigureAwait(false)
      .WithCancellation(cancellationToken)

and for the life of me I don't understand why I cannot handle it gracefully or why the safeguard:

if (cancellationToken.IsCancellationRequested)
{
    yield break;
}

if completely ignored.

Am I doing something wrong here?

Thank you in advance for your help & insights.

[markwallace-microsoft]

@SergeyMenshykh can you take a look at this one?

[SergeyMenshykh]

Hi @MenaceFOP3, I think the reason the TaskCanceledException is thrown by the WithCancellation method and the expected code is not executed is that the enumerator created by the WithCancellation method detects that cancellation was requested well before your code does and throws the exception. As a result, because of the exception, the execution control leaves the await foreach (StreamingChatMessageContent streamingResponse in ChatService) loop and never reaches the if (cancellationToken.IsCancellationRequested) line.

If, for your scenario, you need to detect cancellation via cancellationToken.IsCancellationRequested rather than a cancellation exception, consider this alternative implementation of the GetChatStreamResponseAsync method:

public async IAsyncEnumerable<string> GetChatStreamResponseAsync(ChatHistory chatHistory, string userMessage, [EnumeratorCancellation] CancellationToken cancellationToken)  
{  
    chatHistory.AddUserMessage(userMessage);  
    await using var enumerator = ChatService.GetStreamingChatMessageContentsAsync(chatHistory).GetAsyncEnumerator();  
    while (await enumerator.MoveNextAsync())  
    {  
        if (cancellationToken.IsCancellationRequested)  
        {  
            yield break;  
        }  
        yield return enumerator.Current.Content;  
    }  
}  

[MenaceFOP3]

Thank you I will try it.

The actual usage comes from a SignalR hub (connected to a react app in the front-end) so I cannot really pass the cancellation token from the client, ie can only call streamSubscription.current.dispose();

Again thanks for this :)

[MenaceFOP3]

Worked a treat, thank you @SergeyMenshykh :)

[SergeyMenshykh]

You are welcome @MenaceFOP3.