Addressing PR comments

Signed-off-by: Owais Kazi <[email protected]>
opensearch-project · Nov 29, 2023 · 18e331e · 18e331e
1 parent 7243fab
commit 18e331e
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 9 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/model/WorkflowNode.java b/src/main/java/org/opensearch/flowframework/model/WorkflowNode.java
@@ -26,9 +26,7 @@
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.opensearch.flowframework.common.CommonValue.LLM_FIELD;
-import static org.opensearch.flowframework.util.ParseUtils.buildLLMMap;
-import static org.opensearch.flowframework.util.ParseUtils.buildStringToStringMap;
-import static org.opensearch.flowframework.util.ParseUtils.parseStringToStringMap;
+import static org.opensearch.flowframework.util.ParseUtils.*;
 
 /**
  * This represents a process node (step) in a workflow graph in the {@link Template}.
@@ -100,7 +98,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
                     }
                 }
                 xContentBuilder.endArray();
-            } else if (e.getValue() instanceof Object) {
+            } else if (e.getValue() instanceof LLMSpec) {
                 if (LLM_FIELD.equals(e.getKey())) {
                     xContentBuilder.startObject();
                     buildLLMMap(xContentBuilder, (LLMSpec) e.getValue());
@@ -150,7 +148,7 @@ public static WorkflowNode parse(XContentParser parser) throws IOException {
                                 break;
                             case START_OBJECT:
                                 if (LLM_FIELD.equals(inputFieldName)) {
-                                    userInputs.put(inputFieldName, LLMSpec.parse(parser));
+                                    userInputs.put(inputFieldName, parseLLM(parser));
                                 } else {
                                     userInputs.put(inputFieldName, parseStringToStringMap(parser));
                                 }

diff --git a/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java b/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java
@@ -95,7 +95,7 @@ public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) throws I
 
             MLToolSpec mlToolSpec = builder.build();
 
-            toolFuture.complete(new WorkflowData(Map.ofEntries(Map.entry(TOOLS_FIELD, mlToolSpec))));
+            toolFuture.complete(new WorkflowData(Map.ofEntries(Map.entry(TOOLS_FIELD, mlToolSpec)), data.get(0).getWorkflowId()));
         }
 
         logger.info("Tool registered successfully {}", type);

diff --git a/src/main/resources/mappings/workflow-steps.json b/src/main/resources/mappings/workflow-steps.json
@@ -109,7 +109,7 @@
             "agent_id"
         ]
     },
-    "tool": {
+    "create_tool": {
         "inputs": [
             "type"
         ],

diff --git a/src/test/java/org/opensearch/flowframework/workflow/RegisterAgentTests.java b/src/test/java/org/opensearch/flowframework/workflow/RegisterAgentTests.java
@@ -68,7 +68,8 @@ public void setUp() throws Exception {
                 Map.entry("created_time", 1689793598499L),
                 Map.entry("last_updated_time", 1689793598499L),
                 Map.entry("app_type", "app")
-            )
+            ),
+            "test-id"
         );
     }
 

diff --git a/src/test/java/org/opensearch/flowframework/workflow/ToolStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/ToolStepTests.java
@@ -32,7 +32,8 @@ public void setUp() throws Exception {
                 Map.entry("description", "description"),
                 Map.entry("parameters", Collections.emptyMap()),
                 Map.entry("include_output_in_agent_response", false)
-            )
+            ),
+            "test-id"
         );
     }