diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeleteAgentStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeleteAgentStep.java index 2d54d6498..caaf4efb2 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/DeleteAgentStep.java +++ b/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 REQUIRED_INPUTS = Set.of(AGENT_ID); + /** Optional input keys */ + public static final Set OPTIONAL_INPUTS = Collections.emptySet(); + /** Provided output keys */ + public static final Set PROVIDED_OUTPUTS = Set.of(AGENT_ID); /** * Instantiate this class @@ -60,13 +66,10 @@ public PlainActionFuture execute( ) { PlainActionFuture deleteAgentFuture = PlainActionFuture.newFuture(); - Set requiredKeys = Set.of(AGENT_ID); - Set optionalKeys = Collections.emptySet(); - try { Map inputs = ParseUtils.getInputsFromPreviousSteps( - requiredKeys, - optionalKeys, + REQUIRED_INPUTS, + OPTIONAL_INPUTS, currentNodeInputs, outputs, previousNodeInputs, diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeleteConnectorStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeleteConnectorStep.java index 2ddcb75b8..ee4221d1f 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/DeleteConnectorStep.java +++ b/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 REQUIRED_INPUTS = Set.of(CONNECTOR_ID); + /** Optional input keys */ + public static final Set OPTIONAL_INPUTS = Collections.emptySet(); + /** Provided output keys */ + public static final Set PROVIDED_OUTPUTS = Set.of(CONNECTOR_ID); /** * Instantiate this class @@ -58,13 +64,10 @@ public PlainActionFuture execute( ) { PlainActionFuture deleteConnectorFuture = PlainActionFuture.newFuture(); - Set requiredKeys = Set.of(CONNECTOR_ID); - Set optionalKeys = Collections.emptySet(); - try { Map inputs = ParseUtils.getInputsFromPreviousSteps( - requiredKeys, - optionalKeys, + REQUIRED_INPUTS, + OPTIONAL_INPUTS, currentNodeInputs, outputs, previousNodeInputs, diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeleteModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeleteModelStep.java index 207ea1c76..2560c1033 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/DeleteModelStep.java +++ b/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 REQUIRED_INPUTS = Set.of(MODEL_ID); + /** Optional input keys */ + public static final Set OPTIONAL_INPUTS = Collections.emptySet(); + /** Provided output keys */ + public static final Set PROVIDED_OUTPUTS = Set.of(MODEL_ID); /** * Instantiate this class @@ -60,13 +66,10 @@ public PlainActionFuture execute( ) { PlainActionFuture deleteModelFuture = PlainActionFuture.newFuture(); - Set requiredKeys = Set.of(MODEL_ID); - Set optionalKeys = Collections.emptySet(); - try { Map inputs = ParseUtils.getInputsFromPreviousSteps( - requiredKeys, - optionalKeys, + REQUIRED_INPUTS, + OPTIONAL_INPUTS, currentNodeInputs, outputs, previousNodeInputs, diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java index 599b370aa..65a8560be 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java +++ b/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java @@ -55,6 +55,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; @@ -70,7 +71,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 REQUIRED_INPUTS = Set.of(NAME_FIELD, TYPE); + /** Optional input keys */ + public static final Set 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 PROVIDED_OUTPUTS = Set.of(AGENT_ID); /** The model ID for the LLM */ public static final String MODEL_ID = "model_id"; @@ -121,23 +137,10 @@ public void onFailure(Exception ex) { } }; - Set requiredKeys = Set.of(NAME_FIELD, TYPE); - Set optionalKeys = Set.of( - DESCRIPTION_FIELD, - LLM, - TOOLS_FIELD, - TOOLS_ORDER_FIELD, - PARAMETERS_FIELD, - MEMORY_FIELD, - CREATED_TIME, - LAST_UPDATED_TIME_FIELD, - APP_TYPE_FIELD - ); - try { Map inputs = ParseUtils.getInputsFromPreviousSteps( - requiredKeys, - optionalKeys, + REQUIRED_INPUTS, + OPTIONAL_INPUTS, currentNodeInputs, outputs, previousNodeInputs, diff --git a/src/main/java/org/opensearch/flowframework/workflow/UndeployModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/UndeployModelStep.java index a4764576f..185f69ca1 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/UndeployModelStep.java +++ b/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 REQUIRED_INPUTS = Set.of(MODEL_ID); + /** Optional input keys */ + public static final Set OPTIONAL_INPUTS = Collections.emptySet(); + /** Provided output keys */ + public static final Set PROVIDED_OUTPUTS = Set.of(SUCCESS); /** * Instantiate this class @@ -63,13 +69,10 @@ public PlainActionFuture execute( ) { PlainActionFuture undeployModelFuture = PlainActionFuture.newFuture(); - Set requiredKeys = Set.of(MODEL_ID); - Set optionalKeys = Collections.emptySet(); - try { Map inputs = ParseUtils.getInputsFromPreviousSteps( - requiredKeys, - optionalKeys, + REQUIRED_INPUTS, + OPTIONAL_INPUTS, currentNodeInputs, outputs, previousNodeInputs, diff --git a/src/main/java/org/opensearch/flowframework/workflow/UpdateIndexStep.java b/src/main/java/org/opensearch/flowframework/workflow/UpdateIndexStep.java index b527f4c1d..8f21bf359 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/UpdateIndexStep.java +++ b/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 REQUIRED_INPUTS = Set.of(INDEX_NAME, CONFIGURATIONS); + /** Optional input keys */ + public static final Set OPTIONAL_INPUTS = Collections.emptySet(); + /** Provided output keys */ + public static final Set PROVIDED_OUTPUTS = Set.of(INDEX_NAME); /** * Instantiate this class @@ -69,14 +75,11 @@ public PlainActionFuture execute( ) { PlainActionFuture updateIndexFuture = PlainActionFuture.newFuture(); - Set requiredKeys = Set.of(INDEX_NAME, CONFIGURATIONS); - Set optionalKeys = Collections.emptySet(); - try { Map inputs = ParseUtils.getInputsFromPreviousSteps( - requiredKeys, - optionalKeys, + REQUIRED_INPUTS, + OPTIONAL_INPUTS, currentNodeInputs, outputs, previousNodeInputs, diff --git a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java b/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java index 1978aefff..fce015d73 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java +++ b/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( diff --git a/src/test/java/org/opensearch/flowframework/model/WorkflowStepValidatorTests.java b/src/test/java/org/opensearch/flowframework/model/WorkflowStepValidatorTests.java index f535dedb8..22cc87a9b 100644 --- a/src/test/java/org/opensearch/flowframework/model/WorkflowStepValidatorTests.java +++ b/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 + ); + } + }