add truncate For PPL Tool (#206) (#232)

* add truncate * use head instead * fix typo * use head instead * fix UT --------- (cherry picked from commit 209bb15) 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 16, 2024 · 7546f72 · 7546f72
1 parent 04c6616
commit 7546f72
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 1 deletion.
diff --git a/src/main/java/org/opensearch/agent/tools/PPLTool.java b/src/main/java/org/opensearch/agent/tools/PPLTool.java
@@ -94,6 +94,8 @@ public class PPLTool implements Tool {
 
     private String previousToolKey;
 
+    private int head;
+
     private static Gson gson = new Gson();
 
     private static Map<String, String> DEFAULT_PROMPT_DICT;
@@ -147,7 +149,15 @@ public static PPLModelType from(String value) {
 
     }
 
-    public PPLTool(Client client, String modelId, String contextPrompt, String pplModelType, String previousToolKey, boolean execute) {
+    public PPLTool(
+        Client client,
+        String modelId,
+        String contextPrompt,
+        String pplModelType,
+        String previousToolKey,
+        int head,
+        boolean execute
+    ) {
         this.client = client;
         this.modelId = modelId;
         this.pplModelType = PPLModelType.from(pplModelType);
@@ -157,6 +167,7 @@ public PPLTool(Client client, String modelId, String contextPrompt, String pplMo
             this.contextPrompt = contextPrompt;
         }
         this.previousToolKey = previousToolKey;
+        this.head = head;
         this.execute = execute;
     }
 
@@ -307,6 +318,7 @@ public PPLTool create(Map<String, Object> map) {
                 (String) map.getOrDefault("prompt", ""),
                 (String) map.getOrDefault("model_type", ""),
                 (String) map.getOrDefault("previous_tool_name", ""),
+                Integer.valueOf((String) map.getOrDefault("head", "-1")),
                 Boolean.valueOf((String) map.getOrDefault("execute", "true"))
             );
         }
@@ -487,6 +499,14 @@ private String parseOutput(String llmOutput, String indexName) {
         }
         ppl = ppl.replace("`", "");
         ppl = ppl.replaceAll("\\bSPAN\\(", "span(");
+        if (this.head > 0) {
+            String[] lists = llmOutput.split("\\|");
+            String lastCommand = lists[lists.length - 1].strip();
+            if (!lastCommand.toLowerCase(Locale.ROOT).startsWith("head")) // not handle cases source=...| ... | head 5 | head <head>
+            {
+                ppl = ppl + " | head " + this.head;
+            }
+        }
         return ppl;
     }
 

diff --git a/src/test/java/org/opensearch/agent/tools/PPLToolTests.java b/src/test/java/org/opensearch/agent/tools/PPLToolTests.java
@@ -153,6 +153,36 @@ public void testTool_withPreviousInput() {
 
     }
 
+    @Test
+    public void testTool_withHEADButIgnore() {
+        PPLTool tool = PPLTool.Factory.getInstance().create(ImmutableMap.of("model_id", "modelId", "prompt", "contextPrompt", "head", "5"));
+        assertEquals(PPLTool.TYPE, tool.getName());
+
+        tool.run(ImmutableMap.of("index", "demo", "question", "demo"), ActionListener.<String>wrap(executePPLResult -> {
+            Map<String, String> returnResults = gson.fromJson(executePPLResult, Map.class);
+            assertEquals("ppl result", returnResults.get("executionResult"));
+            assertEquals("source=demo| head 1", returnResults.get("ppl"));
+        }, e -> { log.info(e); }));
+
+    }
+
+    @Test
+    public void testTool_withHEAD() {
+        pplReturns = Collections.singletonMap("response", "source=demo");
+        modelTensor = new ModelTensor("tensor", new Number[0], new long[0], MLResultDataType.STRING, null, null, pplReturns);
+        initMLTensors();
+
+        PPLTool tool = PPLTool.Factory.getInstance().create(ImmutableMap.of("model_id", "modelId", "prompt", "contextPrompt", "head", "5"));
+        assertEquals(PPLTool.TYPE, tool.getName());
+
+        tool.run(ImmutableMap.of("index", "demo", "question", "demo"), ActionListener.<String>wrap(executePPLResult -> {
+            Map<String, String> returnResults = gson.fromJson(executePPLResult, Map.class);
+            assertEquals("ppl result", returnResults.get("executionResult"));
+            assertEquals("source=demo | head 5", returnResults.get("ppl"));
+        }, e -> { log.info(e); }));
+
+    }
+
     @Test
     public void testTool_with_WithoutExecution() {
         PPLTool tool = PPLTool.Factory