From 8d3f42c381e5f21152ef049a1c12cf02d0a3e561 Mon Sep 17 00:00:00 2001 From: martinpkr Date: Thu, 30 May 2024 15:51:20 +0300 Subject: [PATCH] Added a new parse util method to avoid repetition / refactored code with new method --- .../java/org/opensearch/flowframework/util/ParseUtils.java | 5 +++++ .../workflow/AbstractRegisterLocalModelStep.java | 2 +- .../flowframework/workflow/RegisterModelGroupStep.java | 4 +--- .../flowframework/workflow/RegisterRemoteModelStep.java | 2 +- .../java/org/opensearch/flowframework/workflow/ToolStep.java | 4 +--- 5 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/util/ParseUtils.java b/src/main/java/org/opensearch/flowframework/util/ParseUtils.java index b4b908aeb..3d509cb35 100644 --- 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,8 @@ public static Map convertStringToObjectMapToStringToStringMap(Ma return stringToStringMap; } } + + public static Boolean checkIfInputsContainsKey(Map inputs, String enumKey){ + return inputs.containsKey(enumKey) ? Booleans.parseBoolean(inputs.get(enumKey).toString()) : null; + } } diff --git a/src/main/java/org/opensearch/flowframework/workflow/AbstractRegisterLocalModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/AbstractRegisterLocalModelStep.java index 80294a7cd..39d5b19c4 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/AbstractRegisterLocalModelStep.java +++ b/src/main/java/org/opensearch/flowframework/workflow/AbstractRegisterLocalModelStep.java @@ -123,7 +123,7 @@ public PlainActionFuture 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.checkIfInputsContainsKey(inputs, DEPLOY_FIELD); // 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 index 9771d7c77..09b52ac9b 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java +++ b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java @@ -147,9 +147,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.checkIfInputsContainsKey(inputs, ADD_ALL_BACKEND_ROLES); 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 index db5d290ca..b18103381 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/RegisterRemoteModelStep.java +++ b/src/main/java/org/opensearch/flowframework/workflow/RegisterRemoteModelStep.java @@ -100,7 +100,7 @@ public PlainActionFuture 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.checkIfInputsContainsKey(inputs, DEPLOY_FIELD); 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 index 87ecf762e..5a5d4b703 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java +++ b/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java @@ -64,9 +64,7 @@ public PlainActionFuture 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.checkIfInputsContainsKey(inputs ,INCLUDE_OUTPUT_IN_AGENT_RESPONSE); Map parameters = getToolsParametersMap(inputs.get(PARAMETERS_FIELD), previousNodeInputs, outputs); MLToolSpec.MLToolSpecBuilder builder = MLToolSpec.builder();