diff --git a/plugin/src/test/java/org/opensearch/ml/utils/MLNodeUtilsTests.java b/plugin/src/test/java/org/opensearch/ml/utils/MLNodeUtilsTests.java index c0f6d35ee7..40bb230bbf 100644 --- a/plugin/src/test/java/org/opensearch/ml/utils/MLNodeUtilsTests.java +++ b/plugin/src/test/java/org/opensearch/ml/utils/MLNodeUtilsTests.java @@ -6,6 +6,9 @@ package org.opensearch.ml.utils; import static java.util.Collections.emptyMap; +import static org.opensearch.ml.common.utils.ModelInterfaceUtils.BEDROCK_COHERE_EMBED_ENGLISH_V3_MODEL_INTERFACE; +import static org.opensearch.ml.common.utils.ModelInterfaceUtils.BEDROCK_TITAN_EMBED_MULTI_MODAL_V1_MODEL_INTERFACE; +import static org.opensearch.ml.common.utils.ModelInterfaceUtils.BEDROCK_TITAN_EMBED_TEXT_V1_MODEL_INTERFACE; import static org.opensearch.ml.utils.TestHelper.ML_ROLE; import java.io.IOException; @@ -68,139 +71,43 @@ public void testValidateSchema() throws IOException { @Test public void testValidateEmbeddingInputWithGeneralEmbeddingRemoteSchema() 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 schema = BEDROCK_COHERE_EMBED_ENGLISH_V3_MODEL_INTERFACE.get("input"); String json = "{\"text_docs\":[ \"today is sunny\", \"today is sunny\"]}"; 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 schema = BEDROCK_COHERE_EMBED_ENGLISH_V3_MODEL_INTERFACE.get("input"); 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 schema = BEDROCK_TITAN_EMBED_TEXT_V1_MODEL_INTERFACE.get("input"); 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 schema = BEDROCK_TITAN_EMBED_TEXT_V1_MODEL_INTERFACE.get("input"); 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" - + "}"; + public void testValidateEmbeddingInputWithTitanMultiModalRemoteSchema() throws IOException { + String schema = BEDROCK_TITAN_EMBED_MULTI_MODAL_V1_MODEL_INTERFACE.get("input"); 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" + public void testValidateRemoteInputWithTitanMultiModalRemoteSchema() throws IOException { + String schema = BEDROCK_TITAN_EMBED_MULTI_MODAL_V1_MODEL_INTERFACE.get("input"); + String json = "{\n" + " \"parameters\": {\n" + " \"inputText\": \"Say this is a test\",\n" + " \"inputImage\": \"/9jk=\"\n"