Add new attributes field to ToolStep

CHANGELOG.md

@@ -7,6 +7,8 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.1.0/)
 ### Features
 ### Enhancements
 ### Bug Fixes
+- Add new attributes field to ToolStep ([#1113](https://github.com/opensearch-project/flow-framework/pull/1113))
+
 ### Infrastructure
 ### Documentation
 ### Maintenance

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

@@ -172,14 +172,10 @@ 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 */
     public static final String APP_TYPE_FIELD = "app_type";
-    /** To include field for an agent response */
-    public static final String INCLUDE_OUTPUT_IN_AGENT_RESPONSE = "include_output_in_agent_response";
     /** Pipeline ID, also corresponds to pipeline name */
     public static final String PIPELINE_ID = "pipeline_id";
     /** Pipeline Configurations */

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

@@ -23,16 +23,17 @@
 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;
-import static org.opensearch.flowframework.common.CommonValue.PARAMETERS_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.TOOLS_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.TYPE;
 import static org.opensearch.flowframework.common.WorkflowResources.AGENT_ID;
 import static org.opensearch.flowframework.common.WorkflowResources.CONNECTOR_ID;
 import static org.opensearch.flowframework.common.WorkflowResources.MODEL_ID;
+import static org.opensearch.ml.common.agent.MLToolSpec.ATTRIBUTES_FIELD;
+import static org.opensearch.ml.common.agent.MLToolSpec.CONFIG_FIELD;
+import static org.opensearch.ml.common.agent.MLToolSpec.DESCRIPTION_FIELD;
+import static org.opensearch.ml.common.agent.MLToolSpec.INCLUDE_OUTPUT_IN_AGENT_RESPONSE;
+import static org.opensearch.ml.common.agent.MLToolSpec.PARAMETERS_FIELD;
+import static org.opensearch.ml.common.agent.MLToolSpec.TOOL_NAME_FIELD;
+import static org.opensearch.ml.common.agent.MLToolSpec.TOOL_TYPE_FIELD;
 
 /**
  * Step to register a tool for an agent
@@ -45,12 +46,13 @@ public class ToolStep implements WorkflowStep {
     /** 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);
+    public static final Set<String> REQUIRED_INPUTS = Set.of(TOOL_TYPE_FIELD);
     /** Optional input keys */
     public static final Set<String> OPTIONAL_INPUTS = Set.of(
-        NAME_FIELD,
+        TOOL_NAME_FIELD,
         DESCRIPTION_FIELD,
         PARAMETERS_FIELD,
+        ATTRIBUTES_FIELD,
         CONFIG_FIELD,
         INCLUDE_OUTPUT_IN_AGENT_RESPONSE
     );
@@ -76,8 +78,8 @@ public PlainActionFuture<WorkflowData> execute(
                 params
             );
 
-            String type = (String) inputs.get(TYPE);
-            String name = (String) inputs.get(NAME_FIELD);
+            String type = (String) inputs.get(TOOL_TYPE_FIELD);
+            String name = (String) inputs.get(TOOL_NAME_FIELD);
             String description = (String) inputs.get(DESCRIPTION_FIELD);
             Boolean includeOutputInAgentResponse = ParseUtils.parseIfExists(inputs, INCLUDE_OUTPUT_IN_AGENT_RESPONSE, Boolean.class);
 
@@ -89,6 +91,7 @@ public PlainActionFuture<WorkflowData> execute(
                 outputs,
                 toolParameterKeys
             );
+            Map<String, String> attributes = (Map<String, String>) inputs.get(ATTRIBUTES_FIELD);
             @SuppressWarnings("unchecked")
             Map<String, String> config = (Map<String, String>) inputs.getOrDefault(CONFIG_FIELD, Collections.emptyMap());
 
@@ -104,6 +107,9 @@ public PlainActionFuture<WorkflowData> execute(
             if (parameters != null) {
                 builder.parameters(parameters);
             }
+            if (attributes != null) {
+                builder.attributes(attributes);
+            }
             if (includeOutputInAgentResponse != null) {
                 builder.includeOutputInAgentResponse(includeOutputInAgentResponse);
             }

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

@@ -24,6 +24,7 @@
 
 public class ToolStepTests extends OpenSearchTestCase {
     private WorkflowData inputData;
+    private WorkflowData inputDataWithAttributes;
     private WorkflowData inputDataWithConnectorId;
     private WorkflowData inputDataWithModelId;
     private WorkflowData inputDataWithAgentId;
@@ -52,6 +53,18 @@ public void setUp() throws Exception {
             "test-id",
             "test-node-id"
         );
+        inputDataWithAttributes = new WorkflowData(
+            Map.ofEntries(
+                Map.entry("type", "type"),
+                Map.entry("name", "name"),
+                Map.entry("description", "description"),
+                Map.entry("parameters", Collections.emptyMap()),
+                Map.entry("attributes", Map.of("key1", "value1", "key2", "value2")),
+                Map.entry("include_output_in_agent_response", false)
+            ),
+            "test-id",
+            "test-node-id"
+        );
         inputDataWithConnectorId = new WorkflowData(Map.of(CONNECTOR_ID, mockedConnectorId), "test-id", createConnectorNodeId);
         inputDataWithModelId = new WorkflowData(Map.of(MODEL_ID, mockedModelId), "test-id", createModelNodeId);
         inputDataWithAgentId = new WorkflowData(Map.of(AGENT_ID, mockedAgentId), "test-id", createAgentNodeId);
@@ -128,6 +141,24 @@ public void testBoolParseFail() {
         assertEquals(RestStatus.BAD_REQUEST, w.getRestStatus());
     }
 
+    public void testToolWithAttributes() throws ExecutionException, InterruptedException {
+        ToolStep toolStep = new ToolStep();
+
+        PlainActionFuture<WorkflowData> future = toolStep.execute(
+            inputDataWithAttributes.getNodeId(),
+            inputDataWithAttributes,
+            Collections.emptyMap(),
+            Collections.emptyMap(),
+            Collections.emptyMap(),
+            null
+        );
+        assertTrue(future.isDone());
+        Object tools = future.get().getContent().get("tools");
+        assertEquals(MLToolSpec.class, tools.getClass());
+        MLToolSpec mlToolSpec = (MLToolSpec) tools;
+        assertEquals(Map.of("key1", "value1", "key2", "value2"), mlToolSpec.getAttributes());
+    }
+
     public void testToolWithConnectorId() throws ExecutionException, InterruptedException {
         ToolStep toolStep = new ToolStep();
 

src/test/resources/template/agent-framework.json

@@ -85,6 +85,9 @@
             "description": "A general tool to calculate any math problem. The action input must be valid math expression, like 2+3",
             "parameters": {
               "max_iteration": 5
+            },
+            "attributes": {
+                "test_attribute": "test_value"
             }
           }
         },