Add tool param type validation

nathaliellenaa · nathaliellenaa · commit d416d95ef9d1 · 2025-10-28T11:55:02.000-07:00
Signed-off-by: Nathalie Jonathan &lt;nathhjo@amazon.com&gt;
diff --git a/common/src/main/java/org/opensearch/ml/common/input/execute/tool/ToolMLInput.java b/common/src/main/java/org/opensearch/ml/common/input/execute/tool/ToolMLInput.java
@@ -33,6 +33,9 @@ public class ToolMLInput extends MLInput {
     @Setter
     private String toolName;
 
+    @Getter
+    private Map<String, Object> originalParameters;
+
     public ToolMLInput(StreamInput in) throws IOException {
         super(in);
         this.toolName = in.readString();
@@ -66,7 +69,9 @@ public ToolMLInput(XContentParser parser, FunctionName functionName) throws IOEx
                     toolName = parser.text();
                     break;
                 case PARAMETERS_FIELD:
-                    Map<String, String> parameters = StringUtils.getParameterMap(parser.map());
+                    Map<String, Object> rawParams = parser.map();
+                    originalParameters = rawParams;
+                    Map<String, String> parameters = StringUtils.getParameterMap(rawParams);
                     inputDataset = new RemoteInferenceInputDataSet(parameters);
                     break;
                 default:
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/tool/MLToolExecutor.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/tool/MLToolExecutor.java
@@ -85,6 +85,14 @@ public void execute(Input input, ActionListener<Output> listener) {
         try {
             Map<String, String> mutableParams = new HashMap<>(parameters);
             Tool tool = toolFactory.create(mutableParams);
+
+            // Validate original parameter types
+            Map<String, Object> originalParameters = toolMLInput.getOriginalParameters();
+            if (originalParameters != null && !tool.validateParameterTypes(originalParameters)) {
+                listener.onFailure(new IllegalArgumentException("Invalid parameter types for tool: " + toolName));
+                return;
+            }
+
             if (!tool.validate(mutableParams)) {
                 listener.onFailure(new IllegalArgumentException("Invalid parameters for tool: " + toolName));
                 return;
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/AgentTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/AgentTool.java
@@ -117,6 +117,28 @@ public void setName(String s) {
 
     @Override
     public boolean validate(Map<String, String> parameters) {
+        if (parameters == null || parameters.isEmpty()) {
+            return false;
+        }
+
+        // Validate question length
+        String question = parameters.get("question");
+        if (question != null && question.length() > 10000) {
+            throw new IllegalArgumentException("question length cannot exceed 10000 characters");
+        }
+
+        return true;
+    }
+
+    @Override
+    public boolean validateParameterTypes(Map<String, Object> parameters) {
+        // Validate question must be String
+        Object questionObj = parameters.get("question");
+        if (questionObj != null && !(questionObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("question must be a String type, but got %s", questionObj.getClass().getSimpleName())
+            );
+        }
         return true;
     }
 
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/ConnectorTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/ConnectorTool.java
@@ -100,6 +100,18 @@ public boolean validate(Map<String, String> parameters) {
         return parameters != null && !parameters.isEmpty();
     }
 
+    @Override
+    public boolean validateParameterTypes(Map<String, Object> parameters) {
+        // Validate response_filter must be String
+        Object responseFilterObj = parameters.get("response_filter");
+        if (responseFilterObj != null && !(responseFilterObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("response_filter must be a String type, but got %s", responseFilterObj.getClass().getSimpleName())
+            );
+        }
+        return true;
+    }
+
     public static class Factory implements Tool.Factory<ConnectorTool> {
         public static final String TYPE = "ConnectorTool";
         public static final String DEFAULT_DESCRIPTION = "Invokes external service. Required: 'connector_id'. Returns: service response.";
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexInsightTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexInsightTool.java
@@ -113,6 +113,11 @@ public boolean validate(Map<String, String> parameters) {
         return true;
     }
 
+    @Override
+    public boolean validateParameterTypes(Map<String, Object> parameters) {
+        return true;
+    }
+
     public static class Factory implements Tool.Factory<IndexInsightTool> {
         private Client client;
 
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java
@@ -175,7 +175,45 @@ public String getType() {
 
     @Override
     public boolean validate(Map<String, String> parameters) {
-        return parameters != null && !parameters.isEmpty() && parameters.containsKey("index");
+        if (parameters == null || parameters.isEmpty() || !parameters.containsKey("index")) {
+            return false;
+        }
+
+        // Validate question length
+        String question = parameters.get("question");
+        if (question != null && question.length() > 10000) {
+            throw new IllegalArgumentException("question length cannot exceed 10000 characters");
+        }
+
+        return true;
+    }
+
+    @Override
+    public boolean validateParameterTypes(Map<String, Object> parameters) {
+        // Validate question must be String
+        Object questionObj = parameters.get("question");
+        if (questionObj != null && !(questionObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("question must be a String type, but got %s", questionObj.getClass().getSimpleName())
+            );
+        }
+
+        // Validate index must be ArrayList
+        Object indexObj = parameters.get("index");
+        if (indexObj != null && !(indexObj instanceof ArrayList)) {
+            throw new IllegalArgumentException(
+                String.format("index must be an Array type, but got %s", indexObj.getClass().getSimpleName())
+            );
+        }
+
+        // Validate local must be Boolean
+        Object localObj = parameters.get("local");
+        if (localObj != null && !(localObj instanceof Boolean)) {
+            throw new IllegalArgumentException(
+                String.format("local must be a Boolean type, but got %s", localObj.getClass().getSimpleName())
+            );
+        }
+        return true;
     }
 
     /**
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/ListIndexTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/ListIndexTool.java
@@ -415,7 +415,53 @@ public void onFailure(final Exception e) {
 
     @Override
     public boolean validate(Map<String, String> parameters) {
-        return parameters != null && !parameters.isEmpty();
+        if (parameters == null || parameters.isEmpty()) {
+            return false;
+        }
+
+        // Validate question length
+        String question = parameters.get("question");
+        if (question != null && question.length() > 10000) {
+            throw new IllegalArgumentException("question length cannot exceed 10000 characters");
+        }
+
+        return true;
+    }
+
+    @Override
+    public boolean validateParameterTypes(Map<String, Object> parameters) {
+        // Validate question must be String
+        Object questionObj = parameters.get("question");
+        if (questionObj != null && !(questionObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("question must be a String type, but got %s", questionObj.getClass().getSimpleName())
+            );
+        }
+
+        // Validate indices must be ArrayList
+        Object indicesObj = parameters.get("indices");
+        if (indicesObj != null && !(indicesObj instanceof ArrayList)) {
+            throw new IllegalArgumentException(
+                String.format("indices must be an Array type, but got %s", indicesObj.getClass().getSimpleName())
+            );
+        }
+
+        // Validate local must be Boolean
+        Object localObj = parameters.get("local");
+        if (localObj != null && !(localObj instanceof Boolean)) {
+            throw new IllegalArgumentException(
+                String.format("local must be a Boolean type, but got %s", localObj.getClass().getSimpleName())
+            );
+        }
+
+        // Validate page_size must be Integer
+        Object pageSizeObj = parameters.get("page_size");
+        if (pageSizeObj != null && !(pageSizeObj instanceof Integer)) {
+            throw new IllegalArgumentException(
+                String.format("page_size must be an Integer type, but got %s", pageSizeObj.getClass().getSimpleName())
+            );
+        }
+        return true;
     }
 
     /**
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/MLModelTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/MLModelTool.java
@@ -148,6 +148,26 @@ public boolean validate(Map<String, String> parameters) {
         return parameters != null && !parameters.isEmpty();
     }
 
+    @Override
+    public boolean validateParameterTypes(Map<String, Object> parameters) {
+        // Validate prompt must be String
+        Object promptObj = parameters.get("prompt");
+        if (promptObj != null && !(promptObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("prompt must be a String type, but got %s", promptObj.getClass().getSimpleName())
+            );
+        }
+
+        // Validate response_field must be String
+        Object responseFieldObj = parameters.get(RESPONSE_FIELD);
+        if (responseFieldObj != null && !(responseFieldObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("%s must be a String type, but got %s", RESPONSE_FIELD, responseFieldObj.getClass().getSimpleName())
+            );
+        }
+        return true;
+    }
+
     public static class Factory implements WithModelTool.Factory<MLModelTool> {
         private Client client;
 
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/McpSseTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/McpSseTool.java
@@ -99,6 +99,11 @@ public boolean validate(Map<String, String> parameters) {
         return true;
     }
 
+    @Override
+    public boolean validateParameterTypes(Map<String, Object> parameters) {
+        return true;
+    }
+
     public static class Factory implements WithModelTool.Factory<McpSseTool> {
         private static Factory INSTANCE;
 
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/McpStreamableHttpTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/McpStreamableHttpTool.java
@@ -100,6 +100,11 @@ public boolean validate(Map<String, String> parameters) {
         return true;
     }
 
+    @Override
+    public boolean validateParameterTypes(Map<String, Object> parameters) {
+        return true;
+    }
+
     public static class Factory implements WithModelTool.Factory<McpStreamableHttpTool> {
         private static Factory INSTANCE;
 
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningTool.java
@@ -18,6 +18,7 @@
 import static org.opensearch.ml.engine.tools.QueryPlanningPromptTemplate.DEFAULT_TEMPLATE_SELECTION_USER_PROMPT;
 
 import java.io.IOException;
+import java.lang.reflect.Array;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -385,6 +386,91 @@ public boolean validate(Map<String, String> parameters) {
             || !parameters.containsKey(INDEX_NAME_FIELD)) {
             return false;
         }
+
+        // Validate question length
+        String question = parameters.get(QUESTION_FIELD);
+        if (question != null && question.length() > 10000) {
+            throw new IllegalArgumentException("question length cannot exceed 10000 characters");
+        }
+        return true;
+    }
+
+    @Override
+    public boolean validateParameterTypes(Map<String, Object> parameters) {
+        // Validate response_filter must be String
+        Object responseFilterObj = parameters.get("response_filter");
+        if (responseFilterObj != null && !(responseFilterObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("response_filter must be a String type, but got %s", responseFilterObj.getClass().getSimpleName())
+            );
+        }
+
+        // Validate generation_type must be String
+        Object generationTypeObj = parameters.get(GENERATION_TYPE_FIELD);
+        if (generationTypeObj != null && !(generationTypeObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("%s must be a String type, but got %s", GENERATION_TYPE_FIELD, generationTypeObj.getClass().getSimpleName())
+            );
+        }
+
+        // Validate query_planner_system_prompt must be String
+        Object queryPlannerSystemPromptObj = parameters.get(QUERY_PLANNER_SYSTEM_PROMPT_FIELD);
+        if (queryPlannerSystemPromptObj != null && !(queryPlannerSystemPromptObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String
+                    .format(
+                        "%s must be a String type, but got %s",
+                        QUERY_PLANNER_SYSTEM_PROMPT_FIELD,
+                        queryPlannerSystemPromptObj.getClass().getSimpleName()
+                    )
+            );
+        }
+
+        // Validate query_planner_user_prompt must be String
+        Object queryPlannerUserPromptObj = parameters.get(QUERY_PLANNER_USER_PROMPT_FIELD);
+        if (queryPlannerUserPromptObj != null && !(queryPlannerUserPromptObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String
+                    .format(
+                        "%s must be a String type, but got %s",
+                        QUERY_PLANNER_USER_PROMPT_FIELD,
+                        queryPlannerUserPromptObj.getClass().getSimpleName()
+                    )
+            );
+        }
+
+        // Validate search_templates must be Array
+        Object searchTemplatesObj = parameters.get(SEARCH_TEMPLATES_FIELD);
+        if (searchTemplatesObj != null && !(searchTemplatesObj instanceof Array)) {
+            throw new IllegalArgumentException(
+                String.format("%s must be an Array type, but got %s", SEARCH_TEMPLATES_FIELD, searchTemplatesObj.getClass().getSimpleName())
+            );
+        }
+
+        // Validate question must be String
+        Object questionObj = parameters.get(QUESTION_FIELD);
+        if (questionObj != null && !(questionObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("%s must be a String type, but got %s", QUESTION_FIELD, questionObj.getClass().getSimpleName())
+            );
+        }
+
+        // Validate index_name must be String
+        Object indexNameObj = parameters.get(INDEX_NAME_FIELD);
+        if (indexNameObj != null && !(indexNameObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("%s must be a String type, but got %s", INDEX_NAME_FIELD, indexNameObj.getClass().getSimpleName())
+            );
+        }
+
+        // Validate embedding_model_id must be String
+        Object embeddingModelIdObj = parameters.get("embedding_model_id");
+        if (embeddingModelIdObj != null && !(embeddingModelIdObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("embedding_model_id must be a String type, but got %s", embeddingModelIdObj.getClass().getSimpleName())
+            );
+        }
+
         return true;
     }
 
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/ReadFromScratchPadTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/ReadFromScratchPadTool.java
@@ -84,6 +84,18 @@ public boolean validate(Map<String, String> parameters) {
         return true;
     }
 
+    @Override
+    public boolean validateParameterTypes(Map<String, Object> parameters) {
+        // Validate persistent_notes must be String
+        Object persistentNotesObj = parameters.get(PERSISTENT_NOTES_KEY);
+        if (persistentNotesObj != null && !(persistentNotesObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("%s must be a String type, but got %s", PERSISTENT_NOTES_KEY, persistentNotesObj.getClass().getSimpleName())
+            );
+        }
+        return true;
+    }
+
     /**
      * Executes the ReadFromScratchPadTool.
      * This tool retrieves notes from the persistent scratchpad for the current conversation.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/SearchIndexTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/SearchIndexTool.java
@@ -131,6 +131,18 @@ public boolean validate(Map<String, String> parameters) {
         return true;
     }
 
+    @Override
+    public boolean validateParameterTypes(Map<String, Object> parameters) {
+        // Validate input must be String
+        Object inputObj = parameters.get(INPUT_FIELD);
+        if (inputObj != null && !(inputObj instanceof String)) {
+            throw new IllegalArgumentException(
+                String.format("%s must be a String type, but got %s", INPUT_FIELD, inputObj.getClass().getSimpleName())
+            );
+        }
+        return true;
+    }
+
     private SearchRequest getSearchRequest(String index, String query) throws IOException {
         SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
         XContentParser queryParser = XContentType.JSON.xContent().createParser(xContentRegistry, LoggingDeprecationHandler.INSTANCE, query);
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/VisualizationsTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/VisualizationsTool.java
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/WriteToScratchPadTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/WriteToScratchPadTool.java
diff --git a/plugin/src/test/java/org/opensearch/ml/plugin/DummyWrongTool.java b/plugin/src/test/java/org/opensearch/ml/plugin/DummyWrongTool.java
diff --git a/spi/src/main/java/org/opensearch/ml/common/spi/tools/Tool.java b/spi/src/main/java/org/opensearch/ml/common/spi/tools/Tool.java
