Remove other non-required keys and handle nulls

Signed-off-by: Daniel Widdis <widdis@gmail.com>
dbwiddis · Dec 27, 2023 · 9bd5176 · 9bd5176
1 parent 3fadbd5
commit 9bd5176
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 20 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java
@@ -30,6 +30,7 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -155,10 +156,14 @@ public void onFailure(Exception e) {
             String name = (String) inputs.get(NAME_FIELD);
             String description = (String) inputs.get(DESCRIPTION_FIELD);
             String llmModelId = (String) inputs.get(LLM_MODEL_ID);
-            Map<String, String> llmParameters = getStringToStringMap(inputs.get(PARAMETERS_FIELD), LLM_PARAMETERS);
+            Object llmParams = inputs.get(LLM_PARAMETERS);
+            Map<String, String> llmParameters = llmParams == null
+                ? Collections.emptyMap()
+                : getStringToStringMap(llmParams, LLM_PARAMETERS);
             String[] toolsOrder = (String[]) inputs.get(TOOLS_ORDER_FIELD);
             List<MLToolSpec> toolsList = getTools(toolsOrder, previousNodeInputs, outputs);
-            Map<String, String> parameters = getStringToStringMap(inputs.get(PARAMETERS_FIELD), PARAMETERS_FIELD);
+            Object params = inputs.get(PARAMETERS_FIELD);
+            Map<String, String> parameters = params == null ? Collections.emptyMap() : getStringToStringMap(params, PARAMETERS_FIELD);
             MLMemorySpec memory = getMLMemorySpec(inputs.get(MEMORY_FIELD));
             Instant createdTime = Instant.now();
             Instant lastUpdateTime = createdTime;
@@ -223,7 +228,7 @@ private List<MLToolSpec> getTools(@Nullable String[] tools, Map<String, String>
         List<String> sortedNodes = tools == null ? new ArrayList<>() : Arrays.asList(tools);
         previousNodes.removeAll(sortedNodes);
         sortedNodes.addAll(previousNodes);
-        sortedNodes.forEach((node) -> {
+        sortedNodes.forEach(node -> {
             WorkflowData previousNodeOutput = outputs.get(node);
             if (previousNodeOutput != null && previousNodeOutput.getContent().containsKey(TOOLS_FIELD)) {
                 MLToolSpec mlToolSpec = (MLToolSpec) previousNodeOutput.getContent().get(TOOLS_FIELD);
@@ -235,9 +240,6 @@ private List<MLToolSpec> getTools(@Nullable String[] tools, Map<String, String>
     }
 
     private String getLlmModelId(Map<String, String> previousNodeInputs, Map<String, WorkflowData> outputs) {
-        // Case when modelId is already pass in the template
-        String llmModelId = null;
-
         // Case when modelId is passed through previousSteps
         Optional<String> previousNode = previousNodeInputs.entrySet()
             .stream()
@@ -247,11 +249,15 @@ private String getLlmModelId(Map<String, String> previousNodeInputs, Map<String,
 
         if (previousNode.isPresent()) {
             WorkflowData previousNodeOutput = outputs.get(previousNode.get());
-            if (previousNodeOutput != null && previousNodeOutput.getContent().containsKey(MODEL_ID)) {
-                llmModelId = previousNodeOutput.getContent().get(MODEL_ID).toString();
+            if (previousNodeOutput != null) {
+                // Use either llm.model_id (if present) or model_id (backup)
+                Object modelId = previousNodeOutput.getContent().getOrDefault(LLM_MODEL_ID, previousNodeOutput.getContent().get(MODEL_ID));
+                if (modelId != null) {
+                    return modelId.toString();
+                }
             }
         }
-        return llmModelId;
+        return null;
     }
 
     private LLMSpec getLLMSpec(String llmModelId, Map<String, String> llmParameters, String workflowId, String currentNodeId) {
@@ -267,8 +273,7 @@ private LLMSpec getLLMSpec(String llmModelId, Map<String, String> llmParameters,
             builder.parameters(llmParameters);
         }
 
-        LLMSpec llmSpec = builder.build();
-        return llmSpec;
+        return builder.build();
     }
 
     private MLMemorySpec getMLMemorySpec(Object mlMemory) {
@@ -294,9 +299,6 @@ private MLMemorySpec getMLMemorySpec(Object mlMemory) {
             builder.windowSize(windowSize);
         }
 
-        MLMemorySpec mlMemorySpec = builder.build();
-        return mlMemorySpec;
-
+        return builder.build();
     }
-
 }
diff --git a/src/main/resources/mappings/workflow-steps.json b/src/main/resources/mappings/workflow-steps.json
@@ -137,11 +137,7 @@
     "register_agent": {
         "inputs":[
             "name",
-            "type",
-            "tools",
-            "parameters",
-            "memory",
-            "app_type"
+            "type"
         ],
         "outputs":[
             "agent_id"