Update Delete Connector Step with parsing util

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Dec 2, 2023 · ff9d02a · ff9d02a
1 parent 3f073c2
commit ff9d02a
Show file tree

Hide file tree

Showing 8 changed files with 104 additions and 110 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/util/ParseUtils.java b/src/main/java/org/opensearch/flowframework/util/ParseUtils.java
@@ -8,21 +8,6 @@
  */
 package org.opensearch.flowframework.util;
 
-import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
-import static org.opensearch.flowframework.common.CommonValue.MODEL_ID;
-import static org.opensearch.flowframework.common.CommonValue.PARAMETERS_FIELD;
-
-import java.io.IOException;
-import java.time.Instant;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Map.Entry;
-import java.util.Optional;
-import java.util.Set;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.client.Client;
@@ -41,6 +26,21 @@
 import org.opensearch.flowframework.workflow.WorkflowData;
 import org.opensearch.ml.common.agent.LLMSpec;
 
+import java.io.IOException;
+import java.time.Instant;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Optional;
+import java.util.Set;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
+import static org.opensearch.flowframework.common.CommonValue.MODEL_ID;
+import static org.opensearch.flowframework.common.CommonValue.PARAMETERS_FIELD;
+
 /**
  * Utility methods for Template parsing
  */

diff --git a/src/main/java/org/opensearch/flowframework/workflow/CreateConnectorStep.java b/src/main/java/org/opensearch/flowframework/workflow/CreateConnectorStep.java
@@ -8,14 +8,20 @@
  */
 package org.opensearch.flowframework.workflow;
 
-import static org.opensearch.flowframework.common.CommonValue.ACTIONS_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.CREDENTIAL_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.PARAMETERS_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.PROTOCOL_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.VERSION_FIELD;
-import static org.opensearch.flowframework.util.ParseUtils.getStringToStringMap;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.ExceptionsHelper;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.rest.RestStatus;
+import org.opensearch.flowframework.common.WorkflowResources;
+import org.opensearch.flowframework.exception.FlowFrameworkException;
+import org.opensearch.flowframework.indices.FlowFrameworkIndicesHandler;
+import org.opensearch.flowframework.util.ParseUtils;
+import org.opensearch.ml.client.MachineLearningNodeClient;
+import org.opensearch.ml.common.connector.ConnectorAction;
+import org.opensearch.ml.common.connector.ConnectorAction.ActionType;
+import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
+import org.opensearch.ml.common.transport.connector.MLCreateConnectorResponse;
 
 import java.io.IOException;
 import java.security.AccessController;
@@ -31,20 +37,14 @@
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
-import org.opensearch.ExceptionsHelper;
-import org.opensearch.core.action.ActionListener;
-import org.opensearch.core.rest.RestStatus;
-import org.opensearch.flowframework.common.WorkflowResources;
-import org.opensearch.flowframework.exception.FlowFrameworkException;
-import org.opensearch.flowframework.indices.FlowFrameworkIndicesHandler;
-import org.opensearch.flowframework.util.ParseUtils;
-import org.opensearch.ml.client.MachineLearningNodeClient;
-import org.opensearch.ml.common.connector.ConnectorAction;
-import org.opensearch.ml.common.connector.ConnectorAction.ActionType;
-import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
-import org.opensearch.ml.common.transport.connector.MLCreateConnectorResponse;
+import static org.opensearch.flowframework.common.CommonValue.ACTIONS_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.CREDENTIAL_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.PARAMETERS_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.PROTOCOL_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.VERSION_FIELD;
+import static org.opensearch.flowframework.util.ParseUtils.getStringToStringMap;
 
 /**
  * Step to create a connector for a remote model

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeleteConnectorStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeleteConnectorStep.java
@@ -13,13 +13,14 @@
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.action.delete.DeleteResponse;
 import org.opensearch.core.action.ActionListener;
-import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
+import org.opensearch.flowframework.util.ParseUtils;
 import org.opensearch.ml.client.MachineLearningNodeClient;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.Map;
-import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 
 import static org.opensearch.flowframework.common.CommonValue.CONNECTOR_ID;
@@ -72,29 +73,23 @@ public void onFailure(Exception e) {
             }
         };
 
-        String connectorId = null;
-
-        // Previous Node inputs defines which step the connector ID came from
-        Optional<String> previousNode = previousNodeInputs.entrySet()
-            .stream()
-            .filter(e -> CONNECTOR_ID.equals(e.getValue()))
-            .map(Map.Entry::getKey)
-            .findFirst();
-        if (previousNode.isPresent()) {
-            WorkflowData previousNodeOutput = outputs.get(previousNode.get());
-            if (previousNodeOutput != null && previousNodeOutput.getContent().containsKey(CONNECTOR_ID)) {
-                connectorId = previousNodeOutput.getContent().get(CONNECTOR_ID).toString();
-            }
-        }
+        Set<String> requiredKeys = Set.of(CONNECTOR_ID);
+        Set<String> optionalKeys = Collections.emptySet();
 
-        if (connectorId != null) {
-            mlClient.deleteConnector(connectorId, actionListener);
-        } else {
-            deleteConnectorFuture.completeExceptionally(
-                new FlowFrameworkException("Required field " + CONNECTOR_ID + " is not provided", RestStatus.BAD_REQUEST)
+        try {
+            Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
+                requiredKeys,
+                optionalKeys,
+                currentNodeInputs,
+                outputs,
+                previousNodeInputs
             );
-        }
+            String connectorId = (String) inputs.get(CONNECTOR_ID);
 
+            mlClient.deleteConnector(connectorId, actionListener);
+        } catch (FlowFrameworkException e) {
+            deleteConnectorFuture.completeExceptionally(e);
+        }
         return deleteConnectorFuture;
     }
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java
@@ -8,19 +8,6 @@
  */
 package org.opensearch.flowframework.workflow;
 
-import static org.opensearch.flowframework.common.CommonValue.ADD_ALL_BACKEND_ROLES;
-import static org.opensearch.flowframework.common.CommonValue.BACKEND_ROLES_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.MODEL_ACCESS_MODE;
-import static org.opensearch.flowframework.common.CommonValue.MODEL_GROUP_STATUS;
-import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD;
-
-import java.io.IOException;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.CompletableFuture;
-
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
@@ -35,6 +22,19 @@
 import org.opensearch.ml.common.transport.model_group.MLRegisterModelGroupInput.MLRegisterModelGroupInputBuilder;
 import org.opensearch.ml.common.transport.model_group.MLRegisterModelGroupResponse;
 
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+
+import static org.opensearch.flowframework.common.CommonValue.ADD_ALL_BACKEND_ROLES;
+import static org.opensearch.flowframework.common.CommonValue.BACKEND_ROLES_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.MODEL_ACCESS_MODE;
+import static org.opensearch.flowframework.common.CommonValue.MODEL_GROUP_STATUS;
+import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD;
+
 /**
  * Step to register a model group
  */

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterLocalModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterLocalModelStep.java
@@ -8,24 +8,6 @@
  */
 package org.opensearch.flowframework.workflow;
 
-import static org.opensearch.flowframework.common.CommonValue.ALL_CONFIG;
-import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.EMBEDDING_DIMENSION;
-import static org.opensearch.flowframework.common.CommonValue.FRAMEWORK_TYPE;
-import static org.opensearch.flowframework.common.CommonValue.MODEL_CONTENT_HASH_VALUE;
-import static org.opensearch.flowframework.common.CommonValue.MODEL_FORMAT;
-import static org.opensearch.flowframework.common.CommonValue.MODEL_GROUP_ID;
-import static org.opensearch.flowframework.common.CommonValue.MODEL_TYPE;
-import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.REGISTER_MODEL_STATUS;
-import static org.opensearch.flowframework.common.CommonValue.URL;
-import static org.opensearch.flowframework.common.CommonValue.VERSION_FIELD;
-
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.CompletableFuture;
-import java.util.stream.Stream;
-
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
@@ -49,6 +31,24 @@
 import org.opensearch.ml.common.transport.register.MLRegisterModelInput.MLRegisterModelInputBuilder;
 import org.opensearch.ml.common.transport.register.MLRegisterModelResponse;
 
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Stream;
+
+import static org.opensearch.flowframework.common.CommonValue.ALL_CONFIG;
+import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.EMBEDDING_DIMENSION;
+import static org.opensearch.flowframework.common.CommonValue.FRAMEWORK_TYPE;
+import static org.opensearch.flowframework.common.CommonValue.MODEL_CONTENT_HASH_VALUE;
+import static org.opensearch.flowframework.common.CommonValue.MODEL_FORMAT;
+import static org.opensearch.flowframework.common.CommonValue.MODEL_GROUP_ID;
+import static org.opensearch.flowframework.common.CommonValue.MODEL_TYPE;
+import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.REGISTER_MODEL_STATUS;
+import static org.opensearch.flowframework.common.CommonValue.URL;
+import static org.opensearch.flowframework.common.CommonValue.VERSION_FIELD;
+
 /**
  * Step to register a local model
  */

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterRemoteModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterRemoteModelStep.java
@@ -8,18 +8,6 @@
  */
 package org.opensearch.flowframework.workflow;
 
-import static org.opensearch.flowframework.common.CommonValue.CONNECTOR_ID;
-import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.FUNCTION_NAME;
-import static org.opensearch.flowframework.common.CommonValue.MODEL_GROUP_ID;
-import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD;
-import static org.opensearch.flowframework.common.CommonValue.REGISTER_MODEL_STATUS;
-
-import java.util.Locale;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.CompletableFuture;
-
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
@@ -34,6 +22,18 @@
 import org.opensearch.ml.common.transport.register.MLRegisterModelInput.MLRegisterModelInputBuilder;
 import org.opensearch.ml.common.transport.register.MLRegisterModelResponse;
 
+import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+
+import static org.opensearch.flowframework.common.CommonValue.CONNECTOR_ID;
+import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.FUNCTION_NAME;
+import static org.opensearch.flowframework.common.CommonValue.MODEL_GROUP_ID;
+import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.REGISTER_MODEL_STATUS;
+
 /**
  * Step to register a remote model
  */

diff --git a/src/test/java/org/opensearch/flowframework/util/ParseUtilsTests.java b/src/test/java/org/opensearch/flowframework/util/ParseUtilsTests.java
@@ -8,11 +8,6 @@
  */
 package org.opensearch.flowframework.util;
 
-import java.io.IOException;
-import java.time.Instant;
-import java.util.Map;
-import java.util.Set;
-
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -21,6 +16,11 @@
 import org.opensearch.flowframework.workflow.WorkflowData;
 import org.opensearch.test.OpenSearchTestCase;
 
+import java.io.IOException;
+import java.time.Instant;
+import java.util.Map;
+import java.util.Set;
+
 public class ParseUtilsTests extends OpenSearchTestCase {
     public void testToInstant() throws IOException {
         long epochMilli = Instant.now().toEpochMilli();

diff --git a/src/test/java/org/opensearch/flowframework/workflow/DeleteConnectorStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/DeleteConnectorStepTests.java
@@ -13,7 +13,6 @@
 import org.opensearch.core.index.Index;
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.core.rest.RestStatus;
-import org.opensearch.flowframework.common.CommonValue;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.ml.client.MachineLearningNodeClient;
 import org.opensearch.ml.common.transport.connector.MLCreateConnectorResponse;
@@ -44,7 +43,7 @@ public void setUp() throws Exception {
 
         MockitoAnnotations.openMocks(this);
 
-        inputData = new WorkflowData(Map.of(CommonValue.CONNECTOR_ID, "test"), "test-id", "test-node-id");
+        inputData = new WorkflowData(Collections.emptyMap(), "test-id", "test-node-id");
     }
 
     public void testDeleteConnector() throws IOException, ExecutionException, InterruptedException {
@@ -86,7 +85,7 @@ public void testNoConnectorIdInOutput() throws IOException {
         assertTrue(future.isCompletedExceptionally());
         ExecutionException ex = assertThrows(ExecutionException.class, () -> future.get().getContent());
         assertTrue(ex.getCause() instanceof FlowFrameworkException);
-        assertEquals("Required field connector_id is not provided", ex.getCause().getMessage());
+        assertEquals("Missing required inputs [connector_id] in workflow [test-id] node [test-node-id]", ex.getCause().getMessage());
     }
 
     public void testDeleteConnectorFailure() throws IOException {