CreateAnomalyDetectorTool supports empty model_type (#457)

* CreateAnomalyDetectorTool supports empty model_type Signed-off-by: gaobinlong <[email protected]> * Optimize code Signed-off-by: gaobinlong <[email protected]> * Add some comment Signed-off-by: gaobinlong <[email protected]> --------- Signed-off-by: gaobinlong <[email protected]>
opensearch-project · Nov 14, 2024 · cfc88d8 · cfc88d8
1 parent d0c42a8
commit cfc88d8
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/src/main/java/org/opensearch/agent/tools/CreateAnomalyDetectorTool.java b/src/main/java/org/opensearch/agent/tools/CreateAnomalyDetectorTool.java
@@ -429,7 +429,11 @@ public CreateAnomalyDetectorTool create(Map<String, Object> map) {
                 throw new IllegalArgumentException("model_id cannot be empty.");
             }
             String modelType = (String) map.getOrDefault("model_type", ModelType.CLAUDE.toString());
-            if (!ModelType.OPENAI.toString().equalsIgnoreCase(modelType) && !ModelType.CLAUDE.toString().equalsIgnoreCase(modelType)) {
+            // if model type is empty, use the default value
+            if (modelType.isEmpty()) {
+                modelType = ModelType.CLAUDE.toString();
+            } else if (!ModelType.OPENAI.toString().equalsIgnoreCase(modelType)
+                && !ModelType.CLAUDE.toString().equalsIgnoreCase(modelType)) {
                 throw new IllegalArgumentException("Unsupported model_type: " + modelType);
             }
             String prompt = (String) map.getOrDefault("prompt", "");

diff --git a/src/test/java/org/opensearch/agent/tools/CreateAnomalyDetectorToolTests.java b/src/test/java/org/opensearch/agent/tools/CreateAnomalyDetectorToolTests.java
@@ -260,6 +260,22 @@ public void testToolWithCustomPrompt() {
             );
     }
 
+    @Test
+    public void testToolWithEmptyModelType() {
+        CreateAnomalyDetectorTool tool = CreateAnomalyDetectorTool.Factory
+            .getInstance()
+            .create(ImmutableMap.of("model_id", "modelId", "model_type", ""));
+        assertEquals(CreateAnomalyDetectorTool.TYPE, tool.getName());
+        assertEquals("modelId", tool.getModelId());
+        assertEquals("CLAUDE", tool.getModelType().toString());
+
+        tool
+            .run(
+                ImmutableMap.of("index", mockedIndexName),
+                ActionListener.<String>wrap(response -> assertEquals(mockedResult, response), log::info)
+            );
+    }
+
     private void createMappings() {
         indexMappings = new HashMap<>();
         indexMappings