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