[Backport 2.x] Added a new parse util method to avoid repetition / re…

…factored code with new method (opensearch-project#728) Added a new parse util method to avoid repetition / refactored code with new method (opensearch-project#721) * Add user mapping to Workflow State index (opensearch-project#705) * Add user mapping to Workflow State index * Increment schema version --------- * Added a new parse util method to avoid repetition / refactored code with new method * refactored method name and added unit test * made method use generics + added test * fixed javadoc * Added workflow step for ReIndex Step (opensearch-project#718) * Initial commit for reindex workflow step with extra params * Addressed PR comments * Changed request per second to Float * Addressed string array for source indices and removed state index entry * Minor comments --------- * Incorporating parseIfExist method into ReindexStep class * Add param to delete workflow API to clear status even if resources exist (opensearch-project#719) * refactored method to use parseBoolean and parseFloat methods * Adding a missing param in javaDoc * Added workflow step for ReIndex Step (opensearch-project#718) * Initial commit for reindex workflow step with extra params * Addressed PR comments * Changed request per second to Float * Addressed string array for source indices and removed state index entry * Minor comments --------- * Add param to delete workflow API to clear status even if resources exist (opensearch-project#719) * Added a chagelog entry * fixed failing spotless check * Added workflow step for ReIndex Step (opensearch-project#718) * Initial commit for reindex workflow step with extra params * Addressed PR comments * Changed request per second to Float * Addressed string array for source indices and removed state index entry * Minor comments --------- * removed unnecessary changelog info --------- (cherry picked from commit 13b32f1) Signed-off-by: Daniel Widdis <[email protected]> Signed-off-by: martinpkr <[email protected]> Signed-off-by: owaiskazi19 <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Daniel Widdis <[email protected]> Co-authored-by: Owais Kazi <[email protected]>
amitgalitz · Jun 3, 2024 · 9e0fc98 · 9e0fc98
1 parent e06abcb
commit 9e0fc98
Show file tree

Hide file tree

Showing 7 changed files with 62 additions and 13 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -17,7 +17,6 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.1.0/)
 ### Enhancements
 - Add Workflow Step for Reindex from source index to destination ([#718](https://github.com/opensearch-project/flow-framework/pull/718))
 - Add param to delete workflow API to clear status even if resources exist ([#719](https://github.com/opensearch-project/flow-framework/pull/719))
-
 ### Bug Fixes
 - Add user mapping to Workflow State index ([#705](https://github.com/opensearch-project/flow-framework/pull/705))
 

diff --git a/src/main/java/org/opensearch/flowframework/util/ParseUtils.java b/src/main/java/org/opensearch/flowframework/util/ParseUtils.java
@@ -11,6 +11,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.client.Client;
+import org.opensearch.common.Booleans;
 import org.opensearch.common.io.Streams;
 import org.opensearch.common.xcontent.LoggingDeprecationHandler;
 import org.opensearch.common.xcontent.XContentHelper;
@@ -472,4 +473,29 @@ public static Map<String, String> convertStringToObjectMapToStringToStringMap(Ma
             return stringToStringMap;
         }
     }
+
+    /**
+     * Checks if the inputs map contains the specified key and parses the associated value to a generic class.
+     *
+     * @param <T> the type to which the value should be parsed
+     * @param inputs the map containing the input data
+     * @param key the key to check in the map
+     * @param type the class to parse the value to
+     * @throws IllegalArgumentException if the type is not supported
+     * @return the generic type value associated with the key if present, or null if the key is not found
+     */
+    public static <T> T parseIfExists(Map<String, Object> inputs, String key, Class<T> type) {
+        if (!inputs.containsKey(key)) {
+            return null;
+        }
+
+        Object value = inputs.get(key);
+        if (type == Boolean.class) {
+            return type.cast(Booleans.parseBoolean(value.toString()));
+        } else if (type == Float.class) {
+            return type.cast(Float.parseFloat(value.toString()));
+        } else {
+            throw new IllegalArgumentException("Unsupported type: " + type);
+        }
+    }
 }
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,6 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.support.PlainActionFuture;
-import org.opensearch.common.Booleans;
 import org.opensearch.common.xcontent.XContentHelper;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.bytes.BytesArray;
@@ -123,7 +122,7 @@ public PlainActionFuture<WorkflowData> execute(
             String modelGroupId = (String) inputs.get(MODEL_GROUP_ID);
             String allConfig = (String) inputs.get(ALL_CONFIG);
             String modelInterface = (String) inputs.get(INTERFACE_FIELD);
-            final Boolean deploy = inputs.containsKey(DEPLOY_FIELD) ? Booleans.parseBoolean(inputs.get(DEPLOY_FIELD).toString()) : null;
+            final Boolean deploy = ParseUtils.parseIfExists(inputs, DEPLOY_FIELD, Boolean.class);
 
             // Build register model input
             MLRegisterModelInputBuilder mlInputBuilder = MLRegisterModelInput.builder()

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java
@@ -12,7 +12,6 @@
 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;
@@ -147,9 +146,7 @@ public void onFailure(Exception ex) {
             if (inputs.containsKey(MODEL_ACCESS_MODE)) {
                 modelAccessMode = AccessMode.from((inputs.get(MODEL_ACCESS_MODE)).toString().toLowerCase(Locale.ROOT));
             }
-            Boolean isAddAllBackendRoles = inputs.containsKey(ADD_ALL_BACKEND_ROLES)
-                ? Booleans.parseBoolean(inputs.get(ADD_ALL_BACKEND_ROLES).toString())
-                : null;
+            Boolean isAddAllBackendRoles = ParseUtils.parseIfExists(inputs, ADD_ALL_BACKEND_ROLES, Boolean.class);
 
             MLRegisterModelGroupInputBuilder builder = MLRegisterModelGroupInput.builder();
             builder.name(modelGroupName);

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,6 @@
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.support.PlainActionFuture;
 import org.opensearch.action.update.UpdateResponse;
-import org.opensearch.common.Booleans;
 import org.opensearch.common.xcontent.XContentHelper;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.bytes.BytesArray;
@@ -100,7 +99,7 @@ public PlainActionFuture<WorkflowData> execute(
             String connectorId = (String) inputs.get(CONNECTOR_ID);
             Guardrails guardRails = (Guardrails) inputs.get(GUARDRAILS_FIELD);
             String modelInterface = (String) inputs.get(INTERFACE_FIELD);
-            final Boolean deploy = inputs.containsKey(DEPLOY_FIELD) ? Booleans.parseBoolean(inputs.get(DEPLOY_FIELD).toString()) : null;
+            final Boolean deploy = ParseUtils.parseIfExists(inputs, DEPLOY_FIELD, Boolean.class);
 
             MLRegisterModelInputBuilder builder = MLRegisterModelInput.builder()
                 .functionName(FunctionName.REMOTE)

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,6 @@
 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;
@@ -64,9 +63,7 @@ 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 = inputs.containsKey(INCLUDE_OUTPUT_IN_AGENT_RESPONSE)
-                ? Booleans.parseBoolean(inputs.get(INCLUDE_OUTPUT_IN_AGENT_RESPONSE).toString())
-                : null;
+            Boolean includeOutputInAgentResponse = ParseUtils.parseIfExists(inputs, INCLUDE_OUTPUT_IN_AGENT_RESPONSE, Boolean.class);
             Map<String, String> parameters = getToolsParametersMap(inputs.get(PARAMETERS_FIELD), previousNodeInputs, outputs);
 
             MLToolSpec.MLToolSpecBuilder builder = MLToolSpec.builder();

diff --git a/src/test/java/org/opensearch/flowframework/util/ParseUtilsTests.java b/src/test/java/org/opensearch/flowframework/util/ParseUtilsTests.java
@@ -221,4 +221,36 @@ public void testGetInputsFromPreviousSteps() {
         assertEquals("Missing required inputs [not-here] in workflow [workflowId] node [nodeId]", e.getMessage());
         assertEquals(RestStatus.BAD_REQUEST, e.getRestStatus());
     }
+
+    public void testParseIfExistsWithBooleanClass() {
+        Map<String, Object> inputs = new HashMap<>();
+        inputs.put("key1", "true");
+        inputs.put("key2", "false");
+        inputs.put("key3", "true");
+
+        assertEquals(Boolean.TRUE, ParseUtils.parseIfExists(inputs, "key1", Boolean.class));
+        assertEquals(Boolean.FALSE, ParseUtils.parseIfExists(inputs, "key2", Boolean.class));
+        assertNull(ParseUtils.parseIfExists(inputs, "keyThatDoesntExist", Boolean.class));
+
+    }
+
+    public void testParseIfExistsWithFloatClass() {
+        Map<String, Object> inputs = new HashMap<>();
+        inputs.put("key1", "3.14");
+        inputs.put("key2", "0.01");
+        inputs.put("key3", "90.22");
+
+        assertEquals(Float.valueOf("3.14"), ParseUtils.parseIfExists(inputs, "key1", Float.class));
+        assertEquals(Float.valueOf("0.01"), ParseUtils.parseIfExists(inputs, "key2", Float.class));
+        assertNull(ParseUtils.parseIfExists(inputs, "keyThatDoesntExist", Float.class));
+
+    }
+
+    public void testParseIfExistWhenWrongTypeIsPassed() {
+
+        Map<String, Object> inputs = new HashMap<>();
+        inputs.put("key1", "3.14");
+
+        assertThrows(IllegalArgumentException.class, () -> ParseUtils.parseIfExists(inputs, "key1", Integer.class));
+    }
 }