add execute field (opensearch-project#146) (opensearch-project#147)

* add execute field * apply spotless --------- (cherry picked from commit 16a26ce) 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> Signed-off-by: yuye-aws <[email protected]>
yuye-aws · Apr 26, 2024 · c0c9e38 · c0c9e38
1 parent f1945a9
commit c0c9e38
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/src/main/java/org/opensearch/agent/tools/PPLTool.java b/src/main/java/org/opensearch/agent/tools/PPLTool.java
@@ -87,6 +87,8 @@ public class PPLTool implements Tool {
 
     private String contextPrompt;
 
+    private Boolean execute;
+
     private PPLModelType pplModelType;
 
     private static Gson gson = new Gson();
@@ -120,7 +122,7 @@ public static PPLModelType from(String value) {
 
     }
 
-    public PPLTool(Client client, String modelId, String contextPrompt, String pplModelType) {
+    public PPLTool(Client client, String modelId, String contextPrompt, String pplModelType, boolean execute) {
         this.client = client;
         this.modelId = modelId;
         this.pplModelType = PPLModelType.from(pplModelType);
@@ -129,6 +131,7 @@ public PPLTool(Client client, String modelId, String contextPrompt, String pplMo
         } else {
             this.contextPrompt = contextPrompt;
         }
+        this.execute = execute;
     }
 
     @Override
@@ -169,6 +172,10 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
                     ModelTensor modelTensor = modelTensors.getMlModelTensors().get(0);
                     Map<String, String> dataAsMap = (Map<String, String>) modelTensor.getDataAsMap();
                     String ppl = parseOutput(dataAsMap.get("response"), indexName);
+                    if (!this.execute) {
+                        listener.onResponse((T) ppl);
+                        return;
+                    }
                     JSONObject jsonContent = new JSONObject(ImmutableMap.of("query", ppl));
                     PPLQueryRequest pplQueryRequest = new PPLQueryRequest(ppl, jsonContent, null, "jdbc");
                     TransportPPLQueryRequest transportPPLQueryRequest = new TransportPPLQueryRequest(pplQueryRequest);
@@ -253,7 +260,8 @@ public PPLTool create(Map<String, Object> map) {
                 client,
                 (String) map.get("model_id"),
                 (String) map.getOrDefault("prompt", ""),
-                (String) map.getOrDefault("model_type", "")
+                (String) map.getOrDefault("model_type", ""),
+                (boolean) map.getOrDefault("execute", true)
             );
         }
 

diff --git a/src/test/java/org/opensearch/agent/tools/PPLToolTests.java b/src/test/java/org/opensearch/agent/tools/PPLToolTests.java
@@ -138,6 +138,19 @@ public void testTool() {
 
     }
 
+    @Test
+    public void testTool_with_WithoutExecution() {
+        PPLTool tool = PPLTool.Factory
+            .getInstance()
+            .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);
+        }, e -> { log.info(e); }));
+
+    }
+
     @Test
     public void testTool_with_DefaultPrompt() {
         PPLTool tool = PPLTool.Factory.getInstance().create(ImmutableMap.of("model_id", "modelId", "model_type", "claude"));