add bedrock support

Signed-off-by: Jiaping Zeng <[email protected]>

Author: jiapingzeng

common/src/main/java/org/opensearch/ml/common/transport/execute/MLExecuteTaskResponse.java

@@ -26,6 +26,8 @@
 import org.opensearch.ml.common.output.model.ModelTensorOutput;
 import org.opensearch.ml.common.output.model.ModelTensors;
 
+import com.google.gson.Gson;
+
 import lombok.Builder;
 import lombok.Getter;
 import lombok.NonNull;
@@ -103,7 +105,7 @@ private AGUIOutput extractAGUIOutput(ModelTensorOutput modelOutput) {
         try {
             ModelTensors modelTensors = modelOutput.getMlModelOutputs().get(0);
             String eventsJson = modelTensors.getMlModelTensors().get(0).getResult();
-            com.google.gson.Gson gson = new com.google.gson.Gson();
+            Gson gson = new Gson();
             List<Object> events = gson.fromJson(eventsJson, List.class);
             return AGUIOutput.builder().events(events).build();
         } catch (Exception e) {

common/src/test/java/org/opensearch/ml/common/agui/AGUIInputConverterTest.java

@@ -14,6 +14,8 @@
 import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
 import org.opensearch.ml.common.input.execute.agent.AgentMLInput;
 
+import com.google.gson.JsonObject;
+
 public class AGUIInputConverterTest {
 
     @Test
@@ -254,7 +256,7 @@ public void testReconstructAGUIInput() {
         parameters.put("agui_messages", "[]");
         parameters.put("agui_tools", "[]");
 
-        com.google.gson.JsonObject result = AGUIInputConverter.reconstructAGUIInput(parameters);
+        JsonObject result = AGUIInputConverter.reconstructAGUIInput(parameters);
 
         assertNotNull("Reconstructed result should not be null", result);
         assertTrue("Should contain threadId", result.has("threadId"));

common/src/test/java/org/opensearch/ml/common/utils/StringUtilsTest.java

@@ -17,6 +17,7 @@
 
 import java.io.IOException;
 import java.lang.reflect.Constructor;
+import java.lang.reflect.Type;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -1201,9 +1202,9 @@ public void testSerializeFloatNaNAndInfinity_BecomesNull_InPojo() {
     @Test
     public void testDeserializeScientificNotation_ToFloatAndPrimitive() {
         String jsonObj = "{\"fObj\":1.23e-5}";
-        java.lang.reflect.Type mapType = new com.google.gson.reflect.TypeToken<java.util.Map<String, Float>>() {
+        Type mapType = new TypeToken<Map<String, Float>>() {
         }.getType();
-        java.util.Map<String, Float> m = StringUtils.PLAIN_NUMBER_GSON.fromJson(jsonObj, mapType);
+        Map<String, Float> m = StringUtils.PLAIN_NUMBER_GSON.fromJson(jsonObj, mapType);
         assertEquals(1.23e-5f, m.get("fObj"), 1e-9f);
 
         String jsonArr = "[4.56e1]";
@@ -1215,9 +1216,9 @@ public void testDeserializeScientificNotation_ToFloatAndPrimitive() {
     public void testDeserializeNullFloat_ToNull() {
         String json = "{\"fObj\":null,\"fPrim\":1.0}";
 
-        java.lang.reflect.Type mapType = new TypeToken<java.util.Map<String, JsonElement>>() {
+        Type mapType = new TypeToken<Map<String, JsonElement>>() {
         }.getType();
-        java.util.Map<String, JsonElement> m = StringUtils.PLAIN_NUMBER_GSON.fromJson(json, mapType);
+        Map<String, JsonElement> m = StringUtils.PLAIN_NUMBER_GSON.fromJson(json, mapType);
 
         assertTrue(m.containsKey("fObj"));
         assertTrue(m.get("fObj").isJsonNull());

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAGUIAgentRunner.java

@@ -18,7 +18,6 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.commons.text.StringEscapeUtils;
 import org.apache.logging.log4j.Logger;
@@ -41,8 +40,6 @@
 import org.opensearch.transport.StreamTransportService;
 import org.opensearch.transport.client.Client;
 
-import com.google.common.annotations.VisibleForTesting;
-
 import lombok.Getter;
 import lombok.Setter;
 import lombok.extern.log4j.Log4j2;
@@ -73,18 +70,19 @@ public class AwsConnectorExecutor extends AbstractConnectorExecutor {
     @Getter
     private MLGuard mlGuard;
 
-    private final AtomicReference<SdkAsyncHttpClient> httpClientRef = new AtomicReference<>();
+    private SdkAsyncHttpClient httpClient;
 
     @Setter
     @Getter
     private StreamTransportService streamTransportService;
 
-    @Setter
-    private boolean connectorPrivateIpEnabled;
-
     public AwsConnectorExecutor(Connector connector) {
         super.initialize(connector);
         this.connector = (AwsConnector) connector;
+        Duration connectionTimeout = Duration.ofSeconds(super.getConnectorClientConfig().getConnectionTimeout());
+        Duration readTimeout = Duration.ofSeconds(super.getConnectorClientConfig().getReadTimeout());
+        Integer maxConnection = super.getConnectorClientConfig().getMaxConnections();
+        this.httpClient = MLHttpClientFactory.getAsyncHttpClient(connectionTimeout, readTimeout, maxConnection);
     }
 
     @Override
@@ -131,8 +129,7 @@ public void invokeRemoteService(
                     )
                 )
                 .build();
-            AccessController
-                .doPrivileged((PrivilegedExceptionAction<CompletableFuture<Void>>) () -> getHttpClient().execute(executeRequest));
+            AccessController.doPrivileged((PrivilegedExceptionAction<CompletableFuture<Void>>) () -> httpClient.execute(executeRequest));
         } catch (RuntimeException exception) {
             log.error("Failed to execute {} in aws connector: {}", action, exception.getMessage(), exception);
             actionListener.onFailure(exception);
@@ -157,7 +154,7 @@ public void invokeRemoteServiceStream(
             llmInterface = StringEscapeUtils.unescapeJava(llmInterface);
             validateLLMInterface(llmInterface);
 
-            StreamingHandler handler = StreamingHandlerFactory.createHandler(llmInterface, connector, getHttpClient(), null);
+            StreamingHandler handler = StreamingHandlerFactory.createHandler(llmInterface, connector, getHttpClient(), null, parameters);
             handler.startStream(action, parameters, payload, actionListener);
         } catch (Exception e) {
             log.error("Failed to execute streaming", e);
@@ -183,19 +180,4 @@ private void validateLLMInterface(String llmInterface) {
                 throw new IllegalArgumentException(String.format("Unsupported llm interface: %s", llmInterface));
         }
     }
-
-    @VisibleForTesting
-    protected SdkAsyncHttpClient getHttpClient() {
-        if (httpClientRef.get() == null) {
-            Duration connectionTimeout = Duration.ofSeconds(super.getConnectorClientConfig().getConnectionTimeout());
-            Duration readTimeout = Duration.ofSeconds(super.getConnectorClientConfig().getReadTimeout());
-            Integer maxConnection = super.getConnectorClientConfig().getMaxConnections();
-            this.httpClientRef
-                .compareAndSet(
-                    null,
-                    MLHttpClientFactory.getAsyncHttpClient(connectionTimeout, readTimeout, maxConnection, connectorPrivateIpEnabled)
-                );
-        }
-        return httpClientRef.get();
-    }
 }