Addressed PR comments

Signed-off-by: Owais Kazi <[email protected]>
opensearch-project · Oct 11, 2023 · c744902 · c744902
1 parent 0d894bb
commit c744902
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 11 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/common/CommonValue.java b/src/main/java/org/opensearch/flowframework/common/CommonValue.java
@@ -19,4 +19,13 @@ public class CommonValue {
     public static final String GLOBAL_CONTEXT_INDEX = ".plugins-ai-global-context";
     public static final String GLOBAL_CONTEXT_INDEX_MAPPING = "mappings/global-context.json";
     public static final Integer GLOBAL_CONTEXT_INDEX_VERSION = 1;
+    public static final String MODEL_ID = "model_id";
+    public static final String FUNCTION_NAME = "function_name";
+    public static final String MODEL_NAME = "name";
+    public static final String MODEL_VERSION = "model_version";
+    public static final String MODEL_GROUP_ID = "model_group_id";
+    public static final String DESCRIPTION = "description";
+    public static final String CONNECTOR_ID = "connector_id";
+    public static final String MODEL_FORMAT = "model_format";
+    public static final String MODEL_CONFIG = "model_config";
 }
diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java
@@ -20,14 +20,15 @@
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 
+import static org.opensearch.flowframework.common.CommonValue.MODEL_ID;
+
 /**
  * Step to deploy a model
  */
 public class DeployModelStep implements WorkflowStep {
     private static final Logger logger = LogManager.getLogger(DeployModelStep.class);
 
     private Client client;
-    private static final String MODEL_ID = "model_id";
     static final String NAME = "deploy_model";
 
     /**
@@ -66,8 +67,9 @@ public void onFailure(Exception e) {
         for (WorkflowData workflowData : data) {
             Map<String, Object> content = workflowData.getContent();
             for (Map.Entry<String, Object> entry : content.entrySet()) {
-                if (entry.getKey() == MODEL_ID) {
+                if (MODEL_ID.equals(entry.getKey())) {
                     modelId = (String) content.get(MODEL_ID);
+                    break;
                 }
 
             }

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelStep.java
@@ -28,6 +28,15 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Stream;
 
+import static org.opensearch.flowframework.common.CommonValue.CONNECTOR_ID;
+import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION;
+import static org.opensearch.flowframework.common.CommonValue.FUNCTION_NAME;
+import static org.opensearch.flowframework.common.CommonValue.MODEL_CONFIG;
+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_NAME;
+import static org.opensearch.flowframework.common.CommonValue.MODEL_VERSION;
+
 /**
  * Step to register a remote model
  */
@@ -39,15 +48,6 @@ public class RegisterModelStep implements WorkflowStep {
 
     static final String NAME = "register_model";
 
-    private static final String FUNCTION_NAME = "function_name";
-    private static final String MODEL_NAME = "name";
-    private static final String MODEL_VERSION = "model_version";
-    private static final String MODEL_GROUP_ID = "model_group_id";
-    private static final String DESCRIPTION = "description";
-    private static final String CONNECTOR_ID = "connector_id";
-    private static final String MODEL_FORMAT = "model_format";
-    private static final String MODEL_CONFIG = "model_config";
-
     /**
      * Instantiate this class
      * @param client client to instantiate MLClient