Bedrock ConverseStreamResponse provides no easy way to read the streamed response without blocking the calling .NET thread

[tlecomte]

Describe the bug

The common way to read from Bedrock ConverseStreamResponse is apparently to call response.Stread.AsEnumerable(). This achieves the goal of streaming the response completions as they are received, but at the cost of keeping the calling .NET thread busy. So it can lead to thread exhaustion.

Ideally we need a new variant that does not block the thread, so maybe an async enumerator exposed as response.Stread.AsAsyncEnumerable().

Regression Issue

• Select this option if this issue appears to be a regression.

Expected Behavior

I expect to have an easy method that gives an IAsyncEnumerable from ConverseStreamResponse, and which is implemented without blocking the calling thread.

Current Behavior

The current API surface in ConverseStreamResponse encourages the use of response.Stread.AsEnumerable() which blocks the calling thread.

Reproduction Steps

N/A

Possible Solution

This problem was mentioned in #3360 (comment) for a further improvement, with a suggested implementation. I believe the suggested implementation is imperfect, but it can be a basis for a solution.

Additional Information/Context

No response

AWS .NET SDK and/or Package version used

AWSSDK.BedrockRuntime 3.7.404.8

Targeted .NET Platform

.NET 8

Operating System and version

Docker

[tlecomte]

This is in principle how it could be implemented, using an extension, and with System.Threading.Channels:

public static class ConverseStreamOutputExtensions
{
        public async IAsyncEnumerable<IEventStreamEvent> AsAsyncEnumerable(ConverseStreamOutput output, [EnumeratorCancellation] CancellationToken cancellationToken)
        {
            var channel = Channel.CreateUnbounded<IEventStreamEvent>();

            void OnEventReceived(object? sender, EventStreamEventReceivedArgs<IEventStreamEvent> args)
            {
                var _ = channel.Writer.TryWrite(args.EventStreamEvent);

                // the metadata event is the last one of the event stream
                if (args.EventStreamEvent is ConverseStreamMetadataEvent)
                {
                    channel.Writer.Complete();
                }
            }

            void OnExceptionReceived(object? sender, EventStreamExceptionReceivedArgs<BedrockRuntimeEventStreamException> args)
            {
                var _ = channel.Writer.TryComplete(args.EventStreamException);
            }

            Task processingTask;

            try
            {
                output.EventReceived += OnEventReceived;
                output.ExceptionReceived += OnExceptionReceived;

                processingTask = output..StartProcessingAsync();

                await foreach (var e in channel.Reader.ReadAllAsync(cancellationToken))
                {
                    yield return e;
                }
            }
            finally
            {
                output.EventReceived -= OnEventReceived;
                output.ExceptionReceived -= OnExceptionReceived;
            }

            await processingTask.ConfigureAwait(false);
        }
}

Not completely sure how such code could be added to the SDK since all the files say that there are auto-generated from the service model json...

[ashishdhingra]

Needs review with the team.

[tlecomte]

I've submitted a pull request to add support for IAsyncEnumerable in EnumerableEventStream, which makes it available for Bedrock: #3543

It's much easier than with the extension above, fortunately.

[shethaadit]

Hi @ashishdhingra, which SDK version will get this change?

[ashishdhingra]

Hi @ashishdhingra, which SDK version will get this change?

@shethaadit The issue is marked with v4 label. It should be available in next major version v4 of AWS .NET SDK, which is currently in preview.

CC @normj