Allow strings for boolean workflow step parameters (#671)

* Allow strings for boolean workflow step parameters Signed-off-by: Daniel Widdis <[email protected]> * Reduce redundancy in similar tests Signed-off-by: Daniel Widdis <[email protected]> --------- Signed-off-by: Daniel Widdis <[email protected]> (cherry picked from commit b8e4ed0) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Apr 22, 2024 · 45f123c · 45f123c
1 parent fcef4e0
commit 45f123c
Show file tree

Hide file tree

Showing 11 changed files with 402 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,6 +16,8 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.1.0/)
 ### Features
 ### Enhancements
 - Add guardrails to default use case params ([#658](https://github.com/opensearch-project/flow-framework/pull/658))
+- Allow strings for boolean workflow step parameters ([#671](https://github.com/opensearch-project/flow-framework/pull/671))
+
 ### Bug Fixes
 - Reset workflow state to initial state after successful deprovision ([#635](https://github.com/opensearch-project/flow-framework/pull/635))
 - Silently ignore content on APIs that don't require it ([#639](https://github.com/opensearch-project/flow-framework/pull/639))

diff --git a/src/main/java/org/opensearch/flowframework/workflow/AbstractRegisterLocalModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/AbstractRegisterLocalModelStep.java
@@ -12,7 +12,9 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.support.PlainActionFuture;
+import org.opensearch.common.Booleans;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.common.FlowFrameworkSettings;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.exception.WorkflowStepException;
@@ -113,7 +115,7 @@ public PlainActionFuture<WorkflowData> execute(
             String description = (String) inputs.get(DESCRIPTION_FIELD);
             String modelGroupId = (String) inputs.get(MODEL_GROUP_ID);
             String allConfig = (String) inputs.get(ALL_CONFIG);
-            final Boolean deploy = (Boolean) inputs.get(DEPLOY_FIELD);
+            final Boolean deploy = inputs.containsKey(DEPLOY_FIELD) ? Booleans.parseBoolean(inputs.get(DEPLOY_FIELD).toString()) : null;
 
             // Build register model input
             MLRegisterModelInputBuilder mlInputBuilder = MLRegisterModelInput.builder()
@@ -217,6 +219,8 @@ public PlainActionFuture<WorkflowData> execute(
                 logger.error(errorMessage, exception);
                 registerLocalModelFuture.onFailure(new WorkflowStepException(errorMessage, ExceptionsHelper.status(exception)));
             }));
+        } catch (IllegalArgumentException iae) {
+            registerLocalModelFuture.onFailure(new WorkflowStepException(iae.getMessage(), RestStatus.BAD_REQUEST));
         } catch (FlowFrameworkException e) {
             registerLocalModelFuture.onFailure(e);
         }

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java
@@ -12,8 +12,10 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.support.PlainActionFuture;
+import org.opensearch.common.Booleans;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.util.CollectionUtils;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.exception.WorkflowStepException;
 import org.opensearch.flowframework.indices.FlowFrameworkIndicesHandler;
@@ -139,7 +141,9 @@ public void onFailure(Exception e) {
             String description = (String) inputs.get(DESCRIPTION_FIELD);
             List<String> backendRoles = getBackendRoles(inputs);
             AccessMode modelAccessMode = (AccessMode) inputs.get(MODEL_ACCESS_MODE);
-            Boolean isAddAllBackendRoles = (Boolean) inputs.get(ADD_ALL_BACKEND_ROLES);
+            Boolean isAddAllBackendRoles = inputs.containsKey(ADD_ALL_BACKEND_ROLES)
+                ? Booleans.parseBoolean(inputs.get(ADD_ALL_BACKEND_ROLES).toString())
+                : null;
 
             MLRegisterModelGroupInputBuilder builder = MLRegisterModelGroupInput.builder();
             builder.name(modelGroupName);
@@ -158,6 +162,8 @@ public void onFailure(Exception e) {
             MLRegisterModelGroupInput mlInput = builder.build();
 
             mlClient.registerModelGroup(mlInput, actionListener);
+        } catch (IllegalArgumentException iae) {
+            registerModelGroupFuture.onFailure(new WorkflowStepException(iae.getMessage(), RestStatus.BAD_REQUEST));
         } catch (FlowFrameworkException e) {
             registerModelGroupFuture.onFailure(e);
         }

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterRemoteModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterRemoteModelStep.java
@@ -13,7 +13,9 @@
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.action.update.UpdateResponse;
+import org.opensearch.common.Booleans;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.exception.WorkflowStepException;
 import org.opensearch.flowframework.indices.FlowFrameworkIndicesHandler;
@@ -90,7 +92,7 @@ public PlainActionFuture<WorkflowData> execute(
             String description = (String) inputs.get(DESCRIPTION_FIELD);
             String connectorId = (String) inputs.get(CONNECTOR_ID);
             Guardrails guardRails = (Guardrails) inputs.get(GUARDRAILS_FIELD);
-            final Boolean deploy = (Boolean) inputs.get(DEPLOY_FIELD);
+            final Boolean deploy = inputs.containsKey(DEPLOY_FIELD) ? Booleans.parseBoolean(inputs.get(DEPLOY_FIELD).toString()) : null;
 
             MLRegisterModelInputBuilder builder = MLRegisterModelInput.builder()
                 .functionName(FunctionName.REMOTE)
@@ -106,7 +108,6 @@ public PlainActionFuture<WorkflowData> execute(
             if (deploy != null) {
                 builder.deployModel(deploy);
             }
-
             if (guardRails != null) {
                 builder.guardrails(guardRails);
             }
@@ -190,6 +191,8 @@ public void onFailure(Exception e) {
                 }
             });
 
+        } catch (IllegalArgumentException iae) {
+            registerRemoteModelFuture.onFailure(new WorkflowStepException(iae.getMessage(), RestStatus.BAD_REQUEST));
         } catch (FlowFrameworkException e) {
             registerRemoteModelFuture.onFailure(e);
         }

diff --git a/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java b/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java
@@ -11,7 +11,10 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.action.support.PlainActionFuture;
+import org.opensearch.common.Booleans;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
+import org.opensearch.flowframework.exception.WorkflowStepException;
 import org.opensearch.flowframework.util.ParseUtils;
 import org.opensearch.ml.common.agent.MLToolSpec;
 
@@ -61,7 +64,9 @@ public PlainActionFuture<WorkflowData> execute(
             String type = (String) inputs.get(TYPE);
             String name = (String) inputs.get(NAME_FIELD);
             String description = (String) inputs.get(DESCRIPTION_FIELD);
-            Boolean includeOutputInAgentResponse = (Boolean) inputs.get(INCLUDE_OUTPUT_IN_AGENT_RESPONSE);
+            Boolean includeOutputInAgentResponse = inputs.containsKey(INCLUDE_OUTPUT_IN_AGENT_RESPONSE)
+                ? Booleans.parseBoolean(inputs.get(INCLUDE_OUTPUT_IN_AGENT_RESPONSE).toString())
+                : null;
             Map<String, String> parameters = getToolsParametersMap(inputs.get(PARAMETERS_FIELD), previousNodeInputs, outputs);
 
             MLToolSpec.MLToolSpecBuilder builder = MLToolSpec.builder();
@@ -92,6 +97,8 @@ public PlainActionFuture<WorkflowData> execute(
 
             logger.info("Tool registered successfully {}", type);
 
+        } catch (IllegalArgumentException iae) {
+            toolFuture.onFailure(new WorkflowStepException(iae.getMessage(), RestStatus.BAD_REQUEST));
         } catch (FlowFrameworkException e) {
             toolFuture.onFailure(e);
         }

diff --git a/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalCustomModelStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/RegisterLocalCustomModelStepTests.java
@@ -17,8 +17,10 @@
 import org.opensearch.common.util.concurrent.OpenSearchExecutors;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.index.shard.ShardId;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.common.FlowFrameworkSettings;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
+import org.opensearch.flowframework.exception.WorkflowStepException;
 import org.opensearch.flowframework.indices.FlowFrameworkIndicesHandler;
 import org.opensearch.ml.client.MachineLearningNodeClient;
 import org.opensearch.ml.common.MLTask;
@@ -31,6 +33,7 @@
 import org.opensearch.threadpool.ThreadPool;
 import org.junit.AfterClass;
 
+import java.io.IOException;
 import java.util.Collections;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
@@ -40,6 +43,7 @@
 import org.mockito.MockitoAnnotations;
 
 import static org.opensearch.action.DocWriteResponse.Result.UPDATED;
+import static org.opensearch.flowframework.common.CommonValue.DEPLOY_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.FLOW_FRAMEWORK_THREAD_POOL_PREFIX;
 import static org.opensearch.flowframework.common.CommonValue.PROVISION_WORKFLOW_THREAD_POOL;
 import static org.opensearch.flowframework.common.CommonValue.REGISTER_MODEL_STATUS;
@@ -184,6 +188,41 @@ public void testRegisterLocalCustomModelSuccess() throws Exception {
 
         assertEquals(modelId, future.get().getContent().get(MODEL_ID));
         assertEquals(status, future.get().getContent().get(REGISTER_MODEL_STATUS));
+
+        WorkflowData boolStringWorkflowData = new WorkflowData(
+            Map.ofEntries(
+                Map.entry("name", "xyz"),
+                Map.entry("version", "1.0.0"),
+                Map.entry("description", "description"),
+                Map.entry("function_name", "SPARSE_TOKENIZE"),
+                Map.entry("model_format", "TORCH_SCRIPT"),
+                Map.entry(MODEL_GROUP_ID, "abcdefg"),
+                Map.entry("model_content_hash_value", "aiwoeifjoaijeofiwe"),
+                Map.entry("model_type", "bert"),
+                Map.entry("embedding_dimension", "384"),
+                Map.entry("framework_type", "sentence_transformers"),
+                Map.entry("url", "something.com"),
+                Map.entry(DEPLOY_FIELD, "false")
+            ),
+            "test-id",
+            "test-node-id"
+        );
+
+        future = registerLocalModelStep.execute(
+            boolStringWorkflowData.getNodeId(),
+            boolStringWorkflowData,
+            Collections.emptyMap(),
+            Collections.emptyMap(),
+            Collections.emptyMap()
+        );
+
+        future.actionGet();
+
+        verify(machineLearningNodeClient, times(2)).register(any(MLRegisterModelInput.class), any());
+        verify(machineLearningNodeClient, times(2)).getTask(any(), any());
+
+        assertEquals(modelId, future.get().getContent().get(MODEL_ID));
+        assertEquals(status, future.get().getContent().get(REGISTER_MODEL_STATUS));
     }
 
     public void testRegisterLocalCustomModelFailure() {
@@ -283,4 +322,40 @@ public void testMissingInputs() {
         }
         assertTrue(ex.getCause().getMessage().endsWith("] in workflow [test-id] node [test-node-id]"));
     }
+
+    public void testBoolParseFail() throws IOException, ExecutionException, InterruptedException {
+        WorkflowData boolStringWorkflowData = new WorkflowData(
+            Map.ofEntries(
+                Map.entry("name", "xyz"),
+                Map.entry("version", "1.0.0"),
+                Map.entry("description", "description"),
+                Map.entry("function_name", "SPARSE_TOKENIZE"),
+                Map.entry("model_format", "TORCH_SCRIPT"),
+                Map.entry(MODEL_GROUP_ID, "abcdefg"),
+                Map.entry("model_content_hash_value", "aiwoeifjoaijeofiwe"),
+                Map.entry("model_type", "bert"),
+                Map.entry("embedding_dimension", "384"),
+                Map.entry("framework_type", "sentence_transformers"),
+                Map.entry("url", "something.com"),
+                Map.entry(DEPLOY_FIELD, "no")
+            ),
+            "test-id",
+            "test-node-id"
+        );
+
+        PlainActionFuture<WorkflowData> future = registerLocalModelStep.execute(
+            boolStringWorkflowData.getNodeId(),
+            boolStringWorkflowData,
+            Collections.emptyMap(),
+            Collections.emptyMap(),
+            Collections.emptyMap()
+        );
+
+        assertTrue(future.isDone());
+        ExecutionException e = assertThrows(ExecutionException.class, () -> future.get());
+        assertEquals(WorkflowStepException.class, e.getCause().getClass());
+        WorkflowStepException w = (WorkflowStepException) e.getCause();
+        assertEquals("Failed to parse value [no] as only [true] or [false] are allowed.", w.getMessage());
+        assertEquals(RestStatus.BAD_REQUEST, w.getRestStatus());
+    }
 }
diff --git a/...est/java/org/opensearch/flowframework/workflow/RegisterLocalPretrainedModelStepTests.java b/...est/java/org/opensearch/flowframework/workflow/RegisterLocalPretrainedModelStepTests.java
@@ -17,8 +17,10 @@
 import org.opensearch.common.util.concurrent.OpenSearchExecutors;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.index.shard.ShardId;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.common.FlowFrameworkSettings;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
+import org.opensearch.flowframework.exception.WorkflowStepException;
 import org.opensearch.flowframework.indices.FlowFrameworkIndicesHandler;
 import org.opensearch.ml.client.MachineLearningNodeClient;
 import org.opensearch.ml.common.MLTask;
@@ -31,6 +33,7 @@
 import org.opensearch.threadpool.ThreadPool;
 import org.junit.AfterClass;
 
+import java.io.IOException;
 import java.util.Collections;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
@@ -40,6 +43,7 @@
 import org.mockito.MockitoAnnotations;
 
 import static org.opensearch.action.DocWriteResponse.Result.UPDATED;
+import static org.opensearch.flowframework.common.CommonValue.DEPLOY_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.FLOW_FRAMEWORK_THREAD_POOL_PREFIX;
 import static org.opensearch.flowframework.common.CommonValue.PROVISION_WORKFLOW_THREAD_POOL;
 import static org.opensearch.flowframework.common.CommonValue.REGISTER_MODEL_STATUS;
@@ -178,6 +182,35 @@ public void testRegisterLocalPretrainedModelSuccess() throws Exception {
 
         assertEquals(modelId, future.get().getContent().get(MODEL_ID));
         assertEquals(status, future.get().getContent().get(REGISTER_MODEL_STATUS));
+
+        WorkflowData boolStringWorkflowData = new WorkflowData(
+            Map.ofEntries(
+                Map.entry("name", "xyz"),
+                Map.entry("version", "1.0.0"),
+                Map.entry("model_format", "TORCH_SCRIPT"),
+                Map.entry(MODEL_GROUP_ID, "abcdefg"),
+                Map.entry("description", "aiwoeifjoaijeofiwe"),
+                Map.entry(DEPLOY_FIELD, "false")
+            ),
+            "test-id",
+            "test-node-id"
+        );
+
+        future = registerLocalPretrainedModelStep.execute(
+            boolStringWorkflowData.getNodeId(),
+            boolStringWorkflowData,
+            Collections.emptyMap(),
+            Collections.emptyMap(),
+            Collections.emptyMap()
+        );
+
+        future.actionGet();
+
+        verify(machineLearningNodeClient, times(2)).register(any(MLRegisterModelInput.class), any());
+        verify(machineLearningNodeClient, times(2)).getTask(any(), any());
+
+        assertEquals(modelId, future.get().getContent().get(MODEL_ID));
+        assertEquals(status, future.get().getContent().get(REGISTER_MODEL_STATUS));
     }
 
     public void testRegisterLocalPretrainedModelFailure() {
@@ -268,4 +301,34 @@ public void testMissingInputs() {
         }
         assertTrue(ex.getCause().getMessage().endsWith("] in workflow [test-id] node [test-node-id]"));
     }
+
+    public void testBoolParseFail() throws IOException, ExecutionException, InterruptedException {
+        WorkflowData boolStringWorkflowData = new WorkflowData(
+            Map.ofEntries(
+                Map.entry("name", "xyz"),
+                Map.entry("version", "1.0.0"),
+                Map.entry("model_format", "TORCH_SCRIPT"),
+                Map.entry(MODEL_GROUP_ID, "abcdefg"),
+                Map.entry("description", "aiwoeifjoaijeofiwe"),
+                Map.entry(DEPLOY_FIELD, "no")
+            ),
+            "test-id",
+            "test-node-id"
+        );
+
+        PlainActionFuture<WorkflowData> future = registerLocalPretrainedModelStep.execute(
+            boolStringWorkflowData.getNodeId(),
+            boolStringWorkflowData,
+            Collections.emptyMap(),
+            Collections.emptyMap(),
+            Collections.emptyMap()
+        );
+
+        assertTrue(future.isDone());
+        ExecutionException e = assertThrows(ExecutionException.class, () -> future.get());
+        assertEquals(WorkflowStepException.class, e.getCause().getClass());
+        WorkflowStepException w = (WorkflowStepException) e.getCause();
+        assertEquals("Failed to parse value [no] as only [true] or [false] are allowed.", w.getMessage());
+        assertEquals(RestStatus.BAD_REQUEST, w.getRestStatus());
+    }
 }