Remove redundant builder and change the error code check

Signed-off-by: zane-neo <[email protected]>
zane-neo · Mar 13, 2024 · fdb129e · fdb129e
1 parent 98ce807
commit fdb129e
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 13 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/output/model/ModelTensors.java b/common/src/main/java/org/opensearch/ml/common/output/model/ModelTensors.java
@@ -35,12 +35,6 @@ public ModelTensors(List<ModelTensor> mlModelTensors) {
         this.mlModelTensors = mlModelTensors;
     }
 
-    @Builder
-    public ModelTensors(List<ModelTensor> mlModelTensors, Integer statusCode) {
-        this.mlModelTensors = mlModelTensors;
-        this.statusCode = statusCode;
-    }
-
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();

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
@@ -223,13 +223,12 @@ public static ModelTensors processOutput(
             Object filteredResponse = JsonPath.parse(response).read(parameters.get(RESPONSE_FILTER_FIELD));
             connector.parseResponse(filteredResponse, modelTensors, scriptReturnModelTensor);
         }
-        return ModelTensors.builder().statusCode(RestStatus.OK.getStatus()).mlModelTensors(modelTensors).build();
+        return ModelTensors.builder().mlModelTensors(modelTensors).build();
     }
 
-    public static ModelTensors processErrorResponse(Integer statusCode, String errorResponse) {
+    public static ModelTensors processErrorResponse(String errorResponse) {
         return ModelTensors
             .builder()
-            .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(statusCode, "Remote model response is empty!");
+            tensors = processErrorResponse("Remote model response is empty!");
         } else {
             if (statusCode < HttpStatus.SC_OK || statusCode > HttpStatus.SC_MULTIPLE_CHOICES) {
                 log.error("Remote server returned error code: {}", statusCode);
-                tensors = processErrorResponse(statusCode, body);
+                tensors = processErrorResponse(body);
             } else {
                 try {
                     tensors = processOutput(body, connector, scriptService, parameters);
                 } catch (Exception e) {
                     log.error("Failed to process response body: {}", body, e);
-                    tensors = processErrorResponse(statusCode, body);
+                    tensors = processErrorResponse(body);
                 }
             }
         }
@@ -158,7 +158,7 @@ private void reOrderTensorResponses(Map<Integer, ModelTensors> tensorOutputs) {
             // either one fails, we will return a failure response.
             OpenSearchStatusException openSearchStatusException = null;
             for (Map.Entry<Integer, ModelTensors> entry : sortedMap.entrySet()) {
-                if (entry.getValue().getStatusCode() != HttpStatus.SC_OK) {
+                if (entry.getValue().getStatusCode() < HttpStatus.SC_OK || entry.getValue().getStatusCode() > HttpStatus.SC_MULTIPLE_CHOICES) {
                     openSearchStatusException = buildOpenSearchStatusException(entry.getValue());
                     break;
                 }