From 86d459829a3cd8c7c6013c1ca342b14db17c7072 Mon Sep 17 00:00:00 2001
From: Mingshi Liu <mingshl@amazon.com>
Date: Wed, 14 Aug 2024 10:39:20 -0700
Subject: [PATCH] try help escape when the payload is not valid

Signed-off-by: Mingshi Liu <mingshl@amazon.com>
---
 .../ml/common/connector/HttpConnector.java    | 46 ++++---------------
 .../remote/AwsConnectorExecutorTest.java      |  2 +-
 2 files changed, 10 insertions(+), 38 deletions(-)

diff --git a/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java b/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java
index 7f64725057..e7e92ef864 100644
--- a/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java
+++ b/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java
@@ -314,53 +314,25 @@ public  <T> T createPayload(String action, Map<String, String> parameters) {
         Optional<ConnectorAction> connectorAction = findAction(action);
         if (connectorAction.isPresent() && connectorAction.get().getRequestBody() != null) {
             String payload = connectorAction.get().getRequestBody();
-            payload = fillNullParameters(parameters, payload);
-//            parameters = formatArrayParameters(parameters);
             StringSubstitutor substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
             payload = substitutor.replace(payload);
-
-//            if (!isJson(payload)) {
-                try {
-                    JsonObject jsonObject = JsonParser.parseString(payload).getAsJsonObject();
-                    String validJson = jsonObject.toString();
-                    System.out.println("Valid JSON String: " + validJson);
-                } catch (Exception e) {
-                    System.out.println("Invalid JSON String, attempting manual fix...");
-
-                    // Manual fix by escaping double quotes
+            boolean isJson = isJson(payload);
+            if (!isJson) {
+                    // when the payload is not valid, try fix by escaping double quotes
                     String manuallyFixedJson = payload.replace("\"", "\\\"");
-                    System.out.println("Manually Fixed JSON String: " + manuallyFixedJson);
+                if (!isJson(manuallyFixedJson)){
+                    throw new IllegalArgumentException("Invalid payload: " + payload);
+                }
+                else{
+                    payload = manuallyFixedJson;
                 }
 
-//                throw new IllegalArgumentException("Invalid payload: " + payload);
-//            }
+        }
             return (T) payload;
         }
         return (T) parameters.get("http_body");
-    }
 
-    private Map<String,String> formatArrayParameters(Map<String, String> parameters) {
-        Map<String,String> newParameters =  new HashMap<>();
-        for (Map.Entry<String, String> entry : parameters.entrySet()) {
-            String key = entry.getKey();
-            String value = entry.getValue();
-            String escapedValue = escapeJsonArrayIfNeeded(value);
-            newParameters.put(key,escapedValue);
-        }
-        return newParameters;
     }
-    protected String escapeJsonArrayIfNeeded(String value) {
-        if (isJsonArray(value)) {
-            return value.replaceAll("([^\\\\])\"", "$1\\\\\"");
-        }
-        return value;
-    }
-
-    protected boolean isJsonArray(String value) {
-        Pattern jsonArrayPattern = Pattern.compile("^\\[.*\\]$");
-        return jsonArrayPattern.matcher(value).matches();
-    }
-
 
     protected String fillNullParameters(Map<String, String> parameters, String payload) {
         List<String> bodyParams = findStringParametersWithNullDefaultValue(payload);
diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutorTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutorTest.java
index dac471382e..f1a0ddb7c3 100644
--- a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutorTest.java
+++ b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutorTest.java
@@ -377,7 +377,7 @@ public void executePredict_TextDocsInferenceInput_withStepSize_partiallyFailed_t
 
         ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class);
         Mockito.verify(actionListener, times(1)).onFailure(exceptionCaptor.capture());
-//        assert exceptionCaptor.getValue() instanceof OpenSearchStatusException;
+        // assert exceptionCaptor.getValue() instanceof OpenSearchStatusException;
         assertEquals("test failure", exceptionCaptor.getValue().getMessage());
     }