ML Interface poc

Signed-off-by: Sicheng Song <sicheng.song@outlook.com>

Author: b4sjoo

build.gradle

@@ -23,7 +23,7 @@ buildscript {
         }
 
         common_utils_version = System.getProperty("common_utils.version", opensearch_build)
-        kotlin_version = System.getProperty("kotlin.version", "1.8.21")
+        kotlin_version = System.getProperty("kotlin.version", "1.9.23")
     }
 
     repositories {

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

@@ -160,8 +160,12 @@ public class CommonValue {
                         + "\" : {\"type\": \"flat_object\"},\n"
                         + "      \""
                         + AbstractConnector.ACTIONS_FIELD
+                        + "\" : {\"type\": \"flat_object\"},\n"
+                        + "      \""
+                        + AbstractConnector.MODEL_INTERFACE_FIELD
                         + "\" : {\"type\": \"flat_object\"}\n";
 
+
         public static final String ML_MODEL_INDEX_MAPPING = "{\n"
                         + "    \"_meta\": {\"schema_version\": "
                         + ML_MODEL_INDEX_SCHEMA_VERSION

common/src/main/java/org/opensearch/ml/common/connector/AbstractConnector.java

@@ -44,6 +44,7 @@ public abstract class AbstractConnector implements Connector {
     public static final String OWNER_FIELD = "owner";
     public static final String ACCESS_FIELD = "access";
     public static final String CLIENT_CONFIG_FIELD = "client_config";
+    public static final String MODEL_INTERFACE_FIELD = "model_interface";
 
 
     protected String name;
@@ -69,6 +70,7 @@ public abstract class AbstractConnector implements Connector {
     protected Instant lastUpdateTime;
     @Setter
     protected ConnectorClientConfig connectorClientConfig;
+    protected Map<String, String> modelInterface;
 
     protected Map<String, String> createPredictDecryptedHeaders(Map<String, String> headers) {
         if (headers == null) {

common/src/main/java/org/opensearch/ml/common/connector/AwsConnector.java

@@ -32,9 +32,9 @@ public class AwsConnector extends HttpConnector {
     public AwsConnector(String name, String description, String version, String protocol,
                          Map<String, String> parameters, Map<String, String> credential, List<ConnectorAction> actions,
                          List<String> backendRoles, AccessMode accessMode, User owner,
-                         ConnectorClientConfig connectorClientConfig) {
+                         ConnectorClientConfig connectorClientConfig, Map<String, String> modelInterface) {
         super(name, description, version, protocol, parameters, credential, actions, backendRoles, accessMode,
-                owner, connectorClientConfig);
+                owner, connectorClientConfig, modelInterface);
         validate();
     }
 

common/src/main/java/org/opensearch/ml/common/connector/Connector.java

@@ -60,6 +60,8 @@ public interface Connector extends ToXContentObject, Writeable {
 
     String getPredictHttpMethod();
 
+    Map<String, String> getModelInterface();
+
     <T> T createPredictPayload(Map<String, String> parameters);
 
     void decrypt(Function<String, String> function);

common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java

@@ -8,8 +8,7 @@
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.opensearch.ml.common.connector.ConnectorProtocols.HTTP;
 import static org.opensearch.ml.common.connector.ConnectorProtocols.validateProtocol;
-import static org.opensearch.ml.common.utils.StringUtils.getParameterMap;
-import static org.opensearch.ml.common.utils.StringUtils.isJson;
+import static org.opensearch.ml.common.utils.StringUtils.*;
 
 import java.io.IOException;
 import java.time.Instant;
@@ -53,7 +52,7 @@ public class HttpConnector extends AbstractConnector {
     public HttpConnector(String name, String description, String version, String protocol,
                          Map<String, String> parameters, Map<String, String> credential, List<ConnectorAction> actions,
                          List<String> backendRoles, AccessMode accessMode, User owner,
-                         ConnectorClientConfig connectorClientConfig) {
+                         ConnectorClientConfig connectorClientConfig, Map<String, String> modelInterface) {
         validateProtocol(protocol);
         this.name = name;
         this.description = description;
@@ -66,6 +65,7 @@ public HttpConnector(String name, String description, String version, String pro
         this.access = accessMode;
         this.owner = owner;
         this.connectorClientConfig = connectorClientConfig;
+        this.modelInterface = modelInterface;
 
     }
 
@@ -127,6 +127,9 @@ public HttpConnector(String protocol, XContentParser parser) throws IOException
                 case CLIENT_CONFIG_FIELD:
                     connectorClientConfig = ConnectorClientConfig.parse(parser);
                     break;
+                case MODEL_INTERFACE_FIELD:
+                    modelInterface =  getInterfaceMap(parser.map());
+                    break;
                 default:
                     parser.skipChildren();
                     break;
@@ -176,6 +179,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (connectorClientConfig != null) {
             builder.field(CLIENT_CONFIG_FIELD, connectorClientConfig);
         }
+        if (modelInterface != null) {
+            builder.field(MODEL_INTERFACE_FIELD, modelInterface);
+        }
         builder.endObject();
         return builder;
     }
@@ -219,6 +225,9 @@ private void parseFromStream(StreamInput input) throws IOException {
         if (input.readBoolean()) {
             this.connectorClientConfig = new ConnectorClientConfig(input);
         }
+        if (input.readBoolean()) {
+            this.modelInterface = input.readMap(StreamInput::readString, StreamInput::readString);
+        }
     }
 
     @Override
@@ -269,6 +278,12 @@ public void writeTo(StreamOutput out) throws IOException {
         } else {
             out.writeBoolean(false);
         }
+        if (modelInterface != null) {
+            out.writeBoolean(true);
+            out.writeMap(modelInterface, StreamOutput::writeString, StreamOutput::writeString);
+        } else {
+            out.writeBoolean(false);
+        }
     }
 
     @Override
@@ -304,6 +319,9 @@ public void update(MLCreateConnectorInput updateContent, Function<String, String
         if (updateContent.getConnectorClientConfig() != null) {
             this.connectorClientConfig = updateContent.getConnectorClientConfig();
         }
+        if (updateContent.getModelInterface() != null && updateContent.getModelInterface().size() > 0) {
+            this.modelInterface = updateContent.getModelInterface();
+        }
     }
 
     @Override

common/src/main/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInput.java

@@ -27,7 +27,7 @@
 import java.util.Map;
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
-import static org.opensearch.ml.common.utils.StringUtils.getParameterMap;
+import static org.opensearch.ml.common.utils.StringUtils.*;
 
 @Data
 public class MLCreateConnectorInput implements ToXContentObject, Writeable {
@@ -47,6 +47,7 @@ public class MLCreateConnectorInput implements ToXContentObject, Writeable {
     public static final String DRY_RUN_FIELD = "dry_run";
 
     private static final Version MINIMAL_SUPPORTED_VERSION_FOR_CLIENT_CONFIG = Version.V_2_13_0;
+    private static final Version MINIMAL_SUPPORTED_VERSION_FOR_MODEL_INTERFACE = Version.V_2_13_0;
 
     public static final String DRY_RUN_CONNECTOR_NAME = "dryRunConnector";
 
@@ -63,6 +64,7 @@ public class MLCreateConnectorInput implements ToXContentObject, Writeable {
     private boolean dryRun;
     private boolean updateConnector;
     private ConnectorClientConfig connectorClientConfig;
+    private Map<String, String> modelInterface;
 
 
     @Builder(toBuilder = true)
@@ -78,7 +80,8 @@ public MLCreateConnectorInput(String name,
                                   AccessMode access,
                                   boolean dryRun,
                                   boolean updateConnector,
-                                  ConnectorClientConfig connectorClientConfig
+                                  ConnectorClientConfig connectorClientConfig,
+                                  Map<String, String> modelInterface
 
     ) {
         if (!dryRun && !updateConnector) {
@@ -105,6 +108,7 @@ public MLCreateConnectorInput(String name,
         this.dryRun = dryRun;
         this.updateConnector = updateConnector;
         this.connectorClientConfig = connectorClientConfig;
+        this.modelInterface = modelInterface;
 
     }
 
@@ -125,6 +129,7 @@ public static MLCreateConnectorInput parse(XContentParser parser, boolean update
         AccessMode access = null;
         boolean dryRun = false;
         ConnectorClientConfig connectorClientConfig = null;
+        Map<String, String> modelInterface = new HashMap<>();
 
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -176,13 +181,16 @@ public static MLCreateConnectorInput parse(XContentParser parser, boolean update
                 case AbstractConnector.CLIENT_CONFIG_FIELD:
                     connectorClientConfig = ConnectorClientConfig.parse(parser);
                     break;
+                case AbstractConnector.MODEL_INTERFACE_FIELD:
+                    modelInterface = getInterfaceMap(parser.map());
+                    break;
                 default:
                     parser.skipChildren();
                     break;
             }
         }
         return new MLCreateConnectorInput(name, description, version, protocol, parameters, credential, actions,
-                backendRoles, addAllBackendRoles, access, dryRun, updateConnector, connectorClientConfig);
+                backendRoles, addAllBackendRoles, access, dryRun, updateConnector, connectorClientConfig, modelInterface);
     }
 
     @Override
@@ -221,6 +229,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (connectorClientConfig != null) {
             builder.field(AbstractConnector.CLIENT_CONFIG_FIELD, connectorClientConfig);
         }
+        if (modelInterface != null) {
+            builder.field(AbstractConnector.MODEL_INTERFACE_FIELD, modelInterface);
+        }
         builder.endObject();
         return builder;
     }
@@ -276,6 +287,14 @@ public void writeTo(StreamOutput output) throws IOException {
                 output.writeBoolean(false);
             }
         }
+        if (streamOutputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_MODEL_INTERFACE)) {
+            if (modelInterface != null) {
+                output.writeBoolean(true);
+                output.writeMap(modelInterface, StreamOutput::writeString, StreamOutput::writeString);
+            } else {
+                output.writeBoolean(false);
+            }
+        }
     }
 
     public MLCreateConnectorInput(StreamInput input) throws IOException {
@@ -311,6 +330,10 @@ public MLCreateConnectorInput(StreamInput input) throws IOException {
                 this.connectorClientConfig = new ConnectorClientConfig(input);
             }
         }
-
+        if (streamInputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_MODEL_INTERFACE)) {
+            if (input.readBoolean()) {
+                modelInterface = input.readMap(StreamInput::readString, StreamInput::readString);
+            }
+        }
     }
 }

common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java

@@ -19,10 +19,7 @@
 import java.security.AccessController;
 import java.security.PrivilegedActionException;
 import java.security.PrivilegedExceptionAction;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -94,6 +91,28 @@ public static Map<String, Object> fromJson(String jsonStr, String defaultKey) {
         return result;
     }
 
+    public static Map<String, String> getInterfaceMap(Map<String, ?> interfaceObjs) {
+        Map<String, String> parameters = new HashMap<>();
+        for (String key : interfaceObjs.keySet()) {
+            if (Objects.equals(key, "input") || Objects.equals(key, "output")) {
+                Object value = interfaceObjs.get(key);
+                try {
+                    AccessController.doPrivileged((PrivilegedExceptionAction<Void>) () -> {
+                        if (value instanceof String) {
+                            parameters.put(key, (String)value);
+                        } else {
+                            parameters.put(key, gson.toJson(value));
+                        }
+                        return null;
+                    });
+                } catch (PrivilegedActionException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        }
+        return parameters;
+    }
+
     @SuppressWarnings("removal")
     public static Map<String, String> getParameterMap(Map<String, ?> parameterObjs) {
         Map<String, String> parameters = new HashMap<>();

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

@@ -115,6 +115,23 @@ public void getParameterMap() {
         Assert.assertEquals("[1.01,\"abc\"]", parameterMap.get("key5"));
     }
 
+    @Test
+    public void getInterfaceMap() {
+        Map<String, Object> parameters = new HashMap<>();
+        parameters.put("input", "value1");
+        parameters.put("output", 2);
+        parameters.put("key3", 2.1);
+        parameters.put("key4", new int[]{10, 20});
+        parameters.put("key5", new Object[]{1.01, "abc"});
+        Map<String, String> interfaceMap = StringUtils.getInterfaceMap(parameters);
+        Assert.assertEquals(2, interfaceMap.size());
+        Assert.assertEquals("value1", interfaceMap.get("input"));
+        Assert.assertEquals("2", interfaceMap.get("output"));
+        Assert.assertNull(interfaceMap.get("key3"));
+        Assert.assertNull(interfaceMap.get("key4"));
+        Assert.assertNull(interfaceMap.get("key5"));
+    }
+
     @Test
     public void processTextDocs() {
         List<String> processedDocs = StringUtils.processTextDocs(Arrays.asList("abc \n\n123\"4", null, "[1.01,\"abc\"]"));

ml-algorithms/build.gradle

@@ -28,6 +28,7 @@ dependencies {
        exclude group: "org.jetbrains", module: "annotations"
     }
     implementation group: 'org.apache.commons', name: 'commons-text', version: '1.10.0'
+    implementation group: 'com.github.erosb', name: 'json-sKema', version: '0.15.0'
     implementation group: 'org.reflections', name: 'reflections', version: '0.9.12'
     implementation group: 'org.tribuo', name: 'tribuo-clustering-kmeans', version: '4.2.1'
     implementation group: 'org.tribuo', name: 'tribuo-regression-sgd', version: '4.2.1'