Implement Microsoft.Extensions.AI's IChatClient / IEmbeddingGenerator for IAmazonBedrockRuntime

[stephentoub]

Description

This enables AmazonBedrockRuntimeClient to be used as a Microsoft.Extensions.AI.IChatClient, such that it can implicitly be used by any consumer that operates on an IChatClient, and with any middleware written in terms of IChatClient, such as those components in the Microsoft.Extensions.AI package that provide support for automatic function invocation, OpenTelemetry, logging, distributed caching, and more.

https://devblogs.microsoft.com/dotnet/introducing-microsoft-extensions-ai-preview/

Motivation and Context

IChatClient was recently introduced as a core .NET abstraction for representing AI chat completion services. AmazonBedrockRuntimeClient is a natural fit for implementing this interface, such that application and library code can be written in terms of IChatClient and "just work" with AWS.

Testing

Ran manually ad-hoc tests against the implementation, e.g.

Recording.2024-11-05.213108.mp4

Screenshots (if appropriate)

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)

Checklist

• My code follows the code style of this project
• My change requires a change to the documentation
• I have updated the documentation accordingly
• I have read the README document
• I have added tests to cover my changes
• All new and existing tests passed

License

• I confirm that this pull request can be released under the Apache 2 license

[normj]

Thanks @stephentoub for the PR. Looks like a really interesting integration. The AWSSDK. packages are meant to be the low level service clients with limited dependencies. Can you create a new project under the extensions folder https://github.com/aws/aws-sdk-net/tree/main/extensions and put the code changes you put in AWSSDK.Bedrock into that project. The new project would have a project dependency on AWSSDK.Bedrock. I'm thinking the project could be named AWSSDK.Extensions.IChatClient.Bedrock unless you have a better product name.

[stephentoub]

Ok. As a core .NET library I thought the dependency would be acceptable, but I can separate it out into a separate project. I would like to leave the IAsyncEnumerable implementation in place, though; that's important for scale / correctness, and without it, examples like those at https://docs.aws.amazon.com/sdk-for-net/v3/developer-guide/csharp_bedrock-runtime_code_examples.html are problematic, because while they're trying to be async, they're actually doing the bulk of the work doing synchronous I/O.

[normj]

I would like to leave the IAsyncEnumerable implementation in place, though; that's important for scale / correctness, and without it, examples like those at https://docs.aws.amazon.com/sdk-for-net/v3/developer-guide/csharp_bedrock-runtime_code_examples.html are problematic, because while they're trying to be async, they're actually doing the bulk of the work doing synchronous I/O.

Agree the code adding IAsyncEnumerable support to the event streams in core should be left there. It improves our core functionality and doesn't add any new dependencies on the SDK.

[stephentoub]

(Haven't forgotten about this, just been busy with the .NET 9 release this week. Will update the PR soon.)

[normj]

@stephentoub We figured as much. Great talk today!

[stephentoub]

@normj, updated. I created a new AWSSDK.Extensions.Bedrock.MEAI project, which now includes both the IChatClient implementation and an implementation of IEmbeddingGenerator. As they're now not in the same assembly as AmazonBedrockRuntimeClient, I instead exposed them via AsChatClient and AsEmbeddingGenerator extension methods off of IAmazonBedrockRuntime, which is the same general shape we've used for other cases with other providers.

[stephentoub]

(I should note, M.E.AI.Abstractions is still preview, and we do expect minor breaking changes here and there. I'll aim to submit PRs to keep this assembly up-to-date with the latest.)

[stephentoub]

I'd included a video showing using IChatClient. Here's an example using IEmbeddingGenerator:

using Amazon;
using Amazon.BedrockRuntime;
using Microsoft.Extensions.AI;
using System.Numerics.Tensors;

var generator = new AmazonBedrockRuntimeClient(
    Environment.GetEnvironmentVariable("AI:Bedrock:AccessKey"),
    Environment.GetEnvironmentVariable("AI:Bedrock:SecretAccessKey"),
    RegionEndpoint.USEast1).AsEmbeddingGenerator("amazon.titan-embed-text-v1");

var inputEmbedding = await generator.GenerateEmbeddingVectorAsync("What is an amphibian?");

var embeddings = await generator.GenerateAndZipAsync(
[
    "Cos'è un anfibio?",
    "A frog is an amphibian.",
    "Frogs, toads, and salamanders are all examples.",
    "Amphibians are four-limbed and ectothermic vertebrates of the class Amphibia.",
    "They are four-limbed and ectothermic vertebrates.",
    "A frog is green.",
    "A tree is green.",
    "It's not easy bein' green.",
    "A dog is a mammal.",
    "A dog is a man's best friend.",
    "You ain't never had a friend like me.",
    "Rachel, Monica, Phoebe, Joey, Chandler, Ross"
]);

foreach (var e in embeddings
    .Select(e => (e.Value, Similarity: TensorPrimitives.CosineSimilarity(inputEmbedding.Span, e.Embedding.Vector.Span)))
    .OrderByDescending(e => e.Similarity))
{
    Console.WriteLine($"{e.Similarity:0.000}: {e.Value}");
}

resulting in:

0.822: A frog is an amphibian.
0.734: Amphibians are four-limbed and ectothermic vertebrates of the class Amphibia.
0.623: Cos'è un anfibio?
0.609: Frogs, toads, and salamanders are all examples.
0.525: A frog is green.
0.519: They are four-limbed and ectothermic vertebrates.
0.287: A dog is a mammal.
0.197: A tree is green.
0.107: A dog is a man's best friend.
0.103: You ain't never had a friend like me.
0.085: It's not easy bein' green.
-0.031: Rachel, Monica, Phoebe, Joey, Chandler, Ross

[normj]

@stephentoub Thanks for the update. Going through it but bare with me while I get a better understanding of Microsoft.Extensions.AI.

[RogerBarreto]

Small comment only, overall LGTM

[peterrsongg]

My comments are non-blocking, but would like to see what norm thinks on the lang version. That and the nullable reference types are the two main differences I see between this extensions package and the others. Don't see an issue with either

[stephentoub]

get a better understanding of Microsoft.Extensions.AI.

I'd recommend this talk from .NET Conf 2024 last week:
https://www.youtube.com/watch?v=qcp6ufe_XYo

[peterrsongg]

I had to make a few changes to the example code provided in the description for it to work, but other than that the change looks good. Thanks for the contribution!

[stephentoub]

I had to make a few changes to the example code provided in the description for it to work

In the video in the description? Yeah, that was from the original design, before I was asked to move it to a separate library, at which point it needed a method to create the implementation. I assume the only change was needing to call AsChatClient? Or something else as well?

Thanks for the contribution!

You're welcome! Looking forward to being able to reference the nuget package when it's available.

[peterrsongg]

I had to make a few changes to the example code provided in the description for it to work

In the video in the description? Yeah, that was from the original design, before I was asked to move it to a separate library, at which point it needed a method to create the implementation. I assume the only change was needing to call AsChatClient? Or something else as well?

Yup I had to add AsChatClient and another small change I made is changing sb.Append(update) to sb.Append(update.Text);

[stephentoub]

and another small change I made is changing sb.Append(update) to sb.Append(update.Text);

I wouldn't have expected that to actually be necessary... was it?

[peterrsongg]

and another small change I made is changing sb.Append(update) to sb.Append(update.Text);

I wouldn't have expected that to actually be necessary... was it?

ah I could've sworn it was complaining about not being able to convert StreamingChatCompletionUpdate to a string but I think I had something else configured incorrectly. Changing that part actually wasn't necessary👍

[stephentoub]

Cool, thanks

[normj]

I have merged the change into the feature-v4/meai-extension feature branch so we can take care of all of the internal packaging changing and checks before merging into the main v4-development branch. I'll update here the progress on the internal work we have to do and when the package goes out. Thanks so much for the excellent PR.

[stephentoub]

Thanks, @normj! Let me know if there's anything further I can do to help. As there are subsequent releases, I can help follow-up to adapt as well. Also, I'm assuming with the Nova announcements, Bedrock will be updated as well, e.g. with types to represent video, and we'll want to react to those here as well.

[normj]

@stephentoub that would be great to keep us updated on the progress of Microsoft.Extensions.AI and help with the extension would be appreciated.

[normj]

@stephentoub The package was released as part of preview 5 of V4 SDK. I would like to get a blog post written at some point about the library.

https://www.nuget.org/packages/AWSSDK.Extensions.Bedrock.MEAI/4.0.0-preview.5

[stephentoub]

Awesome! I'll check it out.