From 6f2c0885fb27c484ee71046ddf01dc41288a413c Mon Sep 17 00:00:00 2001 From: Hailong Cui Date: Tue, 13 May 2025 14:26:49 +0800 Subject: [PATCH 1/2] 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/2] 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;