fix agent index mapping; fix flow agent

Signed-off-by: Yaliang Wu <[email protected]>
b4sjoo · Nov 13, 2023 · 283dd12 · 283dd12
1 parent affd351
commit 283dd12
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/CommonValue.java b/common/src/main/java/org/opensearch/ml/common/CommonValue.java
@@ -359,6 +359,7 @@ public class CommonValue {
             + MLAgent.MEMORY_FIELD
             + "\" : {\"type\": \"flat_object\"},\n"
             + "      \""
+            + MLAgent.CREATED_TIME_FIELD
             + "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"},\n"
             + "      \""
             + MLAgent.LAST_UPDATED_TIME_FIELD

diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java
@@ -76,7 +76,7 @@ public void run(MLAgent mlAgent, Map<String, String> params, ActionListener<Obje
                 StepListener<Object> nextStepListener = new StepListener<>();
                 int finalI = i;
                 lastStepListener.whenComplete(output -> {
-                    String outputKey = lastToolSpec.getName() + ".output";
+                    String outputKey = lastToolSpec.getType() + ".output";
                     if (lastToolSpec.getName() !=  null) {
                         outputKey = lastToolSpec.getName() + ".output";
                     }
@@ -133,11 +133,13 @@ private Tool createTool(MLToolSpec toolSpec) {
         if (toolSpec.getParameters() != null) {
             toolParams.putAll(toolSpec.getParameters());
         }
-        if (!toolFactories.containsKey(toolSpec.getName())) {
-            throw new IllegalArgumentException("Tool not found: " + toolSpec.getName());
+        if (!toolFactories.containsKey(toolSpec.getType())) {
+            throw new IllegalArgumentException("Tool not found: " + toolSpec.getType());
+        }
+        Tool tool = toolFactories.get(toolSpec.getType()).create(toolParams);
+        if (toolSpec.getName() != null) {
+            tool.setName(toolSpec.getName());
         }
-        Tool tool = toolFactories.get(toolSpec.getName()).create(toolParams);
-        tool.setName(toolSpec.getName());
 
         if (toolSpec.getDescription() != null) {
             tool.setDescription(toolSpec.getDescription());
@@ -152,8 +154,8 @@ private Map<String, String> getToolExecuteParams(MLToolSpec toolSpec, Map<String
         }
         for (String key : params.keySet()) {
             String toBeReplaced = null;
-            if (key.startsWith(toolSpec.getName() + ".")) {
-                toBeReplaced = toolSpec.getName()+".";
+            if (key.startsWith(toolSpec.getType() + ".")) {
+                toBeReplaced = toolSpec.getType() + ".";
             }
             if (toolSpec.getName() != null && key.startsWith(toolSpec.getName() + ".")) {
                 toBeReplaced = toolSpec.getName()+".";