Add Amazon Bedrock Unified Chat Completions support

Author: Evgenii-Kazannik

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/request/completion/AmazonBedrockUnifiedChatCompletionRequest.java

@@ -8,7 +8,11 @@
 package org.elasticsearch.xpack.inference.services.amazonbedrock.request.completion;
 
 import software.amazon.awssdk.core.document.Document;
+import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockDeltaEvent;
+import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockStartEvent;
 import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamRequest;
+import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamResponseHandler;
+import software.amazon.awssdk.services.bedrockruntime.model.MessageStopEvent;
 import software.amazon.awssdk.services.bedrockruntime.model.SpecificToolChoice;
 import software.amazon.awssdk.services.bedrockruntime.model.Tool;
 import software.amazon.awssdk.services.bedrockruntime.model.ToolChoice;
@@ -30,6 +34,8 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Flow;
 
 import static org.elasticsearch.xpack.inference.services.amazonbedrock.request.completion.AmazonBedrockConverseUtils.getUnifiedConverseMessageList;
@@ -54,7 +60,7 @@ public AmazonBedrockUnifiedChatCompletionRequest(
 
     public Flow.Publisher<StreamingUnifiedChatCompletionResults.Results> executeStreamChatCompletionRequest(
         AmazonBedrockBaseClient awsBedrockClient
-    ) {
+    ) throws ExecutionException, InterruptedException {
         var converseStreamRequest = ConverseStreamRequest.builder()
             .messages(getUnifiedConverseMessageList(requestEntity.messages()))
             .modelId(amazonBedrockModel.model());
@@ -86,6 +92,59 @@ public Flow.Publisher<StreamingUnifiedChatCompletionResults.Results> executeStre
             });
         }
 
+        inferenceConfig(requestEntity).ifPresent(converseStreamRequest::inferenceConfig);
+        var response = awsBedrockClient.converseUnifiedStream(converseStreamRequest.build());
+
+        var toolRequested = new CompletableFuture<Boolean>();
+        final String[] toolUseIdHolder = new String[1];
+        final StringBuilder toolJsonArgs = new StringBuilder();
+        final StringBuilder assistantText = new StringBuilder();
+
+        var handler = ConverseStreamResponseHandler.builder()
+            .onEventStream(es -> es.subscribe(event -> {
+                switch (event.sdkEventType()) {
+                    case MESSAGE_START:
+                        break;
+                    case CONTENT_BLOCK_START:
+                        var start = ((ContentBlockStartEvent) event).start();
+                        if (start.toolUse() != null) {
+                            toolUseIdHolder[0] = start.toolUse().toolUseId();
+                        }
+                        break;
+                    case CONTENT_BLOCK_DELTA:
+                        var delta = ((ContentBlockDeltaEvent) event).delta();
+                        if (delta.toolUse() != null && delta.toolUse().input() != null) {
+                            toolJsonArgs.append(delta.toolUse().input());
+                        }
+                        if (delta.text() != null) {
+                            assistantText.append(delta.text());
+                        }
+                        break;
+                    case MESSAGE_STOP:
+                        var stop = ((MessageStopEvent) event).stopReason();
+                        if ("tool_use".equalsIgnoreCase(stop.name())) {
+                            toolRequested.complete(true);
+                        } else {
+                            toolRequested.complete(false);
+                        }
+                        break;
+                    default:
+                }
+            }))
+            .onResponse(r -> toolRequested.complete(true))
+            .onError(toolRequested::completeExceptionally);
+
+        handler.subscriber(converseStreamOutput ->
+            getUnifiedConverseMessageList(requestEntity.messages()).forEach(toolJsonArgs::append));
+
+        if (Boolean.TRUE.equals(toolRequested.get())) {
+            toolJsonArgs.toString().contains("args");
+            Map<String, Object> result = Map.of("tool_use", toolUseIdHolder[0]);
+//            var toolResultBlock = ContentBlock
+//                .fromToolResult(ToolResultContentBlock.builder()
+//                    .document(DocumentBlock.builder().context(result).build()));
+
+        }
         inferenceConfig(requestEntity).ifPresent(converseStreamRequest::inferenceConfig);
         return awsBedrockClient.converseUnifiedStream(converseStreamRequest.build());
     }