[ML] Checking for presence of error object when validating response (#118375) (#118795)

* Refactoring error handling logic

* Refactoring base response handler to remove duplication

* Update docs/changelog/118375.yaml

* Addressing feedback

---------

Co-authored-by: Max Hniebergall <[email protected]>
(cherry picked from commit f4dc716)

# Conflicts:
#	x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/alibabacloudsearch/AlibabaCloudSearchResponseHandler.java
#	x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/anthropic/AnthropicResponseHandler.java
#	x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandler.java
#	x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/elastic/ElasticInferenceServiceResponseHandler.java
#	x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/googleaistudio/GoogleAiStudioResponseHandler.java
#	x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/googlevertexai/GoogleVertexAiResponseHandler.java
#	x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/huggingface/HuggingFaceResponseHandler.java
#	x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/ibmwatsonx/IbmWatsonxResponseHandler.java
#	x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/openai/OpenAiResponseHandler.java

Author: jonathan-buttner

docs/changelog/118375.yaml

@@ -0,0 +1,5 @@
+pr: 118375
+summary: Check for presence of error object when validating streaming responses from integrations in the inference API
+area: Machine Learning
+type: enhancement
+issues: []

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/alibabacloudsearch/AlibabaCloudSearchResponseHandler.java

@@ -7,16 +7,12 @@
 
 package org.elasticsearch.xpack.inference.external.alibabacloudsearch;
 
-import org.apache.logging.log4j.Logger;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.BaseResponseHandler;
 import org.elasticsearch.xpack.inference.external.http.retry.ResponseParser;
 import org.elasticsearch.xpack.inference.external.http.retry.RetryException;
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.external.response.alibabacloudsearch.AlibabaCloudSearchErrorResponseEntity;
-import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
-
-import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
 
 /**
  * Defines how to handle various errors returned from the AlibabaCloudSearch integration.
@@ -27,21 +23,15 @@ public AlibabaCloudSearchResponseHandler(String requestType, ResponseParser pars
         super(requestType, parseFunction, AlibabaCloudSearchErrorResponseEntity::fromResponse);
     }
 
-    @Override
-    public void validateResponse(ThrottlerManager throttlerManager, Logger logger, Request request, HttpResult result)
-        throws RetryException {
-        checkForFailureStatusCode(request, result);
-        checkForEmptyBody(throttlerManager, logger, request, result);
-    }
-
     /**
      * Validates the status code throws an RetryException if not in the range [200, 300).
      *
      * @param request The http request
      * @param result  The http response and body
      * @throws RetryException Throws if status code is {@code >= 300 or < 200 }
      */
-    void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
+    @Override
+    protected void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
         int statusCode = result.response().getStatusLine().getStatusCode();
         if (statusCode >= 200 && statusCode < 300) {
             return;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/anthropic/AnthropicResponseHandler.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.inference.external.anthropic;
 
-import org.apache.logging.log4j.Logger;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.xpack.core.inference.results.StreamingChatCompletionResults;
@@ -19,11 +18,9 @@
 import org.elasticsearch.xpack.inference.external.response.ErrorMessageResponseEntity;
 import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventParser;
 import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventProcessor;
-import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 
 import java.util.concurrent.Flow;
 
-import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
 import static org.elasticsearch.xpack.inference.external.http.retry.ResponseHandlerUtils.getFirstHeaderOrUnknown;
 
 public class AnthropicResponseHandler extends BaseResponseHandler {
@@ -54,13 +51,6 @@ public AnthropicResponseHandler(String requestType, ResponseParser parseFunction
         this.canHandleStreamingResponses = canHandleStreamingResponses;
     }
 
-    @Override
-    public void validateResponse(ThrottlerManager throttlerManager, Logger logger, Request request, HttpResult result)
-        throws RetryException {
-        checkForFailureStatusCode(request, result);
-        checkForEmptyBody(throttlerManager, logger, request, result);
-    }
-
     @Override
     public boolean canHandleStreamingResponses() {
         return canHandleStreamingResponses;
@@ -83,7 +73,8 @@ public InferenceServiceResults parseResult(Request request, Flow.Publisher<HttpR
      * @param result  The http response and body
      * @throws RetryException Throws if status code is {@code >= 300 or < 200 }
      */
-    void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
+    @Override
+    protected void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
         int statusCode = result.response().getStatusLine().getStatusCode();
         if (statusCode >= 200 && statusCode < 300) {
             return;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandler.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.inference.external.cohere;
 
-import org.apache.logging.log4j.Logger;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.xpack.core.inference.results.StreamingChatCompletionResults;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
@@ -17,12 +16,9 @@
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.external.response.cohere.CohereErrorResponseEntity;
 import org.elasticsearch.xpack.inference.external.response.streaming.NewlineDelimitedByteProcessor;
-import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 
 import java.util.concurrent.Flow;
 
-import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
-
 /**
  * Defines how to handle various errors returned from the Cohere integration.
  *
@@ -45,13 +41,6 @@ public CohereResponseHandler(String requestType, ResponseParser parseFunction, b
         this.canHandleStreamingResponse = canHandleStreamingResponse;
     }
 
-    @Override
-    public void validateResponse(ThrottlerManager throttlerManager, Logger logger, Request request, HttpResult result)
-        throws RetryException {
-        checkForFailureStatusCode(request, result);
-        checkForEmptyBody(throttlerManager, logger, request, result);
-    }
-
     @Override
     public boolean canHandleStreamingResponses() {
         return canHandleStreamingResponse;
@@ -73,7 +62,8 @@ public InferenceServiceResults parseResult(Request request, Flow.Publisher<HttpR
      * @param result  The http response and body
      * @throws RetryException Throws if status code is {@code >= 300 or < 200 }
      */
-    void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
+    @Override
+    protected void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
         int statusCode = result.response().getStatusLine().getStatusCode();
         if (statusCode >= 200 && statusCode < 300) {
             return;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/elastic/ElasticInferenceServiceResponseHandler.java

@@ -7,17 +7,13 @@
 
 package org.elasticsearch.xpack.inference.external.elastic;
 
-import org.apache.logging.log4j.Logger;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.BaseResponseHandler;
 import org.elasticsearch.xpack.inference.external.http.retry.ContentTooLargeException;
 import org.elasticsearch.xpack.inference.external.http.retry.ResponseParser;
 import org.elasticsearch.xpack.inference.external.http.retry.RetryException;
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.external.response.elastic.ElasticInferenceServiceErrorResponseEntity;
-import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
-
-import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
 
 public class ElasticInferenceServiceResponseHandler extends BaseResponseHandler {
 
@@ -26,13 +22,7 @@ public ElasticInferenceServiceResponseHandler(String requestType, ResponseParser
     }
 
     @Override
-    public void validateResponse(ThrottlerManager throttlerManager, Logger logger, Request request, HttpResult result)
-        throws RetryException {
-        checkForFailureStatusCode(request, result);
-        checkForEmptyBody(throttlerManager, logger, request, result);
-    }
-
-    void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
+    protected void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
         int statusCode = result.response().getStatusLine().getStatusCode();
         if (statusCode >= 200 && statusCode < 300) {
             return;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/googleaistudio/GoogleAiStudioResponseHandler.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.inference.external.googleaistudio;
 
-import org.apache.logging.log4j.Logger;
 import org.elasticsearch.core.CheckedFunction;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.xcontent.XContentParser;
@@ -20,13 +19,11 @@
 import org.elasticsearch.xpack.inference.external.response.googleaistudio.GoogleAiStudioErrorResponseEntity;
 import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventParser;
 import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventProcessor;
-import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 
 import java.io.IOException;
 import java.util.concurrent.Flow;
 
 import static org.elasticsearch.core.Strings.format;
-import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
 
 public class GoogleAiStudioResponseHandler extends BaseResponseHandler {
 
@@ -52,13 +49,6 @@ public GoogleAiStudioResponseHandler(
         this.content = content;
     }
 
-    @Override
-    public void validateResponse(ThrottlerManager throttlerManager, Logger logger, Request request, HttpResult result)
-        throws RetryException {
-        checkForFailureStatusCode(request, result);
-        checkForEmptyBody(throttlerManager, logger, request, result);
-    }
-
     /**
      * Validates the status code and throws a RetryException if not in the range [200, 300).
      *
@@ -67,7 +57,8 @@ public void validateResponse(ThrottlerManager throttlerManager, Logger logger, R
      * @param result The http response and body
      * @throws RetryException Throws if status code is {@code >= 300 or < 200 }
      */
-    void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
+    @Override
+    protected void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
         int statusCode = result.response().getStatusLine().getStatusCode();
         if (statusCode >= 200 && statusCode < 300) {
             return;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/googlevertexai/GoogleVertexAiResponseHandler.java

@@ -7,17 +7,14 @@
 
 package org.elasticsearch.xpack.inference.external.googlevertexai;
 
-import org.apache.logging.log4j.Logger;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.BaseResponseHandler;
 import org.elasticsearch.xpack.inference.external.http.retry.ResponseParser;
 import org.elasticsearch.xpack.inference.external.http.retry.RetryException;
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.external.response.googlevertexai.GoogleVertexAiErrorResponseEntity;
-import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 
 import static org.elasticsearch.core.Strings.format;
-import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
 
 public class GoogleVertexAiResponseHandler extends BaseResponseHandler {
 
@@ -28,13 +25,7 @@ public GoogleVertexAiResponseHandler(String requestType, ResponseParser parseFun
     }
 
     @Override
-    public void validateResponse(ThrottlerManager throttlerManager, Logger logger, Request request, HttpResult result)
-        throws RetryException {
-        checkForFailureStatusCode(request, result);
-        checkForEmptyBody(throttlerManager, logger, request, result);
-    }
-
-    void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
+    protected void checkForFailureStatusCode(Request request, HttpResult result) throws RetryException {
         int statusCode = result.response().getStatusLine().getStatusCode();
         if (statusCode >= 200 && statusCode < 300) {
             return;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/BaseResponseHandler.java

@@ -7,16 +7,20 @@
 
 package org.elasticsearch.xpack.inference.external.http.retry;
 
+import org.apache.logging.log4j.Logger;
 import org.elasticsearch.ElasticsearchStatusException;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.request.Request;
+import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 
 import java.util.Objects;
 import java.util.function.Function;
 
 import static org.elasticsearch.core.Strings.format;
+import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
 
 public abstract class BaseResponseHandler implements ResponseHandler {
 
@@ -27,14 +31,15 @@ public abstract class BaseResponseHandler implements ResponseHandler {
     public static final String REDIRECTION = "Unhandled redirection";
     public static final String CONTENT_TOO_LARGE = "Received a content too large status code";
     public static final String UNSUCCESSFUL = "Received an unsuccessful status code";
+    public static final String SERVER_ERROR_OBJECT = "Received an error response";
     public static final String BAD_REQUEST = "Received a bad request status code";
     public static final String METHOD_NOT_ALLOWED = "Received a method not allowed status code";
 
     protected final String requestType;
     private final ResponseParser parseFunction;
-    private final Function<HttpResult, ErrorMessage> errorParseFunction;
+    private final Function<HttpResult, ErrorResponse> errorParseFunction;
 
-    public BaseResponseHandler(String requestType, ResponseParser parseFunction, Function<HttpResult, ErrorMessage> errorParseFunction) {
+    public BaseResponseHandler(String requestType, ResponseParser parseFunction, Function<HttpResult, ErrorResponse> errorParseFunction) {
         this.requestType = Objects.requireNonNull(requestType);
         this.parseFunction = Objects.requireNonNull(parseFunction);
         this.errorParseFunction = Objects.requireNonNull(errorParseFunction);
@@ -54,11 +59,42 @@ public String getRequestType() {
         return requestType;
     }
 
+    @Override
+    public void validateResponse(ThrottlerManager throttlerManager, Logger logger, Request request, HttpResult result) {
+        checkForFailureStatusCode(request, result);
+        checkForEmptyBody(throttlerManager, logger, request, result);
+
+        // When the response is streamed the status code could be 200 but the error object will be set
+        // so we need to check for that specifically
+        checkForErrorObject(request, result);
+    }
+
+    protected abstract void checkForFailureStatusCode(Request request, HttpResult result);
+
+    private void checkForErrorObject(Request request, HttpResult result) {
+        var errorEntity = errorParseFunction.apply(result);
+
+        if (errorEntity.errorStructureFound()) {
+            // We don't really know what happened because the status code was 200 so we'll return a failure and let the
+            // client retry if necessary
+            // If we did want to retry here, we'll need to determine if this was a streaming request, if it was
+            // we shouldn't retry because that would replay the entire streaming request and the client would get
+            // duplicate chunks back
+            throw new RetryException(false, buildError(SERVER_ERROR_OBJECT, request, result, errorEntity));
+        }
+    }
+
     protected Exception buildError(String message, Request request, HttpResult result) {
         var errorEntityMsg = errorParseFunction.apply(result);
+        return buildError(message, request, result, errorEntityMsg);
+    }
+
+    protected Exception buildError(String message, Request request, HttpResult result, ErrorResponse errorResponse) {
         var responseStatusCode = result.response().getStatusLine().getStatusCode();
 
-        if (errorEntityMsg == null) {
+        if (errorResponse == null
+            || errorResponse.errorStructureFound() == false
+            || Strings.isNullOrEmpty(errorResponse.getErrorMessage())) {
             return new ElasticsearchStatusException(
                 format(
                     "%s for request from inference entity id [%s] status [%s]",
@@ -76,7 +112,7 @@ protected Exception buildError(String message, Request request, HttpResult resul
                 message,
                 request.getInferenceEntityId(),
                 responseStatusCode,
-                errorEntityMsg.getErrorMessage()
+                errorResponse.getErrorMessage()
             ),
             toRestStatus(responseStatusCode)
         );

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/ErrorMessage.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.external.http.retry;
+
+import java.util.Objects;
+
+public class ErrorResponse {
+
+    // Denotes an error object that was not found
+    public static final ErrorResponse UNDEFINED_ERROR = new ErrorResponse(false);
+
+    private final String errorMessage;
+    private final boolean errorStructureFound;
+
+    public ErrorResponse(String errorMessage) {
+        this.errorMessage = Objects.requireNonNull(errorMessage);
+        this.errorStructureFound = true;
+    }
+
+    private ErrorResponse(boolean errorStructureFound) {
+        this.errorMessage = "";
+        this.errorStructureFound = errorStructureFound;
+    }
+
+    public String getErrorMessage() {
+        return errorMessage;
+    }
+
+    public boolean errorStructureFound() {
+        return errorStructureFound;
+    }
+}