Use of PipelinePolicies in Azure.AI.OpenAI.AzureOpenAIClient

In this post let's see how we can make use of PipelinePolicies in AzureOpenAIClient. Even though we specifically talk about AzureOpenAIClient in this post, the concept will be the same among many Clients in Azure SDK.

When a client in Azure SDK sends a request to Azure Service/API, the request travels through a pipeline which consists of a set of policies that get to modify the request before it's being sent and observe the response after it's received.

In order to add a Custom policy, we need to implement the abstract class PipelinePolicy.

public class CustomLoggingPolicy : PipelinePolicy
{
    // For synchronous processing
    // Not implementing as most of the time we are calling async methods in the client
    public override void Process(PipelineMessage message, IReadOnlyList<PipelinePolicy> pipeline, int currentIndex)
    {
        ProcessNext(message, pipeline, currentIndex);
    }
 
    // For asynchronous processing
    public override async ValueTask ProcessAsync(PipelineMessage message, IReadOnlyList<PipelinePolicy> pipeline, int currentIndex)
    {
        Console.WriteLine($"Sending Request to {message.Request.Method}: {message.Request.Uri}");
 
        await ProcessNextAsync(message, pipeline, currentIndex);
 
        if (message.Response?.IsError == true)
        {
            Console.WriteLine($"Request to {message.Request.Method}: {message.Request.Uri} failed.");
            Console.WriteLine(new
            {
                message.Response.ReasonPhrase,
                message.Response.Content,
            });
        }
        else
        {
            Console.WriteLine($"Request to {message.Request.Method}: {message.Request.Uri} completed.");
        }
    }
}

Here I have added CustomLoggingPolicy for demo purposes. A key thing to remember is, we need to call the  ProcessNext/ProcessNextAsync method to pass the control to the next PipelinePolicy.

We can now make use of the CustomLoggingPolicy as follows.

string deploymentName = "gpt-4o";
 
AzureOpenAIClientOptions azureOpenAIClientOptions = new();
azureOpenAIClientOptions.AddPolicy(new CustomLoggingPolicy(), PipelinePosition.BeforeTransport);
 
AzureOpenAIClient azureOpenAiClient =
    new(new Uri("<openai-endpoint>"), new DefaultAzureCredential(), azureOpenAIClientOptions);
 
ChatClient chatClient = azureOpenAiClient.GetChatClient(deploymentName);
 
List<ChatMessage> chatMessages =
[
    // messages
];
 
ChatCompletionOptions chatCompletionOptions = new()
{
    // options
};
 
ClientResult<ChatCompletion> clientResult
   = await chatClient.CompleteChatAsync(chatMessages, chatCompletionOptions);

CustomLoggingPolicy

Say you want to add a Custom Retry policy, then you can implement a class that inherits from ClientRetryPolicy and override ShouldRetryAsync method.

public class CustomClientRetryPolicy : ClientRetryPolicy
{
    protected override async ValueTask<bool> ShouldRetryAsync(PipelineMessage message, Exception? exception)
    {
        if (message.Response.Status == (int)HttpStatusCode.RequestEntityTooLarge)
        {
            return false;
        }
 
        return await base.ShouldRetryAsync(message, exception);
    }
}

Usage:

AzureOpenAIClientOptions azureOpenAIClientOptions = new();
...
azureOpenAIClientOptions.RetryPolicy = new CustomClientRetryPolicy();

Hope this helps.

Happy Coding.

Regards,

Jaliya