From 6f2c0885fb27c484ee71046ddf01dc41288a413c Mon Sep 17 00:00:00 2001 From: Hailong Cui Date: Tue, 13 May 2025 14:26:49 +0800 Subject: [PATCH 1/5] add more error handling Signed-off-by: Hailong Cui --- .../ml/engine/algorithms/agent/MLChatAgentRunner.java | 3 ++- .../agent/MLPlanExecuteAndReflectAgentRunner.java | 7 ++++++- .../org/opensearch/ml/engine/tools/IndexMappingTool.java | 6 +++++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java index c5f7c22ed6..3925965a41 100644 --- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java +++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java @@ -701,7 +701,8 @@ 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) { diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java index ef12c385c7..ff450df4ea 100644 --- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java +++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java @@ -520,11 +520,16 @@ private 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(); if (!isJson(response)) { - throw new IllegalStateException("Failed to parse LLM output due to invalid JSON"); + throw new IllegalStateException("Failed to parse LLM output due to invalid JSON, response=" + response); } return response; diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java index 3a41dfc2d1..80f497335a 100644 --- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java +++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java @@ -95,7 +95,11 @@ public void run(Map parameters, ActionListener listener) try { List 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) { + indexList.add(parameters.get("index")); + } } if (indexList.isEmpty()) { From 3f00b456ebaf53e64d0e9d82e3da51292cb5745f Mon Sep 17 00:00:00 2001 From: Hailong Cui Date: Wed, 14 May 2025 13:37:12 +0800 Subject: [PATCH 2/5] spotlessApply Signed-off-by: Hailong Cui --- .../ml/engine/algorithms/agent/MLChatAgentRunner.java | 9 +++++++-- .../agent/MLPlanExecuteAndReflectAgentRunner.java | 2 +- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java index 3925965a41..c52cd5427c 100644 --- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java +++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java @@ -701,8 +701,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().replaceAll("\\n", "\n")) + 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) { diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java index ff450df4ea..0c529b4622 100644 --- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java +++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java @@ -529,7 +529,7 @@ private String extractJsonFromMarkdown(String response) { response = response.trim(); if (!isJson(response)) { - throw new IllegalStateException("Failed to parse LLM output due to invalid JSON, response=" + response); + throw new IllegalStateException("Failed to parse LLM output due to invalid JSON"); } return response; From d3a5bfbf106babb6e031e49c73572dd3b95a7a9f Mon Sep 17 00:00:00 2001 From: Hailong Cui Date: Mon, 9 Jun 2025 17:30:39 +0800 Subject: [PATCH 3/5] add unit test Signed-off-by: Hailong Cui --- ...LPlanExecuteAndReflectAgentRunnerTest.java | 8 +++ .../engine/tools/IndexMappingToolTests.java | 64 +++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunnerTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunnerTest.java index e608283ed3..5a14fa381c 100644 --- a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunnerTest.java +++ b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunnerTest.java @@ -446,6 +446,14 @@ 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 testParams = new HashMap<>(); diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java index cc879097b7..c18725f634 100644 --- a/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java +++ b/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java @@ -54,6 +54,7 @@ public class IndexMappingToolTests { private GetIndexResponse getIndexResponse; private Map indexParams; + private Map indexParamsWithRawIndexName; private Map otherParams; private Map 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> 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 future = new CompletableFuture<>(); + ActionListener 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 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(); From 03d40352cd24ecfb9ca50c71e80fd1dffcd08103 Mon Sep 17 00:00:00 2001 From: Hailong Cui Date: Wed, 11 Jun 2025 08:37:49 +0800 Subject: [PATCH 4/5] spotlessApply Signed-off-by: Hailong Cui --- ...LPlanExecuteAndReflectAgentRunnerTest.java | 1 - .../engine/tools/IndexMappingToolTests.java | 38 +++++++++---------- 2 files changed, 19 insertions(+), 20 deletions(-) diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunnerTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunnerTest.java index eda38820c1..6cfb7b09d7 100644 --- a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunnerTest.java +++ b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunnerTest.java @@ -458,7 +458,6 @@ public void testExtractJsonFromMarkdownWithoutJsonPrefix() { assertEquals("{\"key\":\"value\"}", result); } - @Test public void testAddToolsToPrompt() { Map testParams = new HashMap<>(); diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java index c18725f634..331197b5ce 100644 --- a/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java +++ b/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java @@ -189,28 +189,28 @@ public void testRunWithRawIndexNameInput() throws Exception { 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" - } + { + "foo": { + "mappings": { + "year": { + "full_name": "year", + "mapping": { + "year": { + "type": "text" } - }, - "age": { - "full_name": "age", - "mapping": { - "age": { - "type": "integer" - } + } + }, + "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)); @@ -230,8 +230,8 @@ public void testRunWithRawIndexNameInput() throws Exception { assertTrue(responseList.contains("mappings:")); assertTrue( - responseList - .contains("mappings={year={full_name=year, mapping={year={type=text}}}, age={full_name=age, mapping={age={type=integer}}}}") + responseList + .contains("mappings={year={full_name=year, mapping={year={type=text}}}, age={full_name=age, mapping={age={type=integer}}}}") ); assertTrue(responseList.contains("settings:")); From 608810adca4c0d2caab535bd4c9e88ebefef04c7 Mon Sep 17 00:00:00 2001 From: Hailong Cui Date: Wed, 11 Jun 2025 10:36:30 +0800 Subject: [PATCH 5/5] add comments for index mapping tool Signed-off-by: Hailong Cui --- .../java/org/opensearch/ml/engine/tools/IndexMappingTool.java | 1 + 1 file changed, 1 insertion(+) diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java index 80f497335a..e9b46bfbed 100644 --- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java +++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java @@ -98,6 +98,7 @@ public void run(Map parameters, ActionListener listener) 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")); } }