How do I add custom IChatClient middleware to chat client registered with a Kernel?

[jooooel]

I'm trying to figure out how to add middleware to a chat client added to the Kernel. It seems like I need the actual instance of the chat client in order to add the middleware?

This is the documentation for the chat client middleware https://learn.microsoft.com/en-us/dotnet/ai/microsoft-extensions-ai#custom-ichatclient-middleware

This is my setup

services.AddKeyedTransient<Kernel>("Kernel", (sp, _) =>
{
    ...
    var bedrockRuntime = new AmazonBedrockRuntimeClient(awsCredentials, clientConfig);

    var kernelBuilder = Kernel.CreateBuilder();
    kernelBuilder.AddBedrockChatClient(
        modelId: awsBedrockOptions.GetModel(LlmOptions.ModelPurpose.General),
        bedrockRuntime: bedrockRuntime
    );

    return kernelBuilder.Build();
});

This is the middleware:

public sealed class LlmLoggerChatClient(
    IChatClient innerClient,
    LlmLogger llmLogger)
    : DelegatingChatClient(innerClient)
{
    public override async Task<ChatResponse> GetResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null,
        CancellationToken cancellationToken = default)
    {

        var response = await base.GetResponseAsync(messages, options, cancellationToken);
        await llmLogger.LogAsync(...)
        return response;
    }

    public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        await foreach (var update in base.GetStreamingResponseAsync(messages, options, cancellationToken))
        {
            yield return update;
        }
    }
}

And this is how I finally use the kernel

ChatCompletionAgent agent = new()
{
    Name = "Search Agent",
    Kernel = kernel, // Injected via DI
    ...
};

var response = await agent.InvokeAsync(...

At what point should I/can I extend the Bedrock chat client?

[runceel]

Based on the current implementation of AddBedrockChatClient, it seems there is no direct way to inject custom middleware at this point.
As an alternative, how about registering your IChatClient into the DI container using the AddChatClient method provided by Microsoft.Extensions.AI? Here's an example:

builder.Services.AddKeyedTransient<Kernel>("Kernel", (sp, _) =>
{
    var kernelBuilder = Kernel.CreateBuilder();
    kernelBuilder.Services.AddChatClient(sp =>
    {
        var bedrockRuntime = new AmazonBedrockRuntimeClient(credentials, clientConfig);
        return bedrockRuntime.AsIChatClient();
    })
    .UseLogging()
    .UseOpenTelemetry()
    .UseKernelFunctionInvocation()
    .UseYourMiddleware();

    return kernelBuilder.Build();
});

The key point here is to use UseKernelFunctionInvocation to ensure that the middleware responsible for handling automatic plugin invocation in Semantic Kernel is properly configured.

[jooooel]

Yes, that would be an option, but I would be afraid that I'd miss something that AddBedrockChatClient does under the hood. My current approach is this:

// Add the BedrockChatClient as a "base" chat client that we can extend later or (by setting a serviceId)
kernelBuilder.AddBedrockChatClient(
    modelId: "...",
    bedrockRuntime: bedrockRuntime,
    serviceId: "base-client"
);

// Add my custom chat client that wraps the base one with logging functionality, don't use a serviceId here
// so that it will be used as the default chat client in the kernel
kernelBuilder.Services
    .AddChatClient(provider =>
    {
        // This is the base chat client that we registered above with the kernel builder
        var chatClient = provider.GetRequiredKeyedService<IChatClient>("base-client");

        // Create a new chat client that wraps the base one with logging functionality
        return new LlmLoggerChatClient(chatClient);
    });

[runceel]

@jooooel
Your current approach looks solid and unlikely to cause any issues, so I think it's a good direction. In addition to that, it might be a good idea to open an issue requesting support for injecting custom middleware into chat clients. Then we can wait and see if it gets added officially.