|
| 1 | +package aws.community.examples; |
| 2 | + |
| 3 | +import java.nio.charset.StandardCharsets; |
| 4 | +import java.util.concurrent.CompletableFuture; |
| 5 | + |
| 6 | +import org.json.JSONObject; |
| 7 | + |
| 8 | +import software.amazon.awssdk.core.SdkBytes; |
| 9 | +import software.amazon.awssdk.regions.Region; |
| 10 | +import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeAsyncClient; |
| 11 | +import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelWithResponseStreamRequest; |
| 12 | +import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelWithResponseStreamResponseHandler; |
| 13 | +import software.amazon.awssdk.services.bedrockruntime.model.ResponseStream; |
| 14 | + |
| 15 | +public class InvokeBedrockStreamingAsync { |
| 16 | + public static void main(String[] args) { |
| 17 | + BedrockRuntimeAsyncClient runtime = BedrockRuntimeAsyncClient.builder() |
| 18 | + .region(Region.US_EAST_1) |
| 19 | + .build(); |
| 20 | + |
| 21 | + String prompt = "Explain large language models, like Anthropic Claude, in one short paragraph."; |
| 22 | + |
| 23 | + JSONObject jsonBody = new JSONObject() |
| 24 | + .put("prompt", "Human: " + prompt + " Assistant:") |
| 25 | + .put("temperature", 0.8) |
| 26 | + .put("max_tokens_to_sample", 2048); |
| 27 | + |
| 28 | + SdkBytes body = SdkBytes.fromUtf8String( |
| 29 | + jsonBody.toString()); |
| 30 | + |
| 31 | + InvokeModelWithResponseStreamRequest request = InvokeModelWithResponseStreamRequest.builder() |
| 32 | + .modelId("anthropic.claude-v2") |
| 33 | + .body(body) |
| 34 | + .build(); |
| 35 | + |
| 36 | + InvokeModelWithResponseStreamResponseHandler.Visitor visitor = InvokeModelWithResponseStreamResponseHandler.Visitor |
| 37 | + .builder() |
| 38 | + .onChunk((chunk) -> { |
| 39 | + JSONObject jsonObject = new JSONObject( |
| 40 | + chunk.bytes().asString(StandardCharsets.UTF_8)); |
| 41 | + |
| 42 | + System.out.print(jsonObject.getString("completion")); |
| 43 | + |
| 44 | + }) |
| 45 | + .onDefault((event) -> { |
| 46 | + System.out.println("\n\nDefault: " + event.toString()); |
| 47 | + }) |
| 48 | + .build(); |
| 49 | + |
| 50 | + InvokeModelWithResponseStreamResponseHandler responseHandler = InvokeModelWithResponseStreamResponseHandler |
| 51 | + .builder() |
| 52 | + .onComplete( |
| 53 | + () -> System.out.println("\n\nCompleted streaming response.")) |
| 54 | + .onError( |
| 55 | + (error) -> System.out.println( |
| 56 | + "\n\nError streaming response: " + error.getMessage())) |
| 57 | + .onEventStream((stream) -> { |
| 58 | + // print the response stream as it comes in |
| 59 | + stream.subscribe( |
| 60 | + (ResponseStream e) -> { |
| 61 | + e.accept(visitor); |
| 62 | + }); |
| 63 | + }) |
| 64 | + .build(); |
| 65 | + |
| 66 | + CompletableFuture<Void> futureResponse = runtime.invokeModelWithResponseStream( |
| 67 | + request, |
| 68 | + responseHandler); |
| 69 | + |
| 70 | + futureResponse.join(); |
| 71 | + } |
| 72 | +} |
0 commit comments