Includ workflow id and current node id in the exception message (open…

…search-project#262) Includ workflow id and current node id in the exception message during registe agent step Signed-off-by: Jackie Han <[email protected]>
dbwiddis · Dec 15, 2023 · 4123229 · 4123229
1 parent 5682c22
commit 4123229
Showing 1 changed file with 14 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
@@ -83,6 +83,8 @@ public CompletableFuture<WorkflowData> execute(
         Map<String, String> previousNodeInputs
     ) throws IOException {
 
+        String workflowId = currentNodeInputs.getWorkflowId();
+
         CompletableFuture<WorkflowData> registerAgentModelFuture = new CompletableFuture<>();
 
         ActionListener<MLRegisterAgentResponse> actionListener = new ActionListener<>() {
@@ -92,7 +94,7 @@ public void onResponse(MLRegisterAgentResponse mlRegisterAgentResponse) {
                     String resourceName = WorkflowResources.getResourceByWorkflowStep(getName());
                     logger.info("Agent registration successful for the agent {}", mlRegisterAgentResponse.getAgentId());
                     flowFrameworkIndicesHandler.updateResourceInStateIndex(
-                        currentNodeInputs.getWorkflowId(),
+                        workflowId,
                         currentNodeId,
                         getName(),
                         mlRegisterAgentResponse.getAgentId(),
@@ -101,7 +103,7 @@ public void onResponse(MLRegisterAgentResponse mlRegisterAgentResponse) {
                             registerAgentModelFuture.complete(
                                 new WorkflowData(
                                     Map.ofEntries(Map.entry(resourceName, mlRegisterAgentResponse.getAgentId())),
-                                    currentNodeInputs.getWorkflowId(),
+                                    workflowId,
                                     currentNodeId
                                 )
                             );
@@ -168,12 +170,15 @@ public void onFailure(Exception e) {
             // Case when modelId is not present at all
             if (llmModelId == null) {
                 registerAgentModelFuture.completeExceptionally(
-                    new FlowFrameworkException("llm model id is not provided", RestStatus.BAD_REQUEST)
+                    new FlowFrameworkException(
+                        "llm model id is not provided for workflow: " + workflowId + " on node: " + currentNodeId,
+                        RestStatus.BAD_REQUEST
+                    )
                 );
                 return registerAgentModelFuture;
             }
 
-            LLMSpec llmSpec = getLLMSpec(llmModelId, llmParameters);
+            LLMSpec llmSpec = getLLMSpec(llmModelId, llmParameters, workflowId, currentNodeId);
 
             MLAgentBuilder builder = MLAgent.builder().name(name);
 
@@ -246,9 +251,12 @@ private String getLlmModelId(Map<String, String> previousNodeInputs, Map<String,
         return llmModelId;
     }
 
-    private LLMSpec getLLMSpec(String llmModelId, Map<String, String> llmParameters) {
+    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", RestStatus.BAD_REQUEST);
+            throw new FlowFrameworkException(
+                "model id for llm is null for workflow: " + workflowId + " on node: " + currentNodeId,
+                RestStatus.BAD_REQUEST
+            );
         }
         LLMSpec.LLMSpecBuilder builder = LLMSpec.builder();
         builder.modelId(llmModelId);