Add optional config field to tool step

CHANGELOG.md

@@ -19,11 +19,14 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.1.0/)
 
 ## [Unreleased 2.x](https://github.com/opensearch-project/flow-framework/compare/2.17...2.x)
 ### Features
+- Add optional config field to tool step ([#899](https://github.com/opensearch-project/flow-framework/pull/899))
+
 ### Enhancements
 ### Bug Fixes
 ### Infrastructure
 ### Documentation
 - Add query assist data summary agent into sample templates ([#875](https://github.com/opensearch-project/flow-framework/pull/875))
+
 ### Maintenance
 ### Refactoring
 - Update workflow state without using painless script ([#894](https://github.com/opensearch-project/flow-framework/pull/894))

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

@@ -164,6 +164,8 @@ private CommonValue() {}
     public static final String TOOLS_FIELD = "tools";
     /** The tools order field for an agent */
     public static final String TOOLS_ORDER_FIELD = "tools_order";
+    /** The tools config field */
+    public static final String CONFIG_FIELD = "config";
     /** The memory field for an agent */
     public static final String MEMORY_FIELD = "memory";
     /** The app type field for an agent */

src/main/java/org/opensearch/flowframework/workflow/ToolStep.java

@@ -17,10 +17,13 @@
 import org.opensearch.flowframework.util.ParseUtils;
 import org.opensearch.ml.common.agent.MLToolSpec;
 
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 
+import static org.opensearch.flowframework.common.CommonValue.CONFIG_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.INCLUDE_OUTPUT_IN_AGENT_RESPONSE;
 import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD;
@@ -38,7 +41,21 @@ public class ToolStep implements WorkflowStep {
 
     private static final Logger logger = LogManager.getLogger(ToolStep.class);
     PlainActionFuture<WorkflowData> toolFuture = PlainActionFuture.newFuture();
-    static final String NAME = "create_tool";
+
+    /** The name of this step, used as a key in the template and the {@link WorkflowStepFactory} */
+    public static final String NAME = "create_tool";
+    /** Required input keys */
+    public static final Set<String> REQUIRED_INPUTS = Set.of(TYPE);
+    /** Optional input keys */
+    public static final Set<String> OPTIONAL_INPUTS = Set.of(
+        NAME_FIELD,
+        DESCRIPTION_FIELD,
+        PARAMETERS_FIELD,
+        CONFIG_FIELD,
+        INCLUDE_OUTPUT_IN_AGENT_RESPONSE
+    );
+    /** Provided output keys */
+    public static final Set<String> PROVIDED_OUTPUTS = Set.of(TOOLS_FIELD);
 
     @Override
     public PlainActionFuture<WorkflowData> execute(
@@ -48,13 +65,10 @@ public PlainActionFuture<WorkflowData> execute(
         Map<String, String> previousNodeInputs,
         Map<String, String> params
     ) {
-        Set<String> requiredKeys = Set.of(TYPE);
-        Set<String> optionalKeys = Set.of(NAME_FIELD, DESCRIPTION_FIELD, PARAMETERS_FIELD, INCLUDE_OUTPUT_IN_AGENT_RESPONSE);
-
         try {
             Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
-                requiredKeys,
-                optionalKeys,
+                REQUIRED_INPUTS,
+                OPTIONAL_INPUTS,
                 currentNodeInputs,
                 outputs,
                 previousNodeInputs,
@@ -69,11 +83,12 @@ public PlainActionFuture<WorkflowData> execute(
             // parse connector_id, model_id and agent_id from previous node inputs
             Set<String> toolParameterKeys = Set.of(CONNECTOR_ID, MODEL_ID, AGENT_ID);
             Map<String, String> parameters = getToolsParametersMap(
-                inputs.get(PARAMETERS_FIELD),
+                inputs.getOrDefault(PARAMETERS_FIELD, new HashMap<>()),
                 previousNodeInputs,
                 outputs,
                 toolParameterKeys
             );
+            Map<String, String> config = (Map<String, String>) inputs.getOrDefault(CONFIG_FIELD, Collections.emptyMap());
 
             MLToolSpec.MLToolSpecBuilder builder = MLToolSpec.builder();
 
@@ -90,6 +105,8 @@ public PlainActionFuture<WorkflowData> execute(
             if (includeOutputInAgentResponse != null) {
                 builder.includeOutputInAgentResponse(includeOutputInAgentResponse);
             }
+            // TODO https://github.com/opensearch-project/ml-commons/pull/2977/files must be merged for this to compile
+            // builder.configMap(config);
 
             MLToolSpec mlToolSpec = builder.build();
 

src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java

@@ -52,7 +52,6 @@
 import static org.opensearch.flowframework.common.CommonValue.REGISTER_MODEL_STATUS;
 import static org.opensearch.flowframework.common.CommonValue.SOURCE_INDEX;
 import static org.opensearch.flowframework.common.CommonValue.SUCCESS;
-import static org.opensearch.flowframework.common.CommonValue.TOOLS_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.TYPE;
 import static org.opensearch.flowframework.common.CommonValue.URL;
 import static org.opensearch.flowframework.common.CommonValue.VERSION_FIELD;
@@ -231,7 +230,7 @@ public enum WorkflowSteps {
         DELETE_AGENT(DeleteAgentStep.NAME, List.of(AGENT_ID), List.of(AGENT_ID), List.of(OPENSEARCH_ML), null),
 
         /** Create Tool Step */
-        CREATE_TOOL(ToolStep.NAME, List.of(TYPE), List.of(TOOLS_FIELD), List.of(OPENSEARCH_ML), null),
+        CREATE_TOOL(ToolStep.NAME, ToolStep.REQUIRED_INPUTS, ToolStep.PROVIDED_OUTPUTS, List.of(OPENSEARCH_ML), null),
 
         /** Create Ingest Pipeline Step */
         CREATE_INGEST_PIPELINE(

src/test/java/org/opensearch/flowframework/workflow/ToolStepTests.java

@@ -61,6 +61,7 @@ public void setUp() throws Exception {
                 Map.entry("name", "name"),
                 Map.entry("description", "description"),
                 Map.entry("parameters", Collections.emptyMap()),
+                Map.entry("config", Map.of("foo", "bar")),
                 Map.entry("include_output_in_agent_response", "false")
             ),
             "test-id",
@@ -102,6 +103,8 @@ public void testTool() throws ExecutionException, InterruptedException {
         );
         assertTrue(future.isDone());
         assertEquals(MLToolSpec.class, future.get().getContent().get("tools").getClass());
+        // TODO https://github.com/opensearch-project/ml-commons/pull/2977/files must be merged for this to compile
+        // assertEquals(Map.of("foo", "bar"), ((MLToolSpec) future.get().getContent().get("tools")).getConfigMap());
     }
 
     public void testBoolParseFail() {