From 41606a84b0ee20ba444eb91d03c3ff3f66817f62 Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Sun, 4 Feb 2024 10:38:59 +0800
Subject: [PATCH] Add agent framework configuration to make IT pass (#179)
 (#180)

* Add agent framework configuration to make IT pass



* Fix failure UT since the response data structure change



---------


(cherry picked from commit 940ac32f13b1c5141ee4a05f458ad4e79d9efc02)

Signed-off-by: zane-neo <zaniu@amazon.com>
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>
---
 .../integTest/BaseAgentToolsIT.java           |  1 +
 .../integTest/VisualizationsToolIT.java       | 27 ++++++++++++-------
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/src/test/java/org/opensearch/integTest/BaseAgentToolsIT.java b/src/test/java/org/opensearch/integTest/BaseAgentToolsIT.java
index 2fdbe6ab..cc7b2702 100644
--- 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
index e7f54521..2bf0e611 100644
--- 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,7 +88,7 @@ private void prepareVisualization(String title, String id) {
     }
 
     private String extractAdditionalInfo(String responseStr) {
-        return JsonParser
+        JsonArray output = JsonParser
             .parseString(responseStr)
             .getAsJsonObject()
             .get("inference_results")
@@ -94,14 +96,19 @@ private String extractAdditionalInfo(String responseStr) {
             .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;
     }
 }