Skip to content

Add error handling for plan&execute agent #3845

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Jun 11, 2025
+89 −2
Merged

Add error handling for plan&execute agent #3845

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
6f2c088
add more error handling
Hailong-am May 13, 2025
3f00b45
spotlessApply
Hailong-am May 14, 2025
6f5f6ed
Merge branch 'main' into agent_bug_fix
dhrubo-os Jun 3, 2025
a9293d1
Merge branch 'main' into agent_bug_fix
dhrubo-os Jun 5, 2025
d3a5bfb
add unit test
Hailong-am Jun 9, 2025
64b8aaa
Merge branch 'main' into agent_bug_fix
dhrubo-os Jun 9, 2025
7177987
Merge branch 'main' into agent_bug_fix
dhrubo-os Jun 10, 2025
03d4035
spotlessApply
Hailong-am Jun 11, 2025
608810a
add comments for index mapping tool
Hailong-am Jun 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
5 changes: 5 additions & 0 deletions ...in/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -555,6 +555,11 @@ String extractJsonFromMarkdown(String response) {
if (response.contains("```")) {
response = response.substring(0, response.lastIndexOf("```"));
}
} else {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. This method doesn't have any null check for response
  2. Is there any chance that multiple code block can exist?
  3. what about json array?
  4. what happens if there's any string literals like {"data": "{nested}"}
  5. Possible IndexOutOfBoundsException if markers aren't found
  6. why do we need second trim?
  7. Using substring multiple times creates multiple intermediate string objects
    Each operation like contains(), indexOf(), and lastIndexOf() scans through the string

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This method doesn't have any null check for response

response passed to this method could not be null, otherwise it will throw NPE before call it. But it's a good practice to add null check here

Is there any chance that multiple code block can exist?
what about json array?
what happens if there's any string literals like {"data": "{nested}"}

Yes, all of them are possible. the logic here is try our best to extract the expected json from LLM response.

Possible IndexOutOfBoundsException if markers aren't found

if markers are not found, it will not perform substring

why do we need second trim?

the content extracted from json block may have spaces, so trim is needed here before check it's a valid json string.

Using substring multiple times creates multiple intermediate string objects
Each operation like contains(), indexOf(), and lastIndexOf() scans through the string

that's not a big concern. the response is not large, and gc will help to handle intermediate string objects

// extract content from {} block
if (response.contains("{") && response.contains("}")) {
response = response.substring(response.indexOf("{"), response.lastIndexOf("}") + 1);
}
}

response = response.trim();
Expand Down
7 changes: 6 additions & 1 deletion ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()) {
Expand Down
7 changes: 7 additions & 0 deletions ...ava/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunnerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<>();
Expand Down
64 changes: 64 additions & 0 deletions ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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();
}
Expand Down Expand Up @@ -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();
Expand Down
Loading