Add error handling for plan&execute agent (#3845)

Author: Hailong-am

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java

@@ -608,7 +608,13 @@ private static void runTool(
                         );
                     nextStepListener
                         .onResponse(
-                            String.format(Locale.ROOT, "Failed to run the tool %s with the error message %s.", finalAction, e.getMessage())
+                            String
+                                .format(
+                                    Locale.ROOT,
+                                    "Failed to run the tool %s with the error message %s.",
+                                    finalAction,
+                                    e.getMessage().replaceAll("\\n", "\n")
+                                )
                         );
                 });
                 if (tools.get(action) instanceof MLModelTool) {

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java

@@ -555,6 +555,11 @@ String extractJsonFromMarkdown(String response) {
             if (response.contains("```")) {
                 response = response.substring(0, response.lastIndexOf("```"));
             }
+        } else {
+            // extract content from {} block
+            if (response.contains("{") && response.contains("}")) {
+                response = response.substring(response.indexOf("{"), response.lastIndexOf("}") + 1);
+            }
         }
 
         response = response.trim();

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java

@@ -95,7 +95,12 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
         try {
             List<String> indexList = new ArrayList<>();
             if (StringUtils.isNotBlank(parameters.get("index"))) {
-                indexList = gson.fromJson(parameters.get("index"), List.class);
+                try {
+                    indexList = gson.fromJson(parameters.get("index"), List.class);
+                } catch (Exception e) {
+                    // sometimes the input comes from LLM is not a json string, it might a single value of index name
+                    indexList.add(parameters.get("index"));
+                }
             }
 
             if (indexList.isEmpty()) {

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunnerTest.java

@@ -451,6 +451,13 @@ public void testExtractJsonFromMarkdown() {
         assertEquals("{\"key\":\"value\"}", result);
     }
 
+    @Test
+    public void testExtractJsonFromMarkdownWithoutJsonPrefix() {
+        String markdown = "This is the json output {\"key\":\"value\"}\n";
+        String result = mlPlanExecuteAndReflectAgentRunner.extractJsonFromMarkdown(markdown);
+        assertEquals("{\"key\":\"value\"}", result);
+    }
+
     @Test
     public void testAddToolsToPrompt() {
         Map<String, String> testParams = new HashMap<>();

ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java

@@ -54,6 +54,7 @@ public class IndexMappingToolTests {
     private GetIndexResponse getIndexResponse;
 
     private Map<String, String> indexParams;
+    private Map<String, String> indexParamsWithRawIndexName;
     private Map<String, String> otherParams;
     private Map<String, String> emptyParams;
 
@@ -67,6 +68,7 @@ public void setup() {
         IndexMappingTool.Factory.getInstance().init(client);
 
         indexParams = Map.of("index", "[\"foo\"]");
+        indexParamsWithRawIndexName = Map.of("index", "foo");
         otherParams = Map.of("other", "[\"bar\"]");
         emptyParams = Collections.emptyMap();
     }
@@ -175,6 +177,68 @@ public void testRunAsyncIndexMapping() throws Exception {
         assertTrue(responseList.contains("test.int.setting=123"));
     }
 
+    @Test
+    public void testRunWithRawIndexNameInput() throws Exception {
+        String indexName = "foo";
+
+        @SuppressWarnings("unchecked")
+        ArgumentCaptor<ActionListener<GetIndexResponse>> actionListenerCaptor = ArgumentCaptor.forClass(ActionListener.class);
+        doNothing().when(indicesAdminClient).getIndex(any(), actionListenerCaptor.capture());
+
+        when(getIndexResponse.indices()).thenReturn(new String[] { indexName });
+        Settings settings = Settings.builder().put("test.boolean.setting", false).put("test.int.setting", 123).build();
+        when(getIndexResponse.settings()).thenReturn(Map.of(indexName, settings));
+        String source = """
+            {
+                "foo": {
+                    "mappings": {
+                        "year": {
+                            "full_name": "year",
+                            "mapping": {
+                                "year": {
+                                    "type": "text"
+                                }
+                            }
+                        },
+                        "age": {
+                            "full_name": "age",
+                            "mapping": {
+                                "age": {
+                                    "type": "integer"
+                                }
+                            }
+                        }
+                    }
+                }
+            }""";
+        MappingMetadata mapping = new MappingMetadata(indexName, XContentHelper.convertToMap(JsonXContent.jsonXContent, source, true));
+        when(getIndexResponse.mappings()).thenReturn(Map.of(indexName, mapping));
+
+        // Now make the call
+        Tool tool = IndexMappingTool.Factory.getInstance().create(Collections.emptyMap());
+        final CompletableFuture<String> future = new CompletableFuture<>();
+        ActionListener<String> listener = ActionListener.wrap(r -> { future.complete(r); }, e -> { future.completeExceptionally(e); });
+
+        tool.run(indexParamsWithRawIndexName, listener);
+        actionListenerCaptor.getValue().onResponse(getIndexResponse);
+
+        future.orTimeout(10, TimeUnit.SECONDS).join();
+        String response = future.get();
+        List<String> responseList = Arrays.asList(response.trim().split("\\n"));
+
+        assertTrue(responseList.contains("index: foo"));
+
+        assertTrue(responseList.contains("mappings:"));
+        assertTrue(
+            responseList
+                .contains("mappings={year={full_name=year, mapping={year={type=text}}}, age={full_name=age, mapping={age={type=integer}}}}")
+        );
+
+        assertTrue(responseList.contains("settings:"));
+        assertTrue(responseList.contains("test.boolean.setting=false"));
+        assertTrue(responseList.contains("test.int.setting=123"));
+    }
+
     @Test
     public void testTool() {
         Factory instance = IndexMappingTool.Factory.getInstance();