Skip to content

Fix: Address several Bedrock converse streaming issues #4375

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix: Address several Bedrock converse streaming issues #4375

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9cf44ec
Fix converse streaming issues:
emopti-jrufer Sep 11, 2025
b045a1e
Added missing requestMetadata mapping for ConverseStreamRequest.
emopti-jrufer Sep 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 10 additions & 48 deletions ...converse/src/main/java/org/springframework/ai/bedrock/converse/BedrockProxyChatModel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,6 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;
import reactor.core.publisher.Sinks.EmitFailureHandler;
import reactor.core.scheduler.Schedulers;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.core.SdkBytes;
Expand All @@ -51,9 +49,7 @@
import software.amazon.awssdk.services.bedrockruntime.model.ConverseMetrics;
import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamOutput;
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamRequest;
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamResponseHandler;
import software.amazon.awssdk.services.bedrockruntime.model.DocumentBlock;
import software.amazon.awssdk.services.bedrockruntime.model.DocumentSource;
import software.amazon.awssdk.services.bedrockruntime.model.ImageBlock;
Expand All @@ -76,6 +72,7 @@

import org.springframework.ai.bedrock.converse.api.BedrockMediaFormat;
import org.springframework.ai.bedrock.converse.api.ConverseApiUtils;
import org.springframework.ai.bedrock.converse.api.ConverseChatResponseStream;
import org.springframework.ai.bedrock.converse.api.URLValidator;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.MessageType;
Expand All @@ -84,6 +81,7 @@
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.metadata.DefaultUsage;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
Expand Down Expand Up @@ -683,11 +681,17 @@ private Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse perviousCh
.system(converseRequest.system())
.additionalModelRequestFields(converseRequest.additionalModelRequestFields())
.toolConfig(converseRequest.toolConfig())
.requestMetadata(converseRequest.requestMetadata())
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is recently added via 45baa39.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. Dropped commit.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually this is missing. It was added for the ConverseRequest but the mapping to ConverseStreamRequest was missed in that commit.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@emopti-jrufer oh yeah, thanks for catching that!

.build();

Flux<ConverseStreamOutput> response = converseStream(converseStreamRequest);
Usage accumulatedUsage = null;
if (perviousChatResponse != null && perviousChatResponse.getMetadata() != null) {
accumulatedUsage = perviousChatResponse.getMetadata().getUsage();
}

Flux<ChatResponse> chatResponses = ConverseApiUtils.toChatResponse(response, perviousChatResponse);
Flux<ChatResponse> chatResponses = new ConverseChatResponseStream(this.bedrockRuntimeAsyncClient,
converseStreamRequest, accumulatedUsage)
.stream();

Flux<ChatResponse> chatResponseFlux = chatResponses.switchMap(chatResponse -> {

Expand Down Expand Up @@ -734,48 +738,6 @@ private Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse perviousCh
});
}

public static final EmitFailureHandler DEFAULT_EMIT_FAILURE_HANDLER = EmitFailureHandler
.busyLooping(Duration.ofSeconds(10));

/**
* Invoke the model and return the response stream.
*
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html
* https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
* https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/services/bedrockruntime/BedrockRuntimeAsyncClient.html#converseStream
* @param converseStreamRequest Model invocation request.
* @return The model invocation response stream.
*/
public Flux<ConverseStreamOutput> converseStream(ConverseStreamRequest converseStreamRequest) {
Assert.notNull(converseStreamRequest, "'converseStreamRequest' must not be null");

Sinks.Many<ConverseStreamOutput> eventSink = Sinks.many().multicast().onBackpressureBuffer();

ConverseStreamResponseHandler.Visitor visitor = ConverseStreamResponseHandler.Visitor.builder()
.onDefault(output -> {
logger.debug("Received converse stream output:{}", output);
eventSink.emitNext(output, DEFAULT_EMIT_FAILURE_HANDLER);
})
.build();

ConverseStreamResponseHandler responseHandler = ConverseStreamResponseHandler.builder()
.onEventStream(stream -> stream.subscribe(e -> e.accept(visitor)))
.onComplete(() -> {
eventSink.emitComplete(DEFAULT_EMIT_FAILURE_HANDLER);
logger.info("Completed streaming response.");
})
.onError(error -> {
logger.error("Error handling Bedrock converse stream response", error);
eventSink.emitError(error, DEFAULT_EMIT_FAILURE_HANDLER);
})
.build();

this.bedrockRuntimeAsyncClient.converseStream(converseStreamRequest, responseHandler);

return eventSink.asFlux();

}

/**
* Use the provided convention for reporting observation data
* @param observationConvention The provided convention
Expand Down
Loading