如何为AzureOpenAI ChatClient流式API正确传入取消令牌？

我用C#编写了Azure FunctionApp，通过Azure.AI.OpenAI库调用Azure OpenAI聊天补全客户端。调用非流式API（CompleteChatAsync）时，传入的取消令牌能正常生效；但调用流式API（CompleteChatStreamingAsync）返回部分响应时，取消令牌不生效，怀疑每次获取部分响应都会重置取消令牌。完整调用方法如下：

public async Task<string> CompletePromptAsync(double timeout, string promptData)
{
    // Generate prompt messages and options 
    List<ChatMessage> promptMessages = new();
    promptMessages.Add(new UserChatMessage("Please convert the input text data into JSON. Make sure to return JSON array with the elements extracted from the input data."));
    promptMessages.Add(new UserChatMessage($"Here is the input data: {promptData}"));
    
    ChatCompletionOptions promptOptions = new ChatCompletionOptions
    {
        Temperature = 0.1,
        FrequencyPenalty = 0,
        PresencePenalty = 0
    };

    Azure.AI.OpenAI.AzureOpenAIClient _openAIClient = new AzureOpenAIClient(new Uri(openAIEndpoint), credential);
    ChatClient chatClient = _openAIClient.GetChatClient(deployment);

    using var ct = new CancellationTokenSource(TimeSpan.FromSeconds(timeout));

    // To Efficiently handle the larger response, build response using asynchronous streaming
    var responseBuilder = new System.Text.StringBuilder();
    await foreach (StreamingChatCompletionUpdate completionUpdate in chatClient.CompleteChatStreamingAsync(promptMessages, promptOptions, ct.Token))
    {
        foreach (ChatMessageContentPart content in completionUpdate.ContentUpdate)
        {
            responseBuilder.Append(content.Text);
        }
    }

    return responseBuilder.ToString();
}

请问如何让流式API正确响应取消令牌，以及如何通过传入令牌实现整体调用超时后停止请求？

1. 显式绑定取消令牌到异步流 #

Azure OpenAI的流式API迭代过程中不会自动持续检查取消令牌，需要通过WithCancellation扩展方法将令牌与异步流的迭代流程绑定，确保取消信号触发时能立即终止迭代。

2. 关键节点主动检查取消状态 #

在处理每个流式响应块的前后，主动检查取消令牌是否已触发，若已取消则直接抛出OperationCanceledException终止流程，避免继续处理后续响应块。

修改后的完整代码 #

public async Task<string> CompletePromptAsync(double timeout, string promptData)
{
    List<ChatMessage> promptMessages = new();
    promptMessages.Add(new UserChatMessage("Please convert the input text data into JSON. Make sure to return JSON array with the elements extracted from the input data."));
    promptMessages.Add(new UserChatMessage($"Here is the input data: {promptData}"));
    
    ChatCompletionOptions promptOptions = new ChatCompletionOptions
    {
        Temperature = 0.1,
        FrequencyPenalty = 0,
        PresencePenalty = 0
    };

    Azure.AI.OpenAI.AzureOpenAIClient _openAIClient = new AzureOpenAIClient(new Uri(openAIEndpoint), credential);
    ChatClient chatClient = _openAIClient.GetChatClient(deployment);

    using var ct = new CancellationTokenSource(TimeSpan.FromSeconds(timeout));
    ct.Token.ThrowIfCancellationRequested(); // 提前检查是否已超时

    var responseBuilder = new System.Text.StringBuilder();
    // 为异步流绑定取消令牌
    await foreach (StreamingChatCompletionUpdate completionUpdate in chatClient.CompleteChatStreamingAsync(promptMessages, promptOptions, ct.Token).WithCancellation(ct.Token))
    {
        ct.Token.ThrowIfCancellationRequested(); // 处理响应块前检查取消状态
        foreach (ChatMessageContentPart content in completionUpdate.ContentUpdate)
        {
            ct.Token.ThrowIfCancellationRequested(); // 处理内容片段时检查
            responseBuilder.Append(content.Text);
        }
    }

    return responseBuilder.ToString();
}

关键说明 #

• WithCancellation(ct.Token)：强制将取消令牌与异步流的迭代过程绑定，确保等待下一个响应块时能响应取消信号。
• ct.Token.ThrowIfCancellationRequested()：在初始化、处理响应块、处理内容片段等关键节点主动检查，一旦令牌触发，立即终止流程，不会继续处理后续响应。

内容的提问来源于stack exchange，提问作者HereToLearn