add connector tool

Signed-off-by: Yaliang Wu <[email protected]>
ylwu-amzn · Jun 6, 2024 · a196399 · a196399
1 parent f6260c2
commit a196399
Show file tree

Hide file tree

Showing 15 changed files with 180 additions and 144 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/connector/AbstractConnector.java b/common/src/main/java/org/opensearch/ml/common/connector/AbstractConnector.java
@@ -70,7 +70,7 @@ public abstract class AbstractConnector implements Connector {
     @Setter
     protected ConnectorClientConfig connectorClientConfig;
 
-    protected Map<String, String> createPredictDecryptedHeaders(Map<String, String> headers) {
+    protected Map<String, String> createDecryptedHeaders(Map<String, String> headers) {
         if (headers == null) {
             return null;
         }
@@ -116,9 +116,9 @@ public <T> void parseResponse(T response, List<ModelTensor> modelTensors, boolea
     }
 
     @Override
-    public Optional<ConnectorAction> findPredictAction() {
+    public Optional<ConnectorAction> findAction(String action) {
         if (actions != null) {
-            return actions.stream().filter(a -> a.getActionType() == ConnectorAction.ActionType.PREDICT).findFirst();
+            return actions.stream().filter(a -> a.getActionType().name().equalsIgnoreCase(action)).findFirst();
         }
         return Optional.empty();
     }
@@ -131,17 +131,17 @@ public void removeCredential() {
     }
 
     @Override
-    public String getPredictEndpoint(Map<String, String> parameters) {
-        Optional<ConnectorAction> predictAction = findPredictAction();
-        if (!predictAction.isPresent()) {
+    public String getActionEndpoint(String action, Map<String, String> parameters) {
+        Optional<ConnectorAction> connectorAction = findAction(action);
+        if (!connectorAction.isPresent()) {
             return null;
         }
-        String predictEndpoint = predictAction.get().getUrl();
+        String actionEndpoint = connectorAction.get().getUrl();
         if (parameters != null && parameters.size() > 0) {
             StringSubstitutor substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
-            predictEndpoint = substitutor.replace(predictEndpoint);
+            actionEndpoint = substitutor.replace(actionEndpoint);
         }
-        return predictEndpoint;
+        return actionEndpoint;
     }
 
 }
diff --git a/common/src/main/java/org/opensearch/ml/common/connector/Connector.java b/common/src/main/java/org/opensearch/ml/common/connector/Connector.java
@@ -56,18 +56,18 @@ public interface Connector extends ToXContentObject, Writeable {
 
     ConnectorClientConfig getConnectorClientConfig();
 
-    String getPredictEndpoint(Map<String, String> parameters);
+    String getActionEndpoint(String action, Map<String, String> parameters);
 
-    String getPredictHttpMethod();
+    String getActionHttpMethod(String action);
 
-    <T> T createPredictPayload(Map<String, String> parameters);
+    <T> T createPayload(String action, Map<String, String> parameters);
 
-    void decrypt(Function<String, String> function);
+    void decrypt(String action, Function<String, String> function);
     void encrypt(Function<String, String> function);
 
     Connector cloneConnector();
 
-    Optional<ConnectorAction> findPredictAction();
+    Optional<ConnectorAction> findAction(String action);
 
     void removeCredential();
 

diff --git a/common/src/main/java/org/opensearch/ml/common/connector/ConnectorAction.java b/common/src/main/java/org/opensearch/ml/common/connector/ConnectorAction.java
@@ -182,6 +182,7 @@ public static ConnectorAction parse(XContentParser parser) throws IOException {
     }
 
     public enum ActionType {
-        PREDICT
+        PREDICT,
+        EXECUTE
     }
 }
diff --git a/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java b/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java
@@ -307,10 +307,10 @@ public void update(MLCreateConnectorInput updateContent, Function<String, String
     }
 
     @Override
-    public  <T> T createPredictPayload(Map<String, String> parameters) {
-        Optional<ConnectorAction> predictAction = findPredictAction();
-        if (predictAction.isPresent() && predictAction.get().getRequestBody() != null) {
-            String payload = predictAction.get().getRequestBody();
+    public  <T> T createPayload(String action, Map<String, String> parameters) {
+        Optional<ConnectorAction> connectorAction = findAction(action);
+        if (connectorAction.isPresent() && connectorAction.get().getRequestBody() != null) {
+            String payload = connectorAction.get().getRequestBody();
             payload = fillNullParameters(parameters, payload);
             StringSubstitutor substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
             payload = substitutor.replace(payload);
@@ -323,6 +323,7 @@ public  <T> T createPredictPayload(Map<String, String> parameters) {
         return (T) parameters.get("http_body");
     }
 
+
     protected String fillNullParameters(Map<String, String> parameters, String payload) {
         List<String> bodyParams = findStringParametersWithNullDefaultValue(payload);
         String newPayload = payload;
@@ -348,15 +349,15 @@ private List<String> findStringParametersWithNullDefaultValue(String input) {
     }
 
     @Override
-    public void decrypt(Function<String, String> function) {
+    public void decrypt(String action, Function<String, String> function) {
         Map<String, String> decrypted = new HashMap<>();
         for (String key : credential.keySet()) {
             decrypted.put(key, function.apply(credential.get(key)));
         }
         this.decryptedCredential = decrypted;
-        Optional<ConnectorAction> predictAction = findPredictAction();
-        Map<String, String> headers = predictAction.isPresent() ? predictAction.get().getHeaders() : null;
-        this.decryptedHeaders = createPredictDecryptedHeaders(headers);
+        Optional<ConnectorAction> connectorAction = findAction(action);
+        Map<String, String> headers = connectorAction.isPresent() ? connectorAction.get().getHeaders() : null;
+        this.decryptedHeaders = createDecryptedHeaders(headers);
     }
 
     @Override
@@ -378,8 +379,8 @@ public void encrypt(Function<String, String> function) {
         }
     }
 
-    public String getPredictHttpMethod() {
-        return findPredictAction().get().getMethod();
+    public String getActionHttpMethod(String action) {
+        return findAction(action).get().getMethod();
     }
 
 }
diff --git a/...src/main/java/org/opensearch/ml/common/transport/connector/MLExecuteConnectorRequest.java b/...src/main/java/org/opensearch/ml/common/transport/connector/MLExecuteConnectorRequest.java
@@ -34,29 +34,32 @@
 public class MLExecuteConnectorRequest extends MLTaskRequest {
 
     String connectorId;
+    String connectorAction;
     MLInput mlInput;
     @Setter
     User user;
 
     @Builder
-    public MLExecuteConnectorRequest(String connectorId, MLInput mlInput, boolean dispatchTask, User user) {
+    public MLExecuteConnectorRequest(String connectorId, String connectorAction, MLInput mlInput, boolean dispatchTask, User user) {
         super(dispatchTask);
         this.mlInput = mlInput;
+        this.connectorAction = connectorAction == null ? "predict" : connectorAction;
         this.connectorId = connectorId;
         this.user = user;
     }
 
-    public MLExecuteConnectorRequest(String connectorId, MLInput mlInput) {
-        this(connectorId, mlInput, true, null);
+    public MLExecuteConnectorRequest(String connectorId, String connectorAction, MLInput mlInput) {
+        this(connectorId, connectorAction, mlInput, true, null);
     }
 
-    public MLExecuteConnectorRequest(String connectorId, MLInput mlInput, User user) {
-        this(connectorId, mlInput, true, user);
+    public MLExecuteConnectorRequest(String connectorId, String connectorAction, MLInput mlInput, User user) {
+        this(connectorId, connectorAction, mlInput, true, user);
     }
 
     public MLExecuteConnectorRequest(StreamInput in) throws IOException {
         super(in);
-        this.connectorId = in.readOptionalString();
+        this.connectorId = in.readString();
+        this.connectorAction = in.readString();
         this.mlInput = new MLInput(in);
         if (in.readBoolean()) {
             this.user = new User(in);
@@ -66,7 +69,8 @@ public MLExecuteConnectorRequest(StreamInput in) throws IOException {
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
-        out.writeOptionalString(this.connectorId);
+        out.writeString(this.connectorId);
+        out.writeString(this.connectorAction);
         this.mlInput.writeTo(out);
         if (user != null) {
             out.writeBoolean(true);

diff --git a/...rithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutor.java b/...rithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutor.java
@@ -71,7 +71,8 @@ public AwsConnectorExecutor(Connector connector) {
 
     @SuppressWarnings("removal")
     @Override
-    public void invokeRemoteModel(
+    public void invokeRemoteService(
+        String action,
         MLInput mlInput,
         Map<String, String> parameters,
         String payload,
@@ -80,7 +81,7 @@ public void invokeRemoteModel(
         ActionListener<List<ModelTensors>> actionListener
     ) {
         try {
-            SdkHttpFullRequest request = ConnectorUtils.buildSdkRequest(connector, parameters, payload, POST);
+            SdkHttpFullRequest request = ConnectorUtils.buildSdkRequest(action, connector, parameters, payload, POST);
             AsyncExecuteRequest executeRequest = AsyncExecuteRequest
                 .builder()
                 .request(signRequest(request))
@@ -93,17 +94,18 @@ public void invokeRemoteModel(
                         tensorOutputs,
                         connector,
                         scriptService,
-                        mlGuard
+                        mlGuard,
+                        action
                     )
                 )
                 .build();
             AccessController.doPrivileged((PrivilegedExceptionAction<CompletableFuture<Void>>) () -> httpClient.execute(executeRequest));
         } catch (RuntimeException exception) {
-            log.error("Failed to execute predict in aws connector: " + exception.getMessage(), exception);
+            log.error("Failed to execute {} in aws connector: {}", action, exception.getMessage(), exception);
             actionListener.onFailure(exception);
         } catch (Throwable e) {
-            log.error("Failed to execute predict in aws connector", e);
-            actionListener.onFailure(new MLException("Fail to execute predict in aws connector", e));
+            log.error("Failed to execute {} in aws connector", action, e);
+            actionListener.onFailure(new MLException("Fail to execute " + action + " in aws connector", e));
         }
     }
 

diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/ConnectorUtils.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/ConnectorUtils.java
@@ -63,6 +63,7 @@ public class ConnectorUtils {
     }
 
     public static RemoteInferenceInputDataSet processInput(
+        String action,
         MLInput mlInput,
         Connector connector,
         Map<String, String> parameters,
@@ -71,22 +72,23 @@ public static RemoteInferenceInputDataSet processInput(
         if (mlInput == null) {
             throw new IllegalArgumentException("Input is null");
         }
-        Optional<ConnectorAction> predictAction = connector.findPredictAction();
+        Optional<ConnectorAction> predictAction = connector.findAction(action);
         if (predictAction.isEmpty()) {
             throw new IllegalArgumentException("no predict action found");
         }
-        RemoteInferenceInputDataSet inputData = processMLInput(mlInput, connector, parameters, scriptService);
+        RemoteInferenceInputDataSet inputData = processMLInput(action, mlInput, connector, parameters, scriptService);
         escapeRemoteInferenceInputData(inputData);
         return inputData;
     }
 
     private static RemoteInferenceInputDataSet processMLInput(
+        String action,
         MLInput mlInput,
         Connector connector,
         Map<String, String> parameters,
         ScriptService scriptService
     ) {
-        String preProcessFunction = getPreprocessFunction(mlInput, connector);
+        String preProcessFunction = getPreprocessFunction(action, mlInput, connector);
         if (preProcessFunction == null) {
             if (mlInput.getInputDataset() instanceof RemoteInferenceInputDataSet) {
                 return (RemoteInferenceInputDataSet) mlInput.getInputDataset();
@@ -168,8 +170,8 @@ public static void escapeRemoteInferenceInputData(RemoteInferenceInputDataSet in
         }
     }
 
-    private static String getPreprocessFunction(MLInput mlInput, Connector connector) {
-        Optional<ConnectorAction> predictAction = connector.findPredictAction();
+    private static String getPreprocessFunction(String action, MLInput mlInput, Connector connector) {
+        Optional<ConnectorAction> predictAction = connector.findAction(action);
         String preProcessFunction = predictAction.get().getPreProcessFunction();
         if (preProcessFunction != null) {
             return preProcessFunction;
@@ -181,6 +183,7 @@ private static String getPreprocessFunction(MLInput mlInput, Connector connector
     }
 
     public static ModelTensors processOutput(
+        String action,
         String modelResponse,
         Connector connector,
         ScriptService scriptService,
@@ -194,7 +197,7 @@ public static ModelTensors processOutput(
             throw new IllegalArgumentException("guardrails triggered for LLM output");
         }
         List<ModelTensor> modelTensors = new ArrayList<>();
-        Optional<ConnectorAction> predictAction = connector.findPredictAction();
+        Optional<ConnectorAction> predictAction = connector.findAction(action);
         if (predictAction.isEmpty()) {
             throw new IllegalArgumentException("no predict action found");
         }
@@ -263,6 +266,7 @@ public static SdkHttpFullRequest signRequest(
     }
 
     public static SdkHttpFullRequest buildSdkRequest(
+        String action,
         Connector connector,
         Map<String, String> parameters,
         String payload,
@@ -279,7 +283,7 @@ public static SdkHttpFullRequest buildSdkRequest(
             log.error("Content length is 0. Aborting request to remote model");
             throw new IllegalArgumentException("Content length is 0. Aborting request to remote model");
         }
-        String endpoint = connector.getPredictEndpoint(parameters);
+        String endpoint = connector.getActionEndpoint(action, parameters);
         SdkHttpFullRequest.Builder builder = SdkHttpFullRequest
             .builder()
             .method(method)

diff --git a/...s/src/main/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutor.java b/...s/src/main/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutor.java
@@ -75,7 +75,8 @@ public HttpJsonConnectorExecutor(Connector connector) {
 
     @SuppressWarnings("removal")
     @Override
-    public void invokeRemoteModel(
+    public void invokeRemoteService(
+        String action,
         MLInput mlInput,
         Map<String, String> parameters,
         String payload,
@@ -85,15 +86,15 @@ public void invokeRemoteModel(
     ) {
         try {
             SdkHttpFullRequest request;
-            switch (connector.getPredictHttpMethod().toUpperCase(Locale.ROOT)) {
+            switch (connector.getActionHttpMethod(action).toUpperCase(Locale.ROOT)) {
                 case "POST":
                     log.debug("original payload to remote model: " + payload);
-                    validateHttpClientParameters(parameters);
-                    request = ConnectorUtils.buildSdkRequest(connector, parameters, payload, POST);
+                    validateHttpClientParameters(action, parameters);
+                    request = ConnectorUtils.buildSdkRequest(action, connector, parameters, payload, POST);
                     break;
                 case "GET":
-                    validateHttpClientParameters(parameters);
-                    request = ConnectorUtils.buildSdkRequest(connector, parameters, null, GET);
+                    validateHttpClientParameters(action, parameters);
+                    request = ConnectorUtils.buildSdkRequest(action, connector, parameters, null, GET);
                     break;
                 default:
                     throw new IllegalArgumentException("unsupported http method");
@@ -110,7 +111,8 @@ public void invokeRemoteModel(
                         tensorOutputs,
                         connector,
                         scriptService,
-                        mlGuard
+                        mlGuard,
+                        action
                     )
                 )
                 .build();
@@ -124,8 +126,8 @@ public void invokeRemoteModel(
         }
     }
 
-    private void validateHttpClientParameters(Map<String, String> parameters) throws Exception {
-        String endpoint = connector.getPredictEndpoint(parameters);
+    private void validateHttpClientParameters(String action, Map<String, String> parameters) throws Exception {
+        String endpoint = connector.getActionEndpoint(action, parameters);
         URL url = new URL(endpoint);
         String protocol = url.getProtocol();
         String host = url.getHost();
-Original file line number
+Diff line change
@@ Expand Up @@
         }
         public enum ActionType {
-            PREDICT
+            PREDICT,
+            EXECUTE
         }
     }