Log pattern tool improvement

src/main/java/org/opensearch/agent/tools/PPLTool.java

@@ -35,7 +35,6 @@
 import org.opensearch.client.Client;
 import org.opensearch.cluster.metadata.MappingMetadata;
 import org.opensearch.core.action.ActionListener;
-import org.opensearch.core.action.ActionResponse;
 import org.opensearch.index.query.MatchAllQueryBuilder;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
@@ -51,7 +50,6 @@
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.sql.plugin.transport.PPLQueryAction;
 import org.opensearch.sql.plugin.transport.TransportPPLQueryRequest;
-import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse;
 import org.opensearch.sql.ppl.domain.PPLQueryRequest;
 
 import com.google.common.collect.ImmutableMap;
@@ -233,7 +231,7 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
                         .execute(
                             PPLQueryAction.INSTANCE,
                             transportPPLQueryRequest,
-                            getPPLTransportActionListener(ActionListener.wrap(transportPPLQueryResponse -> {
+                            ToolHelper.getPPLTransportActionListener(ActionListener.wrap(transportPPLQueryResponse -> {
                                 String results = transportPPLQueryResponse.getResult();
                                 Map<String, String> returnResults = ImmutableMap.of("ppl", ppl, "executionResult", results);
                                 listener
@@ -444,10 +442,6 @@ private static void extractSamples(Map<String, Object> sampleSource, Map<String,
         }
     }
 
-    private <T extends ActionResponse> ActionListener<T> getPPLTransportActionListener(ActionListener<TransportPPLQueryResponse> listener) {
-        return ActionListener.wrap(r -> { listener.onResponse(TransportPPLQueryResponse.fromActionResponse(r)); }, listener::onFailure);
-    }
-
     @SuppressWarnings("unchecked")
     private void extractFromChatParameters(Map<String, String> parameters) {
         if (parameters.containsKey("input")) {

src/main/java/org/opensearch/agent/tools/utils/ToolHelper.java

@@ -11,7 +11,10 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.action.ActionResponse;
 import org.opensearch.ml.common.utils.StringUtils;
+import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse;
 
 import lombok.extern.log4j.Log4j2;
 
@@ -75,4 +78,15 @@ public static void extractFieldNamesTypes(
             }
         }
     }
+
+    /**
+     * Wrapper to get PPL transport action listener
+     * @param listener input action listener
+     * @return wrapped action listener
+     */
+    public static <T extends ActionResponse> ActionListener<T> getPPLTransportActionListener(
+        ActionListener<TransportPPLQueryResponse> listener
+    ) {
+        return ActionListener.wrap(r -> { listener.onResponse(TransportPPLQueryResponse.fromActionResponse(r)); }, listener::onFailure);
+    }
 }

src/test/java/org/opensearch/agent/tools/LogPatternToolTests.java

@@ -13,24 +13,27 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.opensearch.agent.tools.AbstractRetrieverTool.DOC_SIZE_FIELD;
-import static org.opensearch.agent.tools.LogPatternTool.PATTERN;
+import static org.opensearch.agent.tools.AbstractRetrieverTool.INDEX_FIELD;
+import static org.opensearch.agent.tools.AbstractRetrieverTool.INPUT_FIELD;
+import static org.opensearch.agent.tools.LogPatternTool.PPL_FIELD;
 import static org.opensearch.agent.tools.LogPatternTool.SAMPLE_LOG_SIZE;
 import static org.opensearch.agent.tools.LogPatternTool.TOP_N_PATTERN;
-import static org.opensearch.integTest.BaseAgentToolsIT.gson;
+import static org.opensearch.ml.common.utils.StringUtils.gson;
 
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.regex.Pattern;
 
 import org.hamcrest.MatcherAssert;
+import org.json.JSONObject;
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mock;
@@ -43,6 +46,8 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.SearchHits;
+import org.opensearch.sql.plugin.transport.PPLQueryAction;
+import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.gson.JsonElement;
@@ -52,19 +57,26 @@
 public class LogPatternToolTests {
 
     public static String responseBodyResourceFile = "org/opensearch/agent/tools/expected_flow_agent_of_log_pattern_tool_response_body.json";
-    public static final String TEST_QUERY_TEXT = "123fsd23134sdfouh";
+    public static final String TEST_QUERY_TEXT =
+        """
+            {"size":2,"query":{"bool":{"filter":[{"range":{"timestamp":{"from":"1734404246000||-1d","to":"1734404246000","include_lower":true,"include_upper":true,"format":"epoch_millis","boost":1}}},{"range":{"bytes":{"from":0,"to":null,"include_lower":true,"include_upper":true,"boost":1}}}],"adjust_pure_negative":true,"boost":1}}}""";
     private Map<String, Object> params = new HashMap<>();
     private final Client client = mock(Client.class);
     @Mock
     private SearchResponse searchResponse;
     @Mock
     private SearchHits searchHits;
+    @Mock
+    private TransportPPLQueryResponse pplQueryResponse;
 
     @SneakyThrows
     @Before
     public void setup() {
         MockitoAnnotations.openMocks(this);
         LogPatternTool.Factory.getInstance().init(client, null);
+    }
+
+    private void mockDSLInvocation() throws IOException {
         List<String> fields = List.of("field1", "field2", "field3");
         SearchHit[] hits = new SearchHit[] {
             createHit(0, null, fields, List.of("123", "123.abc-AB * De /", 12345)),
@@ -94,14 +106,26 @@ private BytesReference createSource(List<String> fieldNames, List<Object> fieldC
         return (BytesReference.bytes(builder));
     }
 
+    private void mockPPLInvocation() throws IOException {
+        String pplRawResponse =
+            """
+                {"schema":[{"name":"field1","type":"string"},{"name":"field2","type":"string"},{"name":"field3","type":"long"}],"datarows":[["123","123.abc-AB * De /",12345],["123","45.abc-AB * De /",12345],["123","12.abc_AB * De /",12345],["123","45.ab_AB * De /",12345],["123",".abAB * De /",12345]],"total":5,"size":5}
+                """;
+        doAnswer(invocation -> {
+            ActionListener<TransportPPLQueryResponse> listener = (ActionListener<TransportPPLQueryResponse>) invocation.getArguments()[2];
+            listener.onResponse(pplQueryResponse);
+            return null;
+        }).when(client).execute(eq(PPLQueryAction.INSTANCE), any(), any());
+        when(pplQueryResponse.getResult()).thenReturn(pplRawResponse);
+    }
+
     @Test
     @SneakyThrows
     public void testCreateTool() {
         LogPatternTool tool = LogPatternTool.Factory.getInstance().create(params);
         assertEquals(LogPatternTool.LOG_PATTERN_DEFAULT_DOC_SIZE, (int) tool.docSize);
         assertEquals(LogPatternTool.DEFAULT_TOP_N_PATTERN, tool.getTopNPattern());
         assertEquals(LogPatternTool.DEFAULT_SAMPLE_LOG_SIZE, tool.getSampleLogSize());
-        assertNull(tool.getPattern());
         assertEquals("LogPatternTool", tool.getType());
         assertEquals("LogPatternTool", tool.getName());
         assertEquals(LogPatternTool.DEFAULT_DESCRIPTION, LogPatternTool.Factory.getInstance().getDefaultDescription());
@@ -160,36 +184,31 @@ public void testCreateToolWithNonPositiveSize() {
     @Test
     public void testGetQueryBody() {
         LogPatternTool tool = LogPatternTool.Factory.getInstance().create(params);
-        assertEquals(TEST_QUERY_TEXT, tool.getQueryBody(TEST_QUERY_TEXT));
+        assertEquals(new JSONObject(TEST_QUERY_TEXT).toString(), new JSONObject(tool.getQueryBody(TEST_QUERY_TEXT)).toString());
     }
 
     @Test
     public void testValidate() {
         LogPatternTool tool = LogPatternTool.Factory.getInstance().create(params);
-        assertTrue(tool.validate(Map.of("index", "test1", "input", "input_value")));
+        assertTrue(tool.validate(Map.of(INDEX_FIELD, "test1", INPUT_FIELD, "input_value")));
+        assertTrue(tool.validate(Map.of(INDEX_FIELD, "test1", PPL_FIELD, "ppl_value")));
 
-        // validate failure if no index
-        assertFalse(tool.validate(Map.of("input", "input_value")));
+        // validate failure if no input or ppl
+        assertFalse(tool.validate(Map.of(INDEX_FIELD, "test1")));
 
-        // validate failure if no
-        assertFalse(tool.validate(Map.of("index", "test1")));
+        // validate failure if no index
+        assertFalse(tool.validate(Map.of(INPUT_FIELD, "input_value")));
     }
 
     @Test
     public void testFindLongestField() {
         assertEquals("field2", LogPatternTool.findLongestField(Map.of("field1", "123", "field2", "1234", "filed3", 1234)));
     }
 
-    @Test
-    public void testExtractPattern() {
-        assertEquals("././", LogPatternTool.extractPattern("123.abc/.AB/", null));
-        assertEquals("123.c/.AB/", LogPatternTool.extractPattern("123.abc/.AB/", Pattern.compile("ab")));
-        assertEquals(".abc/.AB/", LogPatternTool.extractPattern("123.abc/.AB/", Pattern.compile("[0-9]")));
-    }
-
     @SneakyThrows
     @Test
     public void testExecutionDefault() {
+        mockDSLInvocation();
         LogPatternTool tool = LogPatternTool.Factory.getInstance().create(params);
         JsonElement expected = gson
             .fromJson(
@@ -210,12 +229,10 @@ public void testExecutionDefault() {
     @SneakyThrows
     @Test
     public void testExecutionWithSpecifiedPatternField() {
+        mockDSLInvocation();
         LogPatternTool tool = LogPatternTool.Factory.getInstance().create(params);
         JsonElement expected = gson
-            .fromJson(
-                "[{\"total count\":5,\"sample logs\":[{\"field1\":\"123\",\"field3\":12345,\"field2\":\"123.abc-AB * De /\"},{\"field1\":\"123\",\"field3\":12345,\"field2\":\"45.abc-AB * De /\"}],\"pattern\":\"\"}]",
-                JsonElement.class
-            );
+            .fromJson("[{\"pattern\":\"<*>\",\"total count\":5,\"sample logs\":[\"123\",\"123\"]}]", JsonElement.class);
         tool
             .run(
                 ImmutableMap.of("index", "index_name", "input", "{}", "pattern_field", "field1", "sample_log_size", "2"),
@@ -227,26 +244,6 @@ public void testExecutionWithSpecifiedPatternField() {
             );
     }
 
-    @SneakyThrows
-    @Test
-    public void testExecutionWithSpecifiedPattern() {
-        LogPatternTool tool = LogPatternTool.Factory.getInstance().create(Map.of(PATTERN, "[a-zA-Z]"));
-        JsonElement expected = gson
-            .fromJson(
-                "[{\"pattern\":\"45.- *  /\",\"sample logs\":[{\"field1\":\"123\",\"field3\":12345,\"field2\":\"45.abc-AB * De /\"}],\"total count\":1},{\"pattern\":\". *  /\",\"sample logs\":[{\"field1\":\"123\",\"field3\":12345,\"field2\":\".abAB * De /\"}],\"total count\":1},{\"pattern\":\"123.- *  /\",\"sample logs\":[{\"field1\":\"123\",\"field3\":12345,\"field2\":\"123.abc-AB * De /\"}],\"total count\":1}]",
-                JsonElement.class
-            );
-        tool
-            .run(
-                ImmutableMap.of("index", "index_name", "input", "{}"),
-                ActionListener
-                    .<String>wrap(
-                        response -> assertEquals(expected, gson.fromJson(response, JsonElement.class)),
-                        e -> fail("Tool runs failed: " + e.getMessage())
-                    )
-            );
-    }
-
     @SneakyThrows
     @Test
     public void testExecutionWithBlankInput() {
@@ -257,7 +254,8 @@ public void testExecutionWithBlankInput() {
                 ActionListener
                     .<String>wrap(
                         response -> fail(),
-                        e -> MatcherAssert.assertThat(e.getMessage(), containsString("[input] is null or empty, can not process it."))
+                        e -> MatcherAssert
+                            .assertThat(e.getMessage(), containsString("Both DSL and PPL input is null or empty, can not process it."))
                     )
             );
     }
@@ -375,4 +373,69 @@ public void testExecutionFailedInSearch() {
                 ActionListener.<String>wrap(response -> fail(), e -> assertEquals("Failed in Search", e.getMessage()))
             );
     }
+
+    @SneakyThrows
+    @Test
+    public void testExecutionWithPPLInput() {
+        mockPPLInvocation();
+        LogPatternTool tool = LogPatternTool.Factory.getInstance().create(params);
+        JsonElement expected = gson
+            .fromJson(
+                Files.readString(Path.of(this.getClass().getClassLoader().getResource(responseBodyResourceFile).toURI())),
+                JsonElement.class
+            );
+        tool
+            .run(
+                ImmutableMap.of(INDEX_FIELD, "index_name", PPL_FIELD, "source"),
+                ActionListener
+                    .<String>wrap(
+                        response -> assertEquals(expected, gson.fromJson(response, JsonElement.class)),
+                        e -> fail("Tool runs failed: " + e.getMessage())
+                    )
+            );
+    }
+
+    @SneakyThrows
+    @Test
+    public void testExecutionWithPPLInputWhenNoDataIsReturned() {
+        String emptyDataPPLResponse =
+            """
+                {"schema":[{"name":"field1","type":"string"},{"name":"field2","type":"string"},{"name":"field3","type":"long"}],"datarows":[],"total":0,"size":0}
+                """;
+        doAnswer(invocation -> {
+            ActionListener<TransportPPLQueryResponse> listener = (ActionListener<TransportPPLQueryResponse>) invocation.getArguments()[2];
+            listener.onResponse(pplQueryResponse);
+            return null;
+        }).when(client).execute(eq(PPLQueryAction.INSTANCE), any(), any());
+        when(pplQueryResponse.getResult()).thenReturn(emptyDataPPLResponse);
+        LogPatternTool tool = LogPatternTool.Factory.getInstance().create(params);
+
+        tool
+            .run(
+                ImmutableMap.of(INDEX_FIELD, "index_name", PPL_FIELD, "source"),
+                ActionListener
+                    .<String>wrap(
+                        response -> assertEquals("Can not get any data row from ppl response.", response),
+                        e -> fail("Tool runs failed: " + e.getMessage())
+                    )
+            );
+    }
+
+    @SneakyThrows
+    @Test
+    public void testExecutionWithPPLFailed() {
+        String pplFailureMessage = "Failed in execute ppl";
+        doAnswer(invocation -> {
+            ActionListener<SearchResponse> listener = (ActionListener<SearchResponse>) invocation.getArguments()[1];
+            listener.onFailure(new Exception(pplFailureMessage));
+            return null;
+        }).when(client).search(any(), any());
+
+        LogPatternTool tool = LogPatternTool.Factory.getInstance().create(params);
+        tool
+            .run(
+                ImmutableMap.of(INDEX_FIELD, "index_name", PPL_FIELD, "source"),
+                ActionListener.<String>wrap(response -> fail(), e -> assertEquals(pplFailureMessage, e.getMessage()))
+            );
+    }
 }