Refactor DeployModel and RegisterRemoteModel steps to use const (#1306)

Refactor to const DeployModel & RegisterRemoteModel Steps

Signed-off-by: kokibas <[email protected]>

Author: kokibas

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

@@ -41,6 +41,12 @@ public class DeployModelStep extends AbstractRetryableWorkflowStep {
 
     /** The name of this step, used as a key in the template and the {@link WorkflowStepFactory} */
     public static final String NAME = "deploy_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
@@ -71,13 +77,10 @@ public PlainActionFuture<WorkflowData> execute(
 
         PlainActionFuture<WorkflowData> deployModelFuture = 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/RegisterRemoteModelStep.java

@@ -41,6 +41,7 @@
 import static org.opensearch.flowframework.common.CommonValue.REGISTER_MODEL_STATUS;
 import static org.opensearch.flowframework.common.WorkflowResources.CONNECTOR_ID;
 import static org.opensearch.flowframework.common.WorkflowResources.MODEL_GROUP_ID;
+import static org.opensearch.flowframework.common.WorkflowResources.MODEL_ID;
 import static org.opensearch.flowframework.common.WorkflowResources.getResourceByWorkflowStep;
 import static org.opensearch.flowframework.exception.WorkflowStepException.getSafeException;
 
@@ -57,6 +58,18 @@ public class RegisterRemoteModelStep implements WorkflowStep {
 
     /** The name of this step, used as a key in the template and the {@link WorkflowStepFactory} */
     public static final String NAME = "register_remote_model";
+    /** Required input keys **/
+    public static final Set<String> REQUIRED_INPUTS = Set.of(NAME_FIELD, CONNECTOR_ID);
+    /** Optional input keys */
+    public static final Set<String> OPTIONAL_INPUTS = Set.of(
+        MODEL_GROUP_ID,
+        DESCRIPTION_FIELD,
+        DEPLOY_FIELD,
+        GUARDRAILS_FIELD,
+        INTERFACE_FIELD
+    );
+    /** Provided output keys */
+    public static final Set<String> PROVIDED_OUTPUTS = Set.of(MODEL_ID, REGISTER_MODEL_STATUS);
 
     /**
      * Instantiate this class
@@ -80,13 +93,10 @@ public PlainActionFuture<WorkflowData> execute(
 
         PlainActionFuture<WorkflowData> registerRemoteModelFuture = PlainActionFuture.newFuture();
 
-        Set<String> requiredKeys = Set.of(NAME_FIELD, CONNECTOR_ID);
-        Set<String> optionalKeys = Set.of(MODEL_GROUP_ID, DESCRIPTION_FIELD, DEPLOY_FIELD, GUARDRAILS_FIELD, INTERFACE_FIELD);
-
         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

@@ -198,8 +198,8 @@ public enum WorkflowSteps {
         /** Register Remote Model Step */
         REGISTER_REMOTE_MODEL(
             RegisterRemoteModelStep.NAME,
-            List.of(NAME_FIELD, CONNECTOR_ID),
-            List.of(MODEL_ID, REGISTER_MODEL_STATUS),
+            RegisterRemoteModelStep.REQUIRED_INPUTS,
+            RegisterRemoteModelStep.PROVIDED_OUTPUTS,
             List.of(OPENSEARCH_ML),
             null
         ),
@@ -214,7 +214,13 @@ public enum WorkflowSteps {
         ),
 
         /** Deploy Model Step */
-        DEPLOY_MODEL(DeployModelStep.NAME, List.of(MODEL_ID), List.of(MODEL_ID), List.of(OPENSEARCH_ML), TimeValue.timeValueSeconds(15)),
+        DEPLOY_MODEL(
+            DeployModelStep.NAME,
+            DeployModelStep.REQUIRED_INPUTS,
+            DeployModelStep.PROVIDED_OUTPUTS,
+            List.of(OPENSEARCH_ML),
+            TimeValue.timeValueSeconds(15)
+        ),
 
         /** Undeploy Model Step */
         UNDEPLOY_MODEL(UndeployModelStep.NAME, List.of(MODEL_ID), List.of(SUCCESS), List.of(OPENSEARCH_ML), null),

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

@@ -11,6 +11,8 @@
 import org.opensearch.flowframework.workflow.CreateConnectorStep;
 import org.opensearch.flowframework.workflow.CreateIndexStep;
 import org.opensearch.flowframework.workflow.DeleteIndexStep;
+import org.opensearch.flowframework.workflow.DeployModelStep;
+import org.opensearch.flowframework.workflow.RegisterRemoteModelStep;
 import org.opensearch.flowframework.workflow.WorkflowStepFactory;
 import org.opensearch.test.OpenSearchTestCase;
 
@@ -71,4 +73,22 @@ public void testCreateIndexStepValidator() throws IOException {
         );
     }
 
+    public void testDeployModelStepValidator() throws IOException {
+        assertStepValidatorMatches(
+            WorkflowStepFactory.WorkflowSteps.DEPLOY_MODEL,
+            DeployModelStep.NAME,
+            DeployModelStep.REQUIRED_INPUTS,
+            DeployModelStep.PROVIDED_OUTPUTS
+        );
+    }
+
+    public void testRegisterRemoteModelStepValidator() throws IOException {
+        assertStepValidatorMatches(
+            WorkflowStepFactory.WorkflowSteps.REGISTER_REMOTE_MODEL,
+            RegisterRemoteModelStep.NAME,
+            RegisterRemoteModelStep.REQUIRED_INPUTS,
+            RegisterRemoteModelStep.PROVIDED_OUTPUTS
+        );
+    }
+
 }