add batch predict job actiontype in connector

Signed-off-by: Xun Zhang <xunzh@amazon.com>

Author: Zhangxunmt

common/src/main/java/org/opensearch/ml/common/PredictMode.java

@@ -0,0 +1,18 @@
+package org.opensearch.ml.common;
+
+import java.util.Locale;
+
+public enum PredictMode {
+    REAL_TIME,
+    BATCH,
+    ASYNC,
+    STREAMING;
+
+    public static PredictMode from(String value) {
+        try {
+            return PredictMode.valueOf(value.toUpperCase(Locale.ROOT));
+        } catch (Exception e) {
+            throw new IllegalArgumentException("Wrong Predict mode");
+        }
+    }
+}

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,
+        BATCH_PREDICT_JOB
     }
 }

common/src/main/java/org/opensearch/ml/common/dataset/remote/RemoteInferenceInputDataSet.java

@@ -10,6 +10,7 @@
 import lombok.Setter;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.ml.common.PredictMode;
 import org.opensearch.ml.common.annotation.InputDataSet;
 import org.opensearch.ml.common.dataset.MLInputDataType;
 import org.opensearch.ml.common.dataset.MLInputDataset;
@@ -23,11 +24,17 @@ public class RemoteInferenceInputDataSet extends MLInputDataset {
 
     @Setter
     private Map<String, String> parameters;
+    private PredictMode predictMode;
 
     @Builder(toBuilder = true)
-    public RemoteInferenceInputDataSet(Map<String, String> parameters) {
+    public RemoteInferenceInputDataSet(Map<String, String> parameters, PredictMode predictMode) {
         super(MLInputDataType.REMOTE);
         this.parameters = parameters;
+        this.predictMode = predictMode;
+    }
+
+    public RemoteInferenceInputDataSet(Map<String, String> parameters) {
+        this(parameters, null);
     }
 
     public RemoteInferenceInputDataSet(StreamInput streamInput) throws IOException {

common/src/main/java/org/opensearch/ml/common/input/remote/RemoteInferenceMLInput.java

@@ -9,6 +9,7 @@
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.ml.common.FunctionName;
+import org.opensearch.ml.common.PredictMode;
 import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
 import org.opensearch.ml.common.input.MLInput;
 import org.opensearch.ml.common.utils.StringUtils;
@@ -21,6 +22,7 @@
 @org.opensearch.ml.common.annotation.MLInput(functionNames = {FunctionName.REMOTE})
 public class RemoteInferenceMLInput extends MLInput {
     public static final String PARAMETERS_FIELD = "parameters";
+    public static final String PREDICT_MODE_FIELD = "mode";
 
     public RemoteInferenceMLInput(StreamInput in) throws IOException {
         super(in);
@@ -34,21 +36,26 @@ public void writeTo(StreamOutput out) throws IOException {
     public RemoteInferenceMLInput(XContentParser parser, FunctionName functionName) throws IOException {
         super();
         this.algorithm = functionName;
+        Map<String, String> parameters = null;
+        PredictMode predictMode = null;
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
             String fieldName = parser.currentName();
             parser.nextToken();
 
             switch (fieldName) {
                 case PARAMETERS_FIELD:
-                    Map<String, String> parameters = StringUtils.getParameterMap(parser.map());
-                    inputDataset = new RemoteInferenceInputDataSet(parameters);
+                    parameters = StringUtils.getParameterMap(parser.map());
                     break;
+                case PREDICT_MODE_FIELD:
+                    predictMode = PredictMode.from(parser.text());
                 default:
                     parser.skipChildren();
                     break;
             }
         }
+        predictMode = predictMode == null? PredictMode.REAL_TIME:predictMode;
+        inputDataset = new RemoteInferenceInputDataSet(parameters, predictMode);
     }
 
 }

plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java

@@ -132,6 +132,7 @@ private MLCommonsSettings() {}
             ImmutableList
                 .of(
                     "^https://runtime\\.sagemaker\\..*[a-z0-9-]\\.amazonaws\\.com/.*$",
+                    "^https://api\\.sagemaker\\..*[a-z0-9-]\\.amazonaws\\.com/.*$",
                     "^https://api\\.openai\\.com/.*$",
                     "^https://api\\.cohere\\.ai/.*$",
                     "^https://bedrock-runtime\\..*[a-z0-9-]\\.amazonaws\\.com/.*$"