Add Amazon Bedrock Unified Chat Completions support

Author: Evgenii-Kazannik

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/client/AmazonBedrockClient.java

@@ -7,6 +7,8 @@
 
 package org.elasticsearch.xpack.inference.services.amazonbedrock.client;
 
+import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
+
 import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
 import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
 import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamRequest;
@@ -26,6 +28,9 @@ public interface AmazonBedrockClient {
     Flow.Publisher<? extends InferenceServiceResults.Result> converseStream(ConverseStreamRequest converseStreamRequest)
         throws ElasticsearchException;
 
+    Flow.Publisher<StreamingUnifiedChatCompletionResults.Results> converseUnifiedStream(ConverseStreamRequest request)
+        throws ElasticsearchException;
+
     void invokeModel(InvokeModelRequest invokeModelRequest, ActionListener<InvokeModelResponse> responseListener)
         throws ElasticsearchException;
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/client/AmazonBedrockInferenceClient.java

@@ -7,6 +7,8 @@
 
 package org.elasticsearch.xpack.inference.services.amazonbedrock.client;
 
+import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
+
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
@@ -101,6 +103,20 @@ public Flow.Publisher<? extends InferenceServiceResults.Result> converseStream(C
         return awsResponseProcessor;
     }
 
+    @Override
+    public Flow.Publisher<StreamingUnifiedChatCompletionResults.Results> converseUnifiedStream(ConverseStreamRequest request)
+        throws ElasticsearchException {
+        var awsResponseProcessor = new AmazonBedrockUnifiedStreamingChatProcessor(threadPool);
+        internalClient.converseStream(
+            request,
+            ConverseStreamResponseHandler.builder().subscriber(() -> FlowAdapters.toSubscriber(awsResponseProcessor)).build()
+        ).exceptionally(e -> {
+            awsResponseProcessor.onError(e);
+            return null; // Void
+        });
+        return awsResponseProcessor;
+    }
+
     private void onFailure(ActionListener<?> listener, Throwable t, String method) {
         ExceptionsHelper.maybeDieOnAnotherThread(t);
         var unwrappedException = t;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/client/AmazonBedrockUnifiedChatCompletionExecutor.java

@@ -10,7 +10,7 @@
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.inference.InferenceServiceResults;
-import org.elasticsearch.xpack.core.inference.results.StreamingChatCompletionResults;
+import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
 import org.elasticsearch.xpack.inference.services.amazonbedrock.request.completion.AmazonBedrockUnifiedChatCompletionRequest;
 import org.elasticsearch.xpack.inference.services.amazonbedrock.response.AmazonBedrockResponseHandler;
 
@@ -35,7 +35,7 @@ protected AmazonBedrockUnifiedChatCompletionExecutor(
     protected void executeClientRequest(AmazonBedrockBaseClient awsBedrockClient) {
         if (chatCompletionRequest.isStreaming()) {
             var publisher = chatCompletionRequest.executeStreamChatCompletionRequest(awsBedrockClient);
-            inferenceResultsListener.onResponse(new StreamingChatCompletionResults(publisher));
+            inferenceResultsListener.onResponse(new StreamingUnifiedChatCompletionResults(publisher));
         }
     }
 }