Refactor error handling in streaming response handlers to use functional interfaces for improved flexibility

Author: Jan-Kazlouski-elastic

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

@@ -19,7 +19,9 @@
 
 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;
@@ -124,7 +126,7 @@ protected Exception buildError(String message, Request request, HttpResult resul
     protected Exception buildError(String message, Request request, HttpResult result, ErrorResponse errorResponse) {
         var responseStatusCode = result.response().getStatusLine().getStatusCode();
         return new ElasticsearchStatusException(
-            errorMessage(message, request, errorResponse, responseStatusCode),
+            extractErrorMessage(message, request, errorResponse, responseStatusCode),
             toRestStatus(responseStatusCode)
         );
     }
@@ -138,22 +140,24 @@ 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 errorResponseClass the class of the expected error response type
+     * @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,
-        Class<? extends ErrorResponse> errorResponseClass
+        Supplier<Class<? extends ErrorResponse>> errorResponseClassSupplier,
+        ChatCompletionErrorBuilder chatCompletionErrorBuilder
     ) {
         assert request.isStreaming() : "Only streaming requests support this format";
         var statusCode = result.response().getStatusLine().getStatusCode();
-        var errorMessage = errorMessage(message, request, errorResponse, statusCode);
+        var errorMessage = extractErrorMessage(message, request, errorResponse, statusCode);
         var restStatus = toRestStatus(statusCode);
 
-        return buildChatCompletionError(errorResponse, errorMessage, restStatus, errorResponseClass);
+        return buildChatCompletionError(errorResponse, errorMessage, restStatus, errorResponseClassSupplier, chatCompletionErrorBuilder);
     }
 
     /**
@@ -164,43 +168,24 @@ protected UnifiedChatCompletionException buildChatCompletionError(
      * @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
+     * @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,
-        Class<? extends ErrorResponse> errorResponseClass
+        Supplier<Class<? extends ErrorResponse>> errorResponseClassSupplier,
+        ChatCompletionErrorBuilder chatCompletionErrorBuilder
     ) {
-        if (errorResponseClass.isInstance(errorResponse)) {
-            return buildProviderSpecificChatCompletionError(errorResponse, errorMessage, restStatus);
+        if (errorResponseClassSupplier.get().isInstance(errorResponse)) {
+            return chatCompletionErrorBuilder.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.
@@ -211,7 +196,7 @@ protected UnifiedChatCompletionException buildProviderSpecificChatCompletionErro
      * @param restStatus    the REST status code of the response
      * @return an instance of {@link UnifiedChatCompletionException} with details from the error response
      */
-    protected UnifiedChatCompletionException buildDefaultChatCompletionError(
+    private static UnifiedChatCompletionException buildDefaultChatCompletionError(
         ErrorResponse errorResponse,
         String errorMessage,
         RestStatus restStatus
@@ -232,21 +217,25 @@ protected 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 errorResponseClass the class of the expected error response type
+     * @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,
-        Class<? extends ErrorResponse> errorResponseClass
+        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 = extractMidStreamChatCompletionErrorResponse(message);
+        var errorResponse = midStreamErrorExtractor.apply(message);
         // Check if the error response matches the expected type
-        if (errorResponseClass.isInstance(errorResponse)) {
+        if (errorResponseClassSupplier.get().isInstance(errorResponse)) {
             // If it matches, we can build a custom mid-stream error exception
-            return buildProviderSpecificMidStreamChatCompletionError(inferenceEntityId, errorResponse);
+            return specificErrorBuilder.apply(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);
@@ -256,26 +245,6 @@ protected UnifiedChatCompletionException buildMidStreamChatCompletionError(
         }
     }
 
-    /**
-     * 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.
@@ -285,7 +254,7 @@ protected UnifiedChatCompletionException buildProviderSpecificMidStreamChatCompl
      * @param errorResponse     the error response extracted from the message
      * @return a {@link UnifiedChatCompletionException} representing the default mid-stream error
      */
-    protected UnifiedChatCompletionException buildDefaultMidStreamChatCompletionError(
+    protected static UnifiedChatCompletionException buildDefaultMidStreamChatCompletionError(
         String inferenceEntityId,
         ErrorResponse errorResponse
     ) {
@@ -297,33 +266,18 @@ protected UnifiedChatCompletionException buildDefaultMidStreamChatCompletionErro
         );
     }
 
-    /**
-     * 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) {
+    private static String createErrorType(ErrorResponse errorResponse) {
         return errorResponse != null ? errorResponse.getClass().getSimpleName() : "unknown";
     }
 
-    protected String errorMessage(String message, Request request, ErrorResponse errorResponse, int statusCode) {
+    private static String extractErrorMessage(String message, Request request, ErrorResponse errorResponse, int statusCode) {
         return (errorResponse == null
             || errorResponse.errorStructureFound() == false
             || Strings.isNullOrEmpty(errorResponse.getErrorMessage()))

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

@@ -0,0 +1,33 @@
+/*
+ * 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;
+
+/**
+ * Functional interface for building provider-specific chat completion errors.
+ * This interface is used to create exceptions that are specific to the chat completion service being used.
+ */
+@FunctionalInterface
+public interface ChatCompletionErrorBuilder {
+
+    /**
+     * Builds a provider-specific chat completion error based on the given parameters.
+     *
+     * @param errorResponse The error response received from the service.
+     * @param errorMessage A custom error message to include in the exception.
+     * @param restStatus The HTTP status code associated with the error.
+     * @return An instance of {@link UnifiedChatCompletionException} representing the error.
+     */
+    UnifiedChatCompletionException buildProviderSpecificChatCompletionError(
+        ErrorResponse errorResponse,
+        String errorMessage,
+        RestStatus restStatus
+    );
+}

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

@@ -59,27 +59,22 @@ 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);
+        return buildChatCompletionError(
+            message,
+            request,
+            result,
+            errorResponse,
+            () -> ErrorResponse.class,
+            ElasticInferenceServiceUnifiedChatCompletionResponseHandler::buildProviderSpecificChatCompletionError
+        );
     }
 
-    /**
-     * Builds a custom {@link UnifiedChatCompletionException} for the Elastic Inference Service.
-     * This method is called when an error response is received from the service.
-     *
-     * @param errorResponse The error response received from the service.
-     * @param errorMessage The error message to include in the exception.
-     * @param restStatus The HTTP status of the error response.
-     * @param errorResponseClass The class of the error response.
-     * @return An instance of {@link UnifiedChatCompletionException} with details from the error response.
-     */
-    @Override
-    protected UnifiedChatCompletionException buildChatCompletionError(
-        ErrorResponse errorResponse,
-        String errorMessage,
-        RestStatus restStatus,
-        Class<? extends ErrorResponse> errorResponseClass
+    private static UnifiedChatCompletionException buildProviderSpecificChatCompletionError(
+        ErrorResponse response,
+        String message,
+        RestStatus restStatus
     ) {
-        return new UnifiedChatCompletionException(restStatus, errorMessage, ERROR_TYPE, restStatus.name().toLowerCase(Locale.ROOT));
+        return new UnifiedChatCompletionException(restStatus, message, ERROR_TYPE, restStatus.name().toLowerCase(Locale.ROOT));
     }
 
     /**
@@ -92,73 +87,24 @@ protected UnifiedChatCompletionException buildChatCompletionError(
      * @return An instance of {@link UnifiedChatCompletionException} representing the mid-stream error.
      */
     private UnifiedChatCompletionException buildMidStreamChatCompletionError(String inferenceEntityId, String message, Exception e) {
-        var errorResponse = extractMidStreamChatCompletionErrorResponse(message);
+        var errorResponse = ElasticInferenceServiceErrorResponseEntity.fromString(message);
         // Check if the error response contains a specific structure
         if (errorResponse.errorStructureFound()) {
-            return buildProviderSpecificMidStreamChatCompletionError(inferenceEntityId, errorResponse);
+            return new UnifiedChatCompletionException(
+                RestStatus.INTERNAL_SERVER_ERROR,
+                format(
+                    "%s for request from inference entity id [%s]. Error message: [%s]",
+                    SERVER_ERROR_OBJECT,
+                    inferenceEntityId,
+                    errorResponse.getErrorMessage()
+                ),
+                ERROR_TYPE,
+                STREAM_ERROR
+            );
         } else if (e != null) {
             return UnifiedChatCompletionException.fromThrowable(e);
         } else {
             return buildDefaultMidStreamChatCompletionError(inferenceEntityId, errorResponse);
         }
     }
-
-    /**
-     * Extracts the error response from the message. This method is specific to the Elastic Inference Service
-     * and should parse the message according to its error response format.
-     *
-     * @param message The message containing the error response.
-     * @return An instance of {@link ErrorResponse} parsed from the message.
-     */
-    @Override
-    protected ErrorResponse extractMidStreamChatCompletionErrorResponse(String message) {
-        return ElasticInferenceServiceErrorResponseEntity.fromString(message);
-    }
-
-    /**
-     * Builds a custom mid-stream {@link UnifiedChatCompletionException} for the Elastic Inference Service.
-     * This method is called when a specific error response structure is found in the message.
-     *
-     * @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.
-     */
-    @Override
-    protected UnifiedChatCompletionException buildProviderSpecificMidStreamChatCompletionError(
-        String inferenceEntityId,
-        ErrorResponse errorResponse
-    ) {
-        return new UnifiedChatCompletionException(
-            RestStatus.INTERNAL_SERVER_ERROR,
-            format(
-                "%s for request from inference entity id [%s]. Error message: [%s]",
-                SERVER_ERROR_OBJECT,
-                inferenceEntityId,
-                errorResponse.getErrorMessage()
-            ),
-            ERROR_TYPE,
-            STREAM_ERROR
-        );
-    }
-
-    /**
-     * Builds a default mid-stream {@link UnifiedChatCompletionException} for the Elastic Inference Service.
-     * This method is called when specific error response structure is NOT found in the message.
-     *
-     * @param inferenceEntityId The ID of the inference entity.
-     * @param errorResponse The error response parsed from the message.
-     * @return An instance of {@link UnifiedChatCompletionException} with a generic error message.
-     */
-    @Override
-    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),
-            ERROR_TYPE,
-            STREAM_ERROR
-        );
-    }
 }