Add status code to process error response

Signed-off-by: zane-neo <[email protected]>
zane-neo · Feb 29, 2024 · a553af4 · a553af4
1 parent fbbf172
commit a553af4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/ConnectorUtils.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/ConnectorUtils.java
@@ -226,10 +226,10 @@ public static ModelTensors processOutput(
         return ModelTensors.builder().statusCode(RestStatus.OK.getStatus()).mlModelTensors(modelTensors).build();
     }
 
-    public static ModelTensors processErrorResponse(String errorResponse) {
+    public static ModelTensors processErrorResponse(Integer statusCode,String errorResponse) {
         return ModelTensors
             .builder()
-            .statusCode(RestStatus.INTERNAL_SERVER_ERROR.getStatus())
+            .statusCode(statusCode)
             .mlModelTensors(List.of(ModelTensor.builder().dataAsMap(Map.of("remote_response", errorResponse)).build()))
             .build();
     }

diff --git a/...c/main/java/org/opensearch/ml/engine/algorithms/remote/MLSdkAsyncHttpResponseHandler.java b/...c/main/java/org/opensearch/ml/engine/algorithms/remote/MLSdkAsyncHttpResponseHandler.java
@@ -121,17 +121,17 @@ private void processResponse(
         ModelTensors tensors;
         if (Strings.isBlank(body)) {
             log.error("Remote model response body is empty!");
-            tensors = processErrorResponse("null");
+            tensors = processErrorResponse(statusCode, "null");
         } else {
             if (statusCode < HttpStatus.SC_OK || statusCode > HttpStatus.SC_MULTIPLE_CHOICES) {
                 log.error("Remote server returned error code: {}", statusCode);
-                tensors = processErrorResponse(body);
+                tensors = processErrorResponse(statusCode, body);
             } else {
                 try {
                     tensors = processOutput(body, connector, scriptService, parameters);
                 } catch (Exception e) {
                     log.error("Failed to process response body: {}", body, e);
-                    tensors = processErrorResponse(body);
+                    tensors = processErrorResponse(statusCode, body);
                 }
             }
         }