Custom header policy with Azure.AI.OpenAI 2.0.0-beta.1

Program.cs

using Azure.AI.OpenAI;
using OpenAI.Chat;
using System.ClientModel.Primitives;

internal class CustomHeaderPolicy : PipelinePolicy
{
    public string NextId { get; set; }

    public override void Process(PipelineMessage message, IReadOnlyList<PipelinePolicy> pipeline, int currentIndex)
    {
        AddHeaderToRequest(message);
        ProcessNext(message, pipeline, currentIndex);
    }

    public override ValueTask ProcessAsync(PipelineMessage message, IReadOnlyList<PipelinePolicy> pipeline, int currentIndex)
    {
        AddHeaderToRequest(message);
        return ProcessNextAsync(message, pipeline, currentIndex);
    }

    private void AddHeaderToRequest(PipelineMessage message)
    {
        if (message?.Request?.Headers?.TryGetValue("x-my-custom-header-name", out _) == false)
        {
            message.Request.Headers.Set("x-my-custom-header-name", NextId);
        }
    }
}

public static class Program
{
    public static async Task Main(string[] args)
    {
        CustomHeaderPolicy customHeaderPolicy = new();
        AzureOpenAIClientOptions clientOptions = new();
        clientOptions.AddPolicy(customHeaderPolicy, PipelinePosition.BeforeTransport);

        AzureOpenAIClient azureClient = new();
        ChatClient client = azureClient.GetChatClient("my-deployment");

        customHeaderPolicy.NextId = "firstIdForTheFirstRequest";
        ChatCompletion completion = await client.CompleteChatAsync(
            [
                new UserChatMessage("Hello, assistant!"),
            ]);

        customHeaderPolicy.NextId = "nextIdForAnotherRequest";
        await client.CompleteChatAsync(
            [
                new UserChatMessage("Hi again, assistant."),
            ]);
    }
}