Making memory and llm spec optional fields (#398)

* Making memory and llm spec optional fields Signed-off-by: Joshua Palis <[email protected]> * Removing second null llmModelId check Signed-off-by: Joshua Palis <[email protected]> --------- Signed-off-by: Joshua Palis <[email protected]> (cherry picked from commit f8e822f) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Jan 11, 2024 · f1292a1 · f1292a1
1 parent a37e83b
commit f1292a1
Showing 1 changed file with 10 additions and 18 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterAgentStep.java
@@ -13,7 +13,6 @@
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.common.Nullable;
 import org.opensearch.core.action.ActionListener;
-import org.opensearch.core.rest.RestStatus;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.indices.FlowFrameworkIndicesHandler;
 import org.opensearch.flowframework.util.ParseUtils;
@@ -176,30 +175,23 @@ public void onFailure(Exception e) {
                 llmModelId = getLlmModelId(previousNodeInputs, outputs);
             }
 
-            // Case when modelId is not present at all
-            if (llmModelId == null) {
-                registerAgentModelFuture.completeExceptionally(
-                    new FlowFrameworkException(
-                        "llm model id is not provided for workflow: " + workflowId + " on node: " + currentNodeId,
-                        RestStatus.BAD_REQUEST
-                    )
-                );
-                return registerAgentModelFuture;
-            }
-
             LLMSpec llmSpec = getLLMSpec(llmModelId, llmParameters, workflowId, currentNodeId);
 
             MLAgentBuilder builder = MLAgent.builder().name(name);
 
             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 +256,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 +268,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;