Refactor mid-stream error handling in response handlers

Author: Jan-Kazlouski-elastic

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

@@ -224,23 +224,6 @@ protected UnifiedChatCompletionException buildDefaultChatCompletionError(
         );
     }
 
-    /**
-     * 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.

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

@@ -91,8 +91,7 @@ protected UnifiedChatCompletionException buildChatCompletionError(
      * @param e The exception that occurred, if any.
      * @return An instance of {@link UnifiedChatCompletionException} representing the mid-stream error.
      */
-    @Override
-    public UnifiedChatCompletionException buildMidStreamChatCompletionError(String inferenceEntityId, String message, Exception e) {
+    private UnifiedChatCompletionException buildMidStreamChatCompletionError(String inferenceEntityId, String message, Exception e) {
         var errorResponse = extractMidStreamChatCompletionErrorResponse(message);
         // Check if the error response contains a specific structure
         if (errorResponse.errorStructureFound()) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiUnifiedChatCompletionResponseHandler.java

@@ -83,8 +83,7 @@ protected UnifiedChatCompletionException buildProviderSpecificChatCompletionErro
         );
     }
 
-    @Override
-    public UnifiedChatCompletionException buildMidStreamChatCompletionError(String inferenceEntityId, String message, Exception e) {
+    private UnifiedChatCompletionException buildMidStreamChatCompletionError(String inferenceEntityId, String message, Exception e) {
         return buildMidStreamChatCompletionError(inferenceEntityId, message, e, GoogleVertexAiErrorResponse.class);
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiUnifiedChatCompletionResponseHandler.java

@@ -91,7 +91,6 @@ protected UnifiedChatCompletionException buildProviderSpecificChatCompletionErro
      * @param e the exception that occurred
      * @return an instance of {@link UnifiedChatCompletionException} with details from the error response
      */
-    @Override
     public UnifiedChatCompletionException buildMidStreamChatCompletionError(String inferenceEntityId, String message, Exception e) {
         // Use the custom type StreamingErrorResponse for mid-stream errors
         return buildMidStreamChatCompletionError(inferenceEntityId, message, e, StreamingErrorResponse.class);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/sagemaker/schema/elastic/ElasticCompletionPayload.java

@@ -40,7 +40,14 @@
  * Each chunk should be in a valid JSON format, as that is the format the Elastic API uses.
  */
 public class ElasticCompletionPayload implements SageMakerStreamSchemaPayload, ElasticPayload {
-    private static final XContentParserConfiguration parserConfig = XContentParserConfiguration.EMPTY.withDeprecationHandler(
+    private static final OpenAiUnifiedChatCompletionResponseHandler ERROR_HANDLER = new OpenAiUnifiedChatCompletionResponseHandler(
+        "sagemaker openai chat completion",
+        ((request, result) -> {
+            assert false : "do not call this";
+            throw new UnsupportedOperationException("SageMaker should not call this object's response parser.");
+        })
+    );
+    private static final XContentParserConfiguration PARSER_CONFIGURATION = XContentParserConfiguration.EMPTY.withDeprecationHandler(
         LoggingDeprecationHandler.INSTANCE
     );
 
@@ -94,19 +101,15 @@ public SdkBytes chatCompletionRequestBytes(SageMakerModel model, UnifiedCompleti
     public StreamingUnifiedChatCompletionResults.Results chatCompletionResponseBody(SageMakerModel model, SdkBytes response) {
         var responseData = response.asUtf8String();
         try {
-            var results = OpenAiUnifiedStreamingProcessor.parse(parserConfig, responseData)
+            var results = OpenAiUnifiedStreamingProcessor.parse(PARSER_CONFIGURATION, responseData)
                 .collect(
                     () -> new ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk>(),
                     ArrayDeque::offer,
                     ArrayDeque::addAll
                 );
             return new StreamingUnifiedChatCompletionResults.Results(results);
         } catch (Exception e) {
-            throw new OpenAiUnifiedChatCompletionResponseHandler(null, null).buildMidStreamChatCompletionError(
-                model.getInferenceEntityId(),
-                responseData,
-                e
-            );
+            throw ERROR_HANDLER.buildMidStreamChatCompletionError(model.getInferenceEntityId(), responseData, e);
         }
     }