diff --git a/src/main/java/org/opensearch/agent/tools/PPLTool.java b/src/main/java/org/opensearch/agent/tools/PPLTool.java index 09fe1d3a..5aaf871f 100644 --- 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 void run(Map parameters, ActionListener listener) Map dataAsMap = (Map) modelTensor.getDataAsMap(); String ppl = parseOutput(dataAsMap.get("response"), indexName); if (!this.execute) { - listener.onResponse((T) ppl); + Map ret = ImmutableMap.of("ppl", ppl); + listener.onResponse((T) AccessController.doPrivileged((PrivilegedExceptionAction) () -> 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 index 129c2411..8ed605a8 100644 --- 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.wrap(executePPLResult -> { - assertEquals("source=demo| head 1", executePPLResult); + Map ret = gson.fromJson(executePPLResult, Map.class); + assertEquals("source=demo| head 1", ret.get("ppl")); }, e -> { log.info(e); })); }