Refactor error handling in response classes to use ChatCompletionErrorResponse for improved consistency and maintainability

Author: Jan-Kazlouski-elastic

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

@@ -19,9 +19,7 @@
 
 import java.util.Locale;
 import java.util.Objects;
-import java.util.function.BiFunction;
 import java.util.function.Function;
-import java.util.function.Supplier;
 
 import static org.elasticsearch.core.Strings.format;
 import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
@@ -38,7 +36,6 @@ 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;
@@ -140,47 +137,22 @@ protected Exception buildError(String message, Request request, HttpResult resul
      * @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 errorResponseClassSupplier the supplier that provides the class of the expected error response type
-     * @param chatCompletionErrorBuilder the builder for creating provider-specific chat completion errors
      * @return an instance of {@link UnifiedChatCompletionException} with details from the error response
      */
     protected UnifiedChatCompletionException buildChatCompletionError(
         String message,
         Request request,
         HttpResult result,
-        ErrorResponse errorResponse,
-        Supplier<Class<? extends ErrorResponse>> errorResponseClassSupplier,
-        ChatCompletionErrorBuilder chatCompletionErrorBuilder
+        ErrorResponse errorResponse
     ) {
         assert request.isStreaming() : "Only streaming requests support this format";
         var statusCode = result.response().getStatusLine().getStatusCode();
         var errorMessage = extractErrorMessage(message, request, errorResponse, statusCode);
         var restStatus = toRestStatus(statusCode);
 
-        return buildChatCompletionError(errorResponse, errorMessage, restStatus, errorResponseClassSupplier, chatCompletionErrorBuilder);
-    }
-
-    /**
-     * 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 errorResponseClassSupplier the supplier that provides the class of the expected error response type
-     * @param chatCompletionErrorBuilder the builder for creating provider-specific chat completion errors
-     * @return an instance of {@link UnifiedChatCompletionException} with details from the error response
-     */
-    protected UnifiedChatCompletionException buildChatCompletionError(
-        ErrorResponse errorResponse,
-        String errorMessage,
-        RestStatus restStatus,
-        Supplier<Class<? extends ErrorResponse>> errorResponseClassSupplier,
-        ChatCompletionErrorBuilder chatCompletionErrorBuilder
-    ) {
-        if (errorResponse.errorStructureFound() && errorResponseClassSupplier.get().isInstance(errorResponse)) {
-            return chatCompletionErrorBuilder.buildProviderSpecificChatCompletionError(errorResponse, errorMessage, restStatus);
+        if (errorResponse.errorStructureFound()
+            && errorResponse instanceof UnifiedChatCompletionExceptionConvertible chatCompletionExceptionConvertible) {
+            return chatCompletionExceptionConvertible.toUnifiedChatCompletionException(errorMessage, restStatus);
         } else {
             return buildDefaultChatCompletionError(errorResponse, errorMessage, restStatus);
         }
@@ -196,7 +168,7 @@ protected UnifiedChatCompletionException buildChatCompletionError(
      * @param restStatus    the REST status code of the response
      * @return an instance of {@link UnifiedChatCompletionException} with details from the error response
      */
-    private static UnifiedChatCompletionException buildDefaultChatCompletionError(
+    protected static UnifiedChatCompletionException buildDefaultChatCompletionError(
         ErrorResponse errorResponse,
         String errorMessage,
         RestStatus restStatus
@@ -217,31 +189,27 @@ private static UnifiedChatCompletionException buildDefaultChatCompletionError(
      * @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 errorResponseClassSupplier a supplier that provides the class of the expected error response type
-     * @param specificErrorBuilder a function that builds a specific error based on the inference entity ID and error response
      * @param midStreamErrorExtractor a function that extracts the mid-stream error response from the message
      * @return a {@link UnifiedChatCompletionException} representing the mid-stream error
      */
     protected UnifiedChatCompletionException buildMidStreamChatCompletionError(
         String inferenceEntityId,
         String message,
         Exception e,
-        Supplier<Class<? extends ErrorResponse>> errorResponseClassSupplier,
-        BiFunction<String, ErrorResponse, UnifiedChatCompletionException> specificErrorBuilder,
         Function<String, ErrorResponse> midStreamErrorExtractor
     ) {
         // Extract the error response from the message using the provided method
-        var errorResponse = midStreamErrorExtractor.apply(message);
+        var error = midStreamErrorExtractor.apply(message);
         // Check if the error response matches the expected type
-        if (errorResponse.errorStructureFound() && errorResponseClassSupplier.get().isInstance(errorResponse)) {
+        if (error.errorStructureFound() && error instanceof MidStreamUnifiedChatCompletionExceptionConvertible midStreamError) {
             // If it matches, we can build a custom mid-stream error exception
-            return specificErrorBuilder.apply(inferenceEntityId, errorResponse);
+            return midStreamError.toUnifiedChatCompletionException(inferenceEntityId);
         } 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);
+            return buildDefaultMidStreamChatCompletionError(inferenceEntityId, error);
         }
     }
 
@@ -277,7 +245,7 @@ private static String createErrorType(ErrorResponse errorResponse) {
         return errorResponse != null ? errorResponse.getClass().getSimpleName() : "unknown";
     }
 
-    private static String extractErrorMessage(String message, Request request, ErrorResponse errorResponse, int statusCode) {
+    protected static String extractErrorMessage(String message, Request request, ErrorResponse errorResponse, int statusCode) {
         return (errorResponse == null
             || errorResponse.errorStructureFound() == false
             || Strings.isNullOrEmpty(errorResponse.getErrorMessage()))
@@ -291,7 +259,7 @@ private static String extractErrorMessage(String message, Request request, Error
                 );
     }
 
-    public static RestStatus toRestStatus(int statusCode) {
+    protected static RestStatus toRestStatus(int statusCode) {
         RestStatus code = null;
         if (statusCode < 500) {
             code = RestStatus.fromCode(statusCode);

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

@@ -0,0 +1,16 @@
+/*
+ * 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 org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
+
+public interface MidStreamUnifiedChatCompletionExceptionConvertible {
+
+    UnifiedChatCompletionException toUnifiedChatCompletionException(String inferenceEntityId);
+
+}

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

@@ -0,0 +1,17 @@
+/*
+ * 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 org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
+
+public interface UnifiedChatCompletionExceptionConvertible {
+
+    UnifiedChatCompletionException toUnifiedChatCompletionException(String errorMessage, RestStatus restStatus);
+
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/streaming/StreamingErrorResponse.java renamed to x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/streaming/OpenAiStreamingChatCompletionErrorResponse.java

@@ -8,19 +8,26 @@
 package org.elasticsearch.xpack.inference.external.response.streaming;
 
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xcontent.XContentParserConfiguration;
 import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.http.retry.MidStreamUnifiedChatCompletionExceptionConvertible;
+import org.elasticsearch.xpack.inference.external.http.retry.UnifiedChatCompletionExceptionConvertible;
 import org.elasticsearch.xpack.inference.external.response.ErrorMessageResponseEntity;
 
 import java.util.Objects;
 import java.util.Optional;
 
+import static org.elasticsearch.core.Strings.format;
+import static org.elasticsearch.xpack.inference.external.http.retry.BaseResponseHandler.SERVER_ERROR_OBJECT;
+
 /**
  * Represents an error response from a streaming inference service.
  * This class extends {@link ErrorResponse} and provides additional fields
@@ -38,17 +45,21 @@
  * </code></pre>
  * TODO: {@link ErrorMessageResponseEntity} is nearly identical to this, but doesn't parse as many fields. We must remove the duplication.
  */
-public class StreamingErrorResponse extends ErrorResponse {
+public class OpenAiStreamingChatCompletionErrorResponse extends ErrorResponse
+    implements
+        UnifiedChatCompletionExceptionConvertible,
+        MidStreamUnifiedChatCompletionExceptionConvertible {
     private static final ConstructingObjectParser<Optional<ErrorResponse>, Void> ERROR_PARSER = new ConstructingObjectParser<>(
         "streaming_error",
         true,
-        args -> Optional.ofNullable((StreamingErrorResponse) args[0])
-    );
-    private static final ConstructingObjectParser<StreamingErrorResponse, Void> ERROR_BODY_PARSER = new ConstructingObjectParser<>(
-        "streaming_error",
-        true,
-        args -> new StreamingErrorResponse((String) args[0], (String) args[1], (String) args[2], (String) args[3])
+        args -> Optional.ofNullable((OpenAiStreamingChatCompletionErrorResponse) args[0])
     );
+    private static final ConstructingObjectParser<OpenAiStreamingChatCompletionErrorResponse, Void> ERROR_BODY_PARSER =
+        new ConstructingObjectParser<>(
+            "streaming_error",
+            true,
+            args -> new OpenAiStreamingChatCompletionErrorResponse((String) args[0], (String) args[1], (String) args[2], (String) args[3])
+        );
 
     static {
         ERROR_BODY_PARSER.declareString(ConstructingObjectParser.constructorArg(), new ParseField("message"));
@@ -105,13 +116,34 @@ public static ErrorResponse fromString(String response) {
     private final String param;
     private final String type;
 
-    StreamingErrorResponse(String errorMessage, @Nullable String code, @Nullable String param, String type) {
+    OpenAiStreamingChatCompletionErrorResponse(String errorMessage, @Nullable String code, @Nullable String param, String type) {
         super(errorMessage);
         this.code = code;
         this.param = param;
         this.type = Objects.requireNonNull(type);
     }
 
+    @Override
+    public UnifiedChatCompletionException toUnifiedChatCompletionException(String errorMessage, RestStatus restStatus) {
+        return new UnifiedChatCompletionException(restStatus, errorMessage, this.type(), this.code(), this.param());
+    }
+
+    @Override
+    public UnifiedChatCompletionException toUnifiedChatCompletionException(String inferenceEntityId) {
+        return new UnifiedChatCompletionException(
+            RestStatus.INTERNAL_SERVER_ERROR,
+            format(
+                "%s for request from inference entity id [%s]. Error message: [%s]",
+                SERVER_ERROR_OBJECT,
+                inferenceEntityId,
+                this.getErrorMessage()
+            ),
+            this.type(),
+            this.code(),
+            this.param()
+        );
+    }
+
     @Nullable
     public String code() {
         return code;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceUnifiedChatCompletionResponseHandler.java

@@ -30,6 +30,8 @@
  * This handler is designed to work with the unified Elastic Inference Service chat completion API.
  */
 public class ElasticInferenceServiceUnifiedChatCompletionResponseHandler extends ElasticInferenceServiceResponseHandler {
+    private static final String ERROR_TYPE = "error";
+
     public ElasticInferenceServiceUnifiedChatCompletionResponseHandler(String requestType, ResponseParser parseFunction) {
         super(requestType, parseFunction, true);
     }
@@ -59,22 +61,16 @@ public InferenceServiceResults parseResult(Request request, Flow.Publisher<HttpR
      */
     @Override
     protected UnifiedChatCompletionException buildError(String message, Request request, HttpResult result, ErrorResponse errorResponse) {
-        return buildChatCompletionError(
-            message,
-            request,
-            result,
-            errorResponse,
-            () -> ErrorResponse.class,
-            ElasticInferenceServiceUnifiedChatCompletionResponseHandler::buildProviderSpecificChatCompletionError
-        );
-    }
+        assert request.isStreaming() : "Only streaming requests support this format";
+        var statusCode = result.response().getStatusLine().getStatusCode();
+        var errorMessage = extractErrorMessage(message, request, errorResponse, statusCode);
+        var restStatus = toRestStatus(statusCode);
 
-    private static UnifiedChatCompletionException buildProviderSpecificChatCompletionError(
-        ErrorResponse response,
-        String message,
-        RestStatus restStatus
-    ) {
-        return new UnifiedChatCompletionException(restStatus, message, ERROR_TYPE, restStatus.name().toLowerCase(Locale.ROOT));
+        if (errorResponse.errorStructureFound()) {
+            return new UnifiedChatCompletionException(restStatus, errorMessage, ERROR_TYPE, restStatus.name().toLowerCase(Locale.ROOT));
+        } else {
+            return buildDefaultChatCompletionError(errorResponse, errorMessage, restStatus);
+        }
     }
 
     /**