Making memory and llm spec optional fields

Signed-off-by: Joshua Palis <jpalis@amazon.com>
opensearch-project · Jan 10, 2024 · 062745b · 062745b
1 parent 815a74e
commit 062745b
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java
@@ -194,12 +194,16 @@ public void onFailure(Exception e) {
             if (description != null) {
                 builder.description(description);
             }
+            if (memory != null) {
+                builder.memory(memory);
+            }
+            if (llmSpec != null) {
+                builder.llm(llmSpec);
+            }
 
             builder.type(type)
-                .llm(llmSpec)
                 .tools(toolsList)
                 .parameters(parameters)
-                .memory(memory)
                 .createdTime(createdTime)
                 .lastUpdateTime(lastUpdateTime)
                 .appType(appType);
@@ -264,10 +268,7 @@ private String getLlmModelId(Map<String, String> previousNodeInputs, Map<String,
 
     private LLMSpec getLLMSpec(String llmModelId, Map<String, String> llmParameters, String workflowId, String currentNodeId) {
         if (llmModelId == null) {
-            throw new FlowFrameworkException(
-                "model id for llm is null for workflow: " + workflowId + " on node: " + currentNodeId,
-                RestStatus.BAD_REQUEST
-            );
+            return null;
         }
         LLMSpec.LLMSpecBuilder builder = LLMSpec.builder();
         builder.modelId(llmModelId);
@@ -279,6 +280,9 @@ private LLMSpec getLLMSpec(String llmModelId, Map<String, String> llmParameters,
     }
 
     private MLMemorySpec getMLMemorySpec(Object mlMemory) {
+        if (mlMemory == null) {
+            return null;
+        }
 
         Map<?, ?> map = (Map<?, ?>) mlMemory;
         String type = null;