Refactor Steps

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

@@ -40,6 +40,12 @@ public class DeleteAgentStep implements WorkflowStep {
 
     /** The name of this step, used as a key in the template and the {@link WorkflowStepFactory} */
     public static final String NAME = "delete_agent";
+    /** Required input keys */
+    public static final Set<String> REQUIRED_INPUTS = Set.of(AGENT_ID);
+    /** Optional input keys */
+    public static final Set<String> OPTIONAL_INPUTS = Collections.emptySet();
+    /** Provided output keys */
+    public static final Set<String> PROVIDED_OUTPUTS = Set.of(AGENT_ID);
 
     /**
      * Instantiate this class
@@ -60,13 +66,10 @@ public PlainActionFuture<WorkflowData> execute(
     ) {
         PlainActionFuture<WorkflowData> deleteAgentFuture = PlainActionFuture.newFuture();
 
-        Set<String> requiredKeys = Set.of(AGENT_ID);
-        Set<String> optionalKeys = Collections.emptySet();
-
         try {
             Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
-                requiredKeys,
-                optionalKeys,
+                REQUIRED_INPUTS,
+                OPTIONAL_INPUTS,
                 currentNodeInputs,
                 outputs,
                 previousNodeInputs,

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

@@ -38,6 +38,12 @@ public class DeleteConnectorStep implements WorkflowStep {
 
     /** The name of this step, used as a key in the template and the {@link WorkflowStepFactory} */
     public static final String NAME = "delete_connector";
+    /** Required input keys **/
+    public static final Set<String> REQUIRED_INPUTS = Set.of(CONNECTOR_ID);
+    /** Optional input keys */
+    public static final Set<String> OPTIONAL_INPUTS = Collections.emptySet();
+    /** Provided output keys */
+    public static final Set<String> PROVIDED_OUTPUTS = Set.of(CONNECTOR_ID);
 
     /**
      * Instantiate this class
@@ -58,13 +64,10 @@ public PlainActionFuture<WorkflowData> execute(
     ) {
         PlainActionFuture<WorkflowData> deleteConnectorFuture = PlainActionFuture.newFuture();
 
-        Set<String> requiredKeys = Set.of(CONNECTOR_ID);
-        Set<String> optionalKeys = Collections.emptySet();
-
         try {
             Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
-                requiredKeys,
-                optionalKeys,
+                REQUIRED_INPUTS,
+                OPTIONAL_INPUTS,
                 currentNodeInputs,
                 outputs,
                 previousNodeInputs,

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

@@ -40,6 +40,12 @@ public class DeleteModelStep implements WorkflowStep {
 
     /** The name of this step, used as a key in the template and the {@link WorkflowStepFactory} */
     public static final String NAME = "delete_model";
+    /** Required input keys **/
+    public static final Set<String> REQUIRED_INPUTS = Set.of(MODEL_ID);
+    /** Optional input keys */
+    public static final Set<String> OPTIONAL_INPUTS = Collections.emptySet();
+    /** Provided output keys */
+    public static final Set<String> PROVIDED_OUTPUTS = Set.of(MODEL_ID);
 
     /**
      * Instantiate this class
@@ -60,13 +66,10 @@ public PlainActionFuture<WorkflowData> execute(
     ) {
         PlainActionFuture<WorkflowData> deleteModelFuture = PlainActionFuture.newFuture();
 
-        Set<String> requiredKeys = Set.of(MODEL_ID);
-        Set<String> optionalKeys = Collections.emptySet();
-
         try {
             Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
-                requiredKeys,
-                optionalKeys,
+                REQUIRED_INPUTS,
+                OPTIONAL_INPUTS,
                 currentNodeInputs,
                 outputs,
                 previousNodeInputs,

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

@@ -60,6 +60,7 @@
 import static org.opensearch.flowframework.common.CommonValue.TOOLS_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.TOOLS_ORDER_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.TYPE;
+import static org.opensearch.flowframework.common.WorkflowResources.AGENT_ID;
 import static org.opensearch.flowframework.exception.WorkflowStepException.getSafeException;
 import static org.opensearch.flowframework.util.ParseUtils.getStringToStringMap;
 
@@ -75,7 +76,22 @@ public class RegisterAgentStep implements WorkflowStep {
 
     /** The name of this step, used as a key in the template and the {@link WorkflowStepFactory} */
     public static final String NAME = "register_agent";
-
+    /** Required input keys **/
+    public static final Set<String> REQUIRED_INPUTS = Set.of(NAME_FIELD, TYPE);
+    /** Optional input keys */
+    public static final Set<String> OPTIONAL_INPUTS = Set.of(
+        DESCRIPTION_FIELD,
+        LLM,
+        TOOLS_FIELD,
+        TOOLS_ORDER_FIELD,
+        PARAMETERS_FIELD,
+        MEMORY_FIELD,
+        CREATED_TIME,
+        LAST_UPDATED_TIME_FIELD,
+        APP_TYPE_FIELD
+    );
+    /** Provided output keys */
+    public static final Set<String> PROVIDED_OUTPUTS = Set.of(AGENT_ID);
     /** The model ID for the LLM */
     public static final String MODEL_ID = "model_id";
 
@@ -144,8 +160,8 @@ public void onFailure(Exception ex) {
 
         try {
             Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
-                requiredKeys,
-                optionalKeys,
+                REQUIRED_INPUTS,
+                OPTIONAL_INPUTS,
                 currentNodeInputs,
                 outputs,
                 previousNodeInputs,

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

@@ -43,6 +43,12 @@ public class UndeployModelStep implements WorkflowStep {
 
     /** The name of this step, used as a key in the template and the {@link WorkflowStepFactory} */
     public static final String NAME = "undeploy_model";
+    /** Required input keys **/
+    public static final Set<String> REQUIRED_INPUTS = Set.of(MODEL_ID);
+    /** Optional input keys */
+    public static final Set<String> OPTIONAL_INPUTS = Collections.emptySet();
+    /** Provided output keys */
+    public static final Set<String> PROVIDED_OUTPUTS = Set.of(SUCCESS);
 
     /**
      * Instantiate this class
@@ -63,13 +69,10 @@ public PlainActionFuture<WorkflowData> execute(
     ) {
         PlainActionFuture<WorkflowData> undeployModelFuture = PlainActionFuture.newFuture();
 
-        Set<String> requiredKeys = Set.of(MODEL_ID);
-        Set<String> optionalKeys = Collections.emptySet();
-
         try {
             Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
-                requiredKeys,
-                optionalKeys,
+                REQUIRED_INPUTS,
+                OPTIONAL_INPUTS,
                 currentNodeInputs,
                 outputs,
                 previousNodeInputs,

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

@@ -48,6 +48,12 @@ public class UpdateIndexStep implements WorkflowStep {
 
     /** The name of this step */
     public static final String NAME = "update_index";
+    /** Required input keys */
+    public static final Set<String> REQUIRED_INPUTS = Set.of(INDEX_NAME, CONFIGURATIONS);
+    /** Optional input keys */
+    public static final Set<String> OPTIONAL_INPUTS = Collections.emptySet();
+    /** Provided output keys */
+    public static final Set<String> PROVIDED_OUTPUTS = Set.of(INDEX_NAME);
 
     /**
      * Instantiate this class
@@ -69,14 +75,11 @@ public PlainActionFuture<WorkflowData> execute(
     ) {
         PlainActionFuture<WorkflowData> updateIndexFuture = PlainActionFuture.newFuture();
 
-        Set<String> requiredKeys = Set.of(INDEX_NAME, CONFIGURATIONS);
-        Set<String> optionalKeys = Collections.emptySet();
-
         try {
 
             Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
-                requiredKeys,
-                optionalKeys,
+                REQUIRED_INPUTS,
+                OPTIONAL_INPUTS,
                 currentNodeInputs,
                 outputs,
                 previousNodeInputs,

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

@@ -44,13 +44,8 @@
 import static org.opensearch.flowframework.common.CommonValue.OPENSEARCH_ML;
 import static org.opensearch.flowframework.common.CommonValue.PIPELINE_ID;
 import static org.opensearch.flowframework.common.CommonValue.REGISTER_MODEL_STATUS;
-import static org.opensearch.flowframework.common.CommonValue.SUCCESS;
-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;
-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.INDEX_NAME;
 import static org.opensearch.flowframework.common.WorkflowResources.MODEL_ID;
 
 /**
@@ -219,19 +214,37 @@ public enum WorkflowSteps {
         ),
 
         /** Undeploy Model Step */
-        UNDEPLOY_MODEL(UndeployModelStep.NAME, List.of(MODEL_ID), List.of(SUCCESS), List.of(OPENSEARCH_ML), null),
+        UNDEPLOY_MODEL(
+            UndeployModelStep.NAME,
+            UndeployModelStep.REQUIRED_INPUTS,
+            UndeployModelStep.PROVIDED_OUTPUTS,
+            List.of(OPENSEARCH_ML),
+            null
+        ),
 
         /** Delete Model Step */
-        DELETE_MODEL(DeleteModelStep.NAME, List.of(MODEL_ID), List.of(MODEL_ID), List.of(OPENSEARCH_ML), null),
+        DELETE_MODEL(DeleteModelStep.NAME, DeleteModelStep.REQUIRED_INPUTS, DeleteModelStep.PROVIDED_OUTPUTS, List.of(OPENSEARCH_ML), null),
 
         /** Delete Connector Step */
-        DELETE_CONNECTOR(DeleteConnectorStep.NAME, List.of(CONNECTOR_ID), List.of(CONNECTOR_ID), List.of(OPENSEARCH_ML), null),
+        DELETE_CONNECTOR(
+            DeleteConnectorStep.NAME,
+            DeleteConnectorStep.REQUIRED_INPUTS,
+            DeleteConnectorStep.PROVIDED_OUTPUTS,
+            List.of(OPENSEARCH_ML),
+            null
+        ),
 
         /** Register Agent Step */
-        REGISTER_AGENT(RegisterAgentStep.NAME, List.of(NAME_FIELD, TYPE), List.of(AGENT_ID), List.of(OPENSEARCH_ML), null),
+        REGISTER_AGENT(
+            RegisterAgentStep.NAME,
+            RegisterAgentStep.REQUIRED_INPUTS,
+            RegisterAgentStep.PROVIDED_OUTPUTS,
+            List.of(OPENSEARCH_ML),
+            null
+        ),
 
         /** Delete Agent Step */
-        DELETE_AGENT(DeleteAgentStep.NAME, List.of(AGENT_ID), List.of(AGENT_ID), List.of(OPENSEARCH_ML), null),
+        DELETE_AGENT(DeleteAgentStep.NAME, DeleteAgentStep.REQUIRED_INPUTS, DeleteAgentStep.PROVIDED_OUTPUTS, List.of(OPENSEARCH_ML), null),
 
         /** Create Tool Step */
         CREATE_TOOL(ToolStep.NAME, ToolStep.REQUIRED_INPUTS, ToolStep.PROVIDED_OUTPUTS, List.of(OPENSEARCH_ML), null),
@@ -282,7 +295,13 @@ public enum WorkflowSteps {
         ),
 
         /** Update Index Step */
-        UPDATE_INDEX(UpdateIndexStep.NAME, List.of(INDEX_NAME, CONFIGURATIONS), List.of(INDEX_NAME), Collections.emptyList(), null),
+        UPDATE_INDEX(
+            UpdateIndexStep.NAME,
+            UpdateIndexStep.REQUIRED_INPUTS,
+            UpdateIndexStep.PROVIDED_OUTPUTS,
+            Collections.emptyList(),
+            null
+        ),
 
         /** Delete Search Pipeline Step */
         DELETE_SEARCH_PIPELINE(

src/test/java/org/opensearch/flowframework/model/WorkflowStepValidatorTests.java

@@ -10,11 +10,17 @@
 
 import org.opensearch.flowframework.workflow.CreateConnectorStep;
 import org.opensearch.flowframework.workflow.CreateIndexStep;
+import org.opensearch.flowframework.workflow.DeleteAgentStep;
+import org.opensearch.flowframework.workflow.DeleteConnectorStep;
 import org.opensearch.flowframework.workflow.DeleteIndexStep;
+import org.opensearch.flowframework.workflow.DeleteModelStep;
 import org.opensearch.flowframework.workflow.DeployModelStep;
+import org.opensearch.flowframework.workflow.RegisterAgentStep;
 import org.opensearch.flowframework.workflow.RegisterModelGroupStep;
 import org.opensearch.flowframework.workflow.RegisterRemoteModelStep;
 import org.opensearch.flowframework.workflow.ReindexStep;
+import org.opensearch.flowframework.workflow.UndeployModelStep;
+import org.opensearch.flowframework.workflow.UpdateIndexStep;
 import org.opensearch.flowframework.workflow.WorkflowStepFactory;
 import org.opensearch.test.OpenSearchTestCase;
 
@@ -111,4 +117,59 @@ public void testRegisterModelGroupStepValidator() throws IOException {
         );
     }
 
+    public void testUndeployModelStepValidator() throws IOException {
+        assertStepValidatorMatches(
+            WorkflowStepFactory.WorkflowSteps.UNDEPLOY_MODEL,
+            UndeployModelStep.NAME,
+            UndeployModelStep.REQUIRED_INPUTS,
+            UndeployModelStep.PROVIDED_OUTPUTS
+        );
+    }
+
+    public void testDeleteModelStepValidator() throws IOException {
+        assertStepValidatorMatches(
+            WorkflowStepFactory.WorkflowSteps.DELETE_MODEL,
+            DeleteModelStep.NAME,
+            DeleteModelStep.REQUIRED_INPUTS,
+            DeleteModelStep.PROVIDED_OUTPUTS
+        );
+    }
+
+    public void testDeleteConnectorStepValidator() throws IOException {
+        assertStepValidatorMatches(
+            WorkflowStepFactory.WorkflowSteps.DELETE_CONNECTOR,
+            DeleteConnectorStep.NAME,
+            DeleteConnectorStep.REQUIRED_INPUTS,
+            DeleteConnectorStep.PROVIDED_OUTPUTS
+        );
+
+    }
+
+    public void testRegisterAgentStepValidator() throws IOException {
+        assertStepValidatorMatches(
+            WorkflowStepFactory.WorkflowSteps.REGISTER_AGENT,
+            RegisterAgentStep.NAME,
+            RegisterAgentStep.REQUIRED_INPUTS,
+            RegisterAgentStep.PROVIDED_OUTPUTS
+        );
+    }
+
+    public void testDeleteAgentStepValidator() throws IOException {
+        assertStepValidatorMatches(
+            WorkflowStepFactory.WorkflowSteps.DELETE_AGENT,
+            DeleteAgentStep.NAME,
+            DeleteAgentStep.REQUIRED_INPUTS,
+            DeleteAgentStep.PROVIDED_OUTPUTS
+        );
+    }
+
+    public void testUpdateIndexStepValidator() throws IOException {
+        assertStepValidatorMatches(
+            WorkflowStepFactory.WorkflowSteps.UPDATE_INDEX,
+            UpdateIndexStep.NAME,
+            UpdateIndexStep.REQUIRED_INPUTS,
+            UpdateIndexStep.PROVIDED_OUTPUTS
+        );
+    }
+
 }