Refactor response handlers to improve error handling and streamline mid-stream error processing

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

@@ -12,10 +12,12 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
 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.Locale;
 import java.util.Objects;
 import java.util.function.Function;
 
@@ -34,17 +36,23 @@ public abstract class BaseResponseHandler implements ResponseHandler {
     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 static final String ERROR_TYPE = "error";
+    protected static final String STREAM_ERROR = "stream_error";
 
     protected final String requestType;
     protected final ResponseParser parseFunction;
     private final Function<HttpResult, ErrorResponse> errorParseFunction;
     private final boolean canHandleStreamingResponses;
 
-    public BaseResponseHandler(String requestType, ResponseParser parseFunction, Function<HttpResult, ErrorResponse> errorParseFunction) {
+    protected BaseResponseHandler(
+        String requestType,
+        ResponseParser parseFunction,
+        Function<HttpResult, ErrorResponse> errorParseFunction
+    ) {
         this(requestType, parseFunction, errorParseFunction, false);
     }
 
-    public BaseResponseHandler(
+    protected BaseResponseHandler(
         String requestType,
         ResponseParser parseFunction,
         Function<HttpResult, ErrorResponse> errorParseFunction,
@@ -109,19 +117,230 @@ private void checkForErrorObject(Request request, HttpResult result) {
     }
 
     protected Exception buildError(String message, Request request, HttpResult result) {
-        var errorEntityMsg = errorParseFunction.apply(result);
-        return buildError(message, request, result, errorEntityMsg);
+        var errorResponse = errorParseFunction.apply(result);
+        return buildError(message, request, result, errorResponse);
     }
 
     protected Exception buildError(String message, Request request, HttpResult result, ErrorResponse errorResponse) {
         var responseStatusCode = result.response().getStatusLine().getStatusCode();
         return new ElasticsearchStatusException(
-            errorMessage(message, request, result, errorResponse, responseStatusCode),
+            errorMessage(message, request, errorResponse, responseStatusCode),
             toRestStatus(responseStatusCode)
         );
     }
 
-    protected String errorMessage(String message, Request request, HttpResult result, ErrorResponse errorResponse, int statusCode) {
+    /**
+     * Builds an error for a streaming request with a custom error type.
+     * This method is used when an error response is received from the external service.
+     * Only streaming requests support this format, and it should be used when the error response.
+     *
+     * @param message            the error message to include in the exception
+     * @param request            the request that caused the error
+     * @param result             the HTTP result containing the error response
+     * @param errorResponse      the parsed error response from the HTTP result
+     * @param errorResponseClass the class of the expected error response type
+     * @return an instance of {@link UnifiedChatCompletionException} with details from the error response
+     */
+    protected UnifiedChatCompletionException buildChatCompletionError(
+        String message,
+        Request request,
+        HttpResult result,
+        ErrorResponse errorResponse,
+        Class<? extends ErrorResponse> errorResponseClass
+    ) {
+        assert request.isStreaming() : "Only streaming requests support this format";
+        var statusCode = result.response().getStatusLine().getStatusCode();
+        var errorMessage = errorMessage(message, request, errorResponse, statusCode);
+        var restStatus = toRestStatus(statusCode);
+
+        return buildChatCompletionError(errorResponse, errorMessage, restStatus, errorResponseClass);
+    }
+
+    /**
+     * Builds a {@link UnifiedChatCompletionException} for a streaming request.
+     * This method is used when an error response is received from the external service.
+     * Only streaming requests should use this method.
+     *
+     * @param errorResponse      the error response parsed from the HTTP result
+     * @param errorMessage       the error message to include in the exception
+     * @param restStatus         the REST status code of the response
+     * @param errorResponseClass the class of the expected error response type
+     * @return an instance of {@link UnifiedChatCompletionException} with details from the error response
+     */
+    protected UnifiedChatCompletionException buildChatCompletionError(
+        ErrorResponse errorResponse,
+        String errorMessage,
+        RestStatus restStatus,
+        Class<? extends ErrorResponse> errorResponseClass
+    ) {
+        if (errorResponseClass.isInstance(errorResponse)) {
+            return buildProviderSpecificChatCompletionError(errorResponse, errorMessage, restStatus);
+        } else {
+            return buildDefaultChatCompletionError(errorResponse, errorMessage, restStatus);
+        }
+    }
+
+    /**
+     * Builds a custom {@link UnifiedChatCompletionException} for a streaming request.
+     * This method is called when a specific error response is found in the HTTP result.
+     * It must be implemented by subclasses to handle specific error response formats.
+     * Only streaming requests should use this method.
+     *
+     * @param errorResponse the error response parsed from the HTTP result
+     * @param errorMessage  the error message to include in the exception
+     * @param restStatus    the REST status code of the response
+     * @return an instance of {@link UnifiedChatCompletionException} with details from the error response
+     */
+    protected UnifiedChatCompletionException buildProviderSpecificChatCompletionError(
+        ErrorResponse errorResponse,
+        String errorMessage,
+        RestStatus restStatus
+    ) {
+        throw new UnsupportedOperationException(
+            "Custom error handling is not implemented. Please override buildProviderSpecificChatCompletionError method."
+        );
+    }
+
+    /**
+     * Builds a default {@link UnifiedChatCompletionException} for a streaming request.
+     * This method is used when an error response is received but no specific error handling is implemented.
+     * Only streaming requests should use this method.
+     *
+     * @param errorResponse the error response parsed from the HTTP result
+     * @param errorMessage  the error message to include in the exception
+     * @param restStatus    the REST status code of the response
+     * @return an instance of {@link UnifiedChatCompletionException} with details from the error response
+     */
+    protected UnifiedChatCompletionException buildDefaultChatCompletionError(
+        ErrorResponse errorResponse,
+        String errorMessage,
+        RestStatus restStatus
+    ) {
+        return new UnifiedChatCompletionException(
+            restStatus,
+            errorMessage,
+            createErrorType(errorResponse),
+            restStatus.name().toLowerCase(Locale.ROOT)
+        );
+    }
+
+    /**
+     * Builds a mid-stream error for a streaming request.
+     * This method is used when an error occurs while processing a streaming response.
+     * It must be implemented by subclasses to handle specific error response formats.
+     * Only streaming requests should use this method.
+     *
+     * @param inferenceEntityId the ID of the inference entity
+     * @param message           the error message
+     * @param e                the exception that caused the error, can be null
+     * @return a {@link UnifiedChatCompletionException} representing the mid-stream error
+     */
+    public UnifiedChatCompletionException buildMidStreamChatCompletionError(String inferenceEntityId, String message, Exception e) {
+        throw new UnsupportedOperationException(
+            "Mid-stream error handling is not implemented. Please override buildMidStreamChatCompletionError method."
+        );
+    }
+
+    /**
+     * Builds a mid-stream error for a streaming request with a custom error type.
+     * This method is used when an error occurs while processing a streaming response and allows for custom error handling.
+     * Only streaming requests should use this method.
+     *
+     * @param inferenceEntityId the ID of the inference entity
+     * @param message           the error message
+     * @param e                the exception that caused the error, can be null
+     * @param errorResponseClass the class of the expected error response type
+     * @return a {@link UnifiedChatCompletionException} representing the mid-stream error
+     */
+    protected UnifiedChatCompletionException buildMidStreamChatCompletionError(
+        String inferenceEntityId,
+        String message,
+        Exception e,
+        Class<? extends ErrorResponse> errorResponseClass
+    ) {
+        // Extract the error response from the message using the provided method
+        var errorResponse = extractMidStreamChatCompletionErrorResponse(message);
+        // Check if the error response matches the expected type
+        if (errorResponseClass.isInstance(errorResponse)) {
+            // If it matches, we can build a custom mid-stream error exception
+            return buildProviderSpecificMidStreamChatCompletionError(inferenceEntityId, errorResponse);
+        } else if (e != null) {
+            // If the error response does not match, we can still return an exception based on the original throwable
+            return UnifiedChatCompletionException.fromThrowable(e);
+        } else {
+            // If no specific error response is found, we return a default mid-stream error
+            return buildDefaultMidStreamChatCompletionError(inferenceEntityId, errorResponse);
+        }
+    }
+
+    /**
+     * Builds a custom mid-stream {@link UnifiedChatCompletionException} for a streaming request.
+     * This method is called when a specific error response is found in the message.
+     * It must be implemented by subclasses to handle specific error response formats.
+     * Only streaming requests should use this method.
+     *
+     * @param inferenceEntityId the ID of the inference entity
+     * @param errorResponse     the error response parsed from the message
+     * @return an instance of {@link UnifiedChatCompletionException} with details from the error response
+     */
+    protected UnifiedChatCompletionException buildProviderSpecificMidStreamChatCompletionError(
+        String inferenceEntityId,
+        ErrorResponse errorResponse
+    ) {
+        throw new UnsupportedOperationException(
+            "Mid-stream error handling is not implemented for this response handler. "
+                + "Please override buildProviderSpecificMidStreamChatCompletionError method."
+        );
+    }
+
+    /**
+     * Builds a default mid-stream error for a streaming request.
+     * This method is used when no specific error response is found in the message.
+     * Only streaming requests should use this method.
+     *
+     * @param inferenceEntityId the ID of the inference entity
+     * @param errorResponse     the error response extracted from the message
+     * @return a {@link UnifiedChatCompletionException} representing the default mid-stream error
+     */
+    protected UnifiedChatCompletionException buildDefaultMidStreamChatCompletionError(
+        String inferenceEntityId,
+        ErrorResponse errorResponse
+    ) {
+        return new UnifiedChatCompletionException(
+            RestStatus.INTERNAL_SERVER_ERROR,
+            format("%s for request from inference entity id [%s]", SERVER_ERROR_OBJECT, inferenceEntityId),
+            createErrorType(errorResponse),
+            STREAM_ERROR
+        );
+    }
+
+    /**
+     * Extracts the mid-stream error response from the message.
+     * This method is used to parse the error response from a streaming message.
+     * It must be implemented by subclasses to handle specific error response formats.
+     * Only streaming requests should use this method.
+     *
+     * @param message the message containing the error response
+     * @return an {@link ErrorResponse} object representing the mid-stream error
+     */
+    protected ErrorResponse extractMidStreamChatCompletionErrorResponse(String message) {
+        throw new UnsupportedOperationException(
+            "Mid-stream error extraction is not implemented. Please override extractMidStreamChatCompletionErrorResponse method."
+        );
+    }
+
+    /**
+     * Creates a string representation of the error type based on the provided ErrorResponse.
+     * This method is used to generate a human-readable error type for logging or exception messages.
+     *
+     * @param errorResponse the ErrorResponse object
+     * @return a string representing the error type
+     */
+    protected static String createErrorType(ErrorResponse errorResponse) {
+        return errorResponse != null ? errorResponse.getClass().getSimpleName() : "unknown";
+    }
+
+    protected String errorMessage(String message, Request request, ErrorResponse errorResponse, int statusCode) {
         return (errorResponse == null
             || errorResponse.errorStructureFound() == false
             || Strings.isNullOrEmpty(errorResponse.getErrorMessage()))