diff --git a/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java b/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java index 22165c487f..b70b7a1027 100644 --- a/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java +++ b/common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java @@ -139,11 +139,16 @@ public static boolean isJson(String json) { * prepareJsonValue("{\"key\":123}") → {\"key\":123} (valid JSON object, unchanged) * * @param input + * @param escapeJson * @return */ - public static String prepareJsonValue(String input) { + public static String prepareJsonValue(String input, boolean escapeJson) { if (isJson(input)) { - return input; + if (!escapeJson) { + return input; + } else { + return escapeJson(input); + } } return escapeJson(input); } diff --git a/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java b/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java index becf53e3c8..d80b8c6794 100644 --- a/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java +++ b/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java @@ -33,6 +33,7 @@ public class ToolUtils { public static final String TOOL_OUTPUT_FILTERS_FIELD = "output_filter"; + public static final String TOOL_ESCAPE_OUTPUT = "escape_output"; public static final String TOOL_REQUIRED_PARAMS = "required_parameters"; public static final String NO_ESCAPE_PARAMS = "no_escape_params"; diff --git a/common/src/test/java/org/opensearch/ml/common/utils/StringUtilsTest.java b/common/src/test/java/org/opensearch/ml/common/utils/StringUtilsTest.java index a64630de40..f2bb9aa4fd 100644 --- a/common/src/test/java/org/opensearch/ml/common/utils/StringUtilsTest.java +++ b/common/src/test/java/org/opensearch/ml/common/utils/StringUtilsTest.java @@ -861,14 +861,21 @@ public void testValidateFields_InvalidCharacterSet() { @Test public void prepareJsonValue_returnsRawIfJson() { String json = "{\"key\": 123}"; - String result = StringUtils.prepareJsonValue(json); + String result = StringUtils.prepareJsonValue(json, false); assertSame(json, result); // branch where isJson(input)==true } + @Test + public void prepareJsonValue_returnEscapeJsonIfForce() { + String json = "{\"key\": 123}"; + String result = StringUtils.prepareJsonValue(json, true); + assertEquals("{\\\"key\\\": 123}", result); + } + @Test public void prepareJsonValue_escapesBadCharsOtherwise() { String input = "Tom & Jerry \"