fix UT bug (#172) (#175)

(cherry picked from commit 6070c65) Signed-off-by: xinyual <[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>
opensearch-project · Feb 2, 2024 · e3d6711 · e3d6711
1 parent bc56d78
commit e3d6711
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/src/main/java/org/opensearch/agent/tools/PPLTool.java b/src/main/java/org/opensearch/agent/tools/PPLTool.java
@@ -173,7 +173,8 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
                     Map<String, String> dataAsMap = (Map<String, String>) modelTensor.getDataAsMap();
                     String ppl = parseOutput(dataAsMap.get("response"), indexName);
                     if (!this.execute) {
-                        listener.onResponse((T) ppl);
+                        Map<String, String> ret = ImmutableMap.of("ppl", ppl);
+                        listener.onResponse((T) AccessController.doPrivileged((PrivilegedExceptionAction<String>) () -> gson.toJson(ret)));
                         return;
                     }
                     JSONObject jsonContent = new JSONObject(ImmutableMap.of("query", ppl));

diff --git a/src/test/java/org/opensearch/agent/tools/PPLToolTests.java b/src/test/java/org/opensearch/agent/tools/PPLToolTests.java
@@ -142,11 +142,12 @@ public void testTool() {
     public void testTool_with_WithoutExecution() {
         PPLTool tool = PPLTool.Factory
             .getInstance()
-            .create(ImmutableMap.of("model_id", "modelId", "model_type", "claude", "execute", false));
+            .create(ImmutableMap.of("model_id", "modelId", "model_type", "claude", "execute", "false"));
         assertEquals(PPLTool.TYPE, tool.getName());
 
         tool.run(ImmutableMap.of("index", "demo", "question", "demo"), ActionListener.<String>wrap(executePPLResult -> {
-            assertEquals("source=demo| head 1", executePPLResult);
+            Map<String, String> ret = gson.fromJson(executePPLResult, Map.class);
+            assertEquals("source=demo| head 1", ret.get("ppl"));
         }, e -> { log.info(e); }));
 
     }