Add agent framework configuration to make IT pass (opensearch-project…

…#179) (opensearch-project#180) * Add agent framework configuration to make IT pass * Fix failure UT since the response data structure change --------- (cherry picked from commit 940ac32) Signed-off-by: zane-neo <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Signed-off-by: yuye-aws <[email protected]>
yuye-aws · Apr 26, 2024 · 4f27764 · 4f27764
1 parent ad44ab9
commit 4f27764
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git a/src/test/java/org/opensearch/integTest/BaseAgentToolsIT.java b/src/test/java/org/opensearch/integTest/BaseAgentToolsIT.java
@@ -61,6 +61,7 @@ public void updateClusterSettings() {
         updateClusterSettings("plugins.ml_commons.native_memory_threshold", 100);
         updateClusterSettings("plugins.ml_commons.jvm_heap_memory_threshold", 100);
         updateClusterSettings("plugins.ml_commons.allow_registering_model_via_url", true);
+        updateClusterSettings("plugins.ml_commons.agent_framework_enabled", true);
     }
 
     @SneakyThrows

diff --git a/src/test/java/org/opensearch/integTest/VisualizationsToolIT.java b/src/test/java/org/opensearch/integTest/VisualizationsToolIT.java
@@ -16,6 +16,8 @@
 import org.opensearch.client.Response;
 import org.opensearch.core.rest.RestStatus;
 
+import com.google.gson.JsonArray;
+import com.google.gson.JsonElement;
 import com.google.gson.JsonParser;
 
 import lombok.extern.log4j.Log4j2;
@@ -86,22 +88,27 @@ private void prepareVisualization(String title, String id) {
     }
 
     private String extractAdditionalInfo(String responseStr) {
-        return JsonParser
+        JsonArray output = JsonParser
             .parseString(responseStr)
             .getAsJsonObject()
             .get("inference_results")
             .getAsJsonArray()
             .get(0)
             .getAsJsonObject()
             .get("output")
-            .getAsJsonArray()
-            .get(0)
-            .getAsJsonObject()
-            .get("dataAsMap")
-            .getAsJsonObject()
-            .get("additional_info")
-            .getAsJsonObject()
-            .get(String.format(Locale.ROOT, "%s.output", toolType()))
-            .getAsString();
+            .getAsJsonArray();
+        for (JsonElement element : output) {
+            if ("response".equals(element.getAsJsonObject().get("name").getAsString())) {
+                return element
+                    .getAsJsonObject()
+                    .get("dataAsMap")
+                    .getAsJsonObject()
+                    .get("additional_info")
+                    .getAsJsonObject()
+                    .get(String.format(Locale.ROOT, "%s.output", toolType()))
+                    .getAsString();
+            }
+        }
+        return null;
     }
 }