Add UT for both embedding and remote cases for all remote embedding s…

…chema Signed-off-by: b4sjoo <[email protected]>
opensearch-project · Dec 20, 2024 · 9ff2d5d · 9ff2d5d
1 parent 1f91102
commit 9ff2d5d
Showing 1 changed file with 119 additions and 1 deletion.
diff --git a/plugin/src/test/java/org/opensearch/ml/utils/MLNodeUtilsTests.java b/plugin/src/test/java/org/opensearch/ml/utils/MLNodeUtilsTests.java
@@ -67,7 +67,7 @@ public void testValidateSchema() throws IOException {
     }
 
     @Test
-    public void testValidateEmbeddingInputWithRemoteSchema() throws IOException {
+    public void testValidateEmbeddingInputWithGeneralEmbeddingRemoteSchema() throws IOException {
         String schema = "{\n"
             + "    \"type\": \"object\",\n"
             + "    \"properties\": {\n"
@@ -91,6 +91,124 @@ public void testValidateEmbeddingInputWithRemoteSchema() throws IOException {
         MLNodeUtils.validateSchema(schema, json);
     }
 
+    @Test
+    public void testValidateRemoteInputWithGeneralEmbeddingRemoteSchema() throws IOException {
+        String schema = "{\n"
+            + "    \"type\": \"object\",\n"
+            + "    \"properties\": {\n"
+            + "        \"parameters\": {\n"
+            + "            \"type\": \"object\",\n"
+            + "            \"properties\": {\n"
+            + "                \"texts\": {\n"
+            + "                    \"type\": \"array\",\n"
+            + "                    \"items\": {\n"
+            + "                        \"type\": \"string\"\n"
+            + "                    }\n"
+            + "                }\n"
+            + "            },\n"
+            + "            \"required\": [\n"
+            + "                \"texts\"\n"
+            + "            ]\n"
+            + "        }\n"
+            + "    }\n"
+            + "}";
+        String json = "{\"parameters\": {\"texts\": [\"Hello\",\"world\"]}}";
+        MLNodeUtils.validateSchema(schema, json);
+    }
+
+    @Test
+    public void testValidateEmbeddingInputWithTitanTextRemoteSchema() throws IOException {
+        String schema = "{\n"
+            + "    \"type\": \"object\",\n"
+            + "    \"properties\": {\n"
+            + "        \"parameters\": {\n"
+            + "            \"type\": \"object\",\n"
+            + "            \"properties\": {\n"
+            + "                \"inputText\": {\n"
+            + "                    \"type\": \"string\"\n"
+            + "                }\n"
+            + "            },\n"
+            + "            \"required\": [\n"
+            + "                \"inputText\"\n"
+            + "            ]\n"
+            + "        }\n"
+            + "    }\n"
+            + "}";
+        String json = "{\"text_docs\":[ \"today is sunny\", \"today is sunny\"]}";
+        MLNodeUtils.validateSchema(schema, json);
+    }
+
+    @Test
+    public void testValidateRemoteInputWithTitanTextRemoteSchema() throws IOException {
+        String schema = "{\n"
+            + "    \"type\": \"object\",\n"
+            + "    \"properties\": {\n"
+            + "        \"parameters\": {\n"
+            + "            \"type\": \"object\",\n"
+            + "            \"properties\": {\n"
+            + "                \"inputText\": {\n"
+            + "                    \"type\": \"string\"\n"
+            + "                }\n"
+            + "            },\n"
+            + "            \"required\": [\n"
+            + "                \"inputText\"\n"
+            + "            ]\n"
+            + "        }\n"
+            + "    }\n"
+            + "}";
+        String json = "{\"parameters\": {\"inputText\": \"Say this is a test\"}}";
+        MLNodeUtils.validateSchema(schema, json);
+    }
+
+    @Test
+    public void testValidateEmbeddingInputWithTitanImageRemoteSchema() throws IOException {
+        String schema = "{\n"
+            + "    \"type\": \"object\",\n"
+            + "    \"properties\": {\n"
+            + "        \"parameters\": {\n"
+            + "            \"type\": \"object\",\n"
+            + "            \"properties\": {\n"
+            + "                \"inputText\": {\n"
+            + "                    \"type\": \"string\"\n"
+            + "                },\n"
+            + "                \"inputImage\": {\n"
+            + "                    \"type\": \"string\"\n"
+            + "                }\n"
+            + "            }\n"
+            + "        }\n"
+            + "    }\n"
+            + "}";
+        String json = "{\"text_docs\":[ \"today is sunny\", \"today is sunny\"]}";
+        MLNodeUtils.validateSchema(schema, json);
+    }
+
+    @Test
+    public void testValidateRemoteInputWithTitanImageRemoteSchema() throws IOException {
+        String schema = "{\n"
+            + "    \"type\": \"object\",\n"
+            + "    \"properties\": {\n"
+            + "        \"parameters\": {\n"
+            + "            \"type\": \"object\",\n"
+            + "            \"properties\": {\n"
+            + "                \"inputText\": {\n"
+            + "                    \"type\": \"string\"\n"
+            + "                },\n"
+            + "                \"inputImage\": {\n"
+            + "                    \"type\": \"string\"\n"
+            + "                }\n"
+            + "            }\n"
+            + "        }\n"
+            + "    }\n"
+            + "}";
+        String json = "{{\n"
+            + "  \"parameters\": {\n"
+            + "    \"inputText\": \"Say this is a test\",\n"
+            + "    \"inputImage\": \"/9jk=\"\n"
+            + "  }\n"
+            + "}";
+        MLNodeUtils.validateSchema(schema, json);
+    }
+
     @Test
     public void testProcessRemoteInferenceInputDataSetParametersValueNoParameters() throws IOException {
         String json = "{\"key1\":\"foo\",\"key2\":123,\"key3\":true}";