Refactor error handling and add unit tests for Ai21 service

Author: Jan-Kazlouski-elastic

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

@@ -68,4 +68,20 @@ public static ErrorResponse fromResponse(HttpResult response) {
 
         return ErrorResponse.UNDEFINED_ERROR;
     }
+
+    /**
+     * Parses a string response into an ErrorResponse.
+     * If the string is not blank, creates a new ErrorResponse with the string as the error message.
+     * If the string is blank, returns UNDEFINED_ERROR.
+     *
+     * @param response the error response as a string
+     * @return an ErrorResponse instance
+     */
+    public static ErrorResponse fromString(String response) {
+        if (Objects.nonNull(response) && response.isBlank() == false) {
+            return new ErrorResponse(response);
+        } else {
+            return ErrorResponse.UNDEFINED_ERROR;
+        }
+    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ai21/completion/Ai21ChatCompletionResponseHandler.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.inference.services.ai21.completion;
 
+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.http.retry.ErrorResponse;
@@ -16,6 +17,8 @@
 
 import java.util.Locale;
 
+import static org.elasticsearch.core.Strings.format;
+
 /**
  * Handles streaming chat completion responses and error parsing for Mistral inference endpoints.
  * Adapts the OpenAI handler to support Mistral's error schema.
@@ -36,4 +39,30 @@ protected Exception buildError(String message, Request request, HttpResult resul
         var restStatus = toRestStatus(responseStatusCode);
         return new UnifiedChatCompletionException(restStatus, errorMessage, AI_21_ERROR, restStatus.name().toLowerCase(Locale.ROOT));
     }
+
+    protected Exception buildMidStreamError(Request request, String message, Exception e) {
+        var errorResponse = ErrorResponse.fromString(message);
+        if (errorResponse.errorStructureFound()) {
+            return new UnifiedChatCompletionException(
+                RestStatus.INTERNAL_SERVER_ERROR,
+                format(
+                    "%s for request from inference entity id [%s]. Error message: [%s]",
+                    SERVER_ERROR_OBJECT,
+                    request.getInferenceEntityId(),
+                    errorResponse.getErrorMessage()
+                ),
+                AI_21_ERROR,
+                "stream_error"
+            );
+        } else if (e != null) {
+            return UnifiedChatCompletionException.fromThrowable(e);
+        } else {
+            return new UnifiedChatCompletionException(
+                RestStatus.INTERNAL_SERVER_ERROR,
+                format("%s for request from inference entity id [%s]", SERVER_ERROR_OBJECT, request.getInferenceEntityId()),
+                createErrorType(errorResponse),
+                "stream_error"
+            );
+        }
+    }
 }

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

@@ -149,7 +149,7 @@ public static ErrorResponse fromResponse(HttpResult response) {
             }
         }
 
-        static ErrorResponse fromString(String response) {
+        public static ErrorResponse fromString(String response) {
             try (
                 XContentParser parser = XContentFactory.xContent(XContentType.JSON)
                     .createParser(XContentParserConfiguration.EMPTY, response)

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceChatCompletionResponseHandler.java

@@ -141,7 +141,7 @@ private static class StreamingHuggingFaceErrorResponseEntity extends ErrorRespon
          * @param response the raw JSON string representing an error
          * @return a parsed {@link ErrorResponse} or {@link ErrorResponse#UNDEFINED_ERROR} if parsing fails
          */
-        private static ErrorResponse fromString(String response) {
+        public static ErrorResponse fromString(String response) {
             try (
                 XContentParser parser = XContentFactory.xContent(XContentType.JSON)
                     .createParser(XContentParserConfiguration.EMPTY, response)

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/AbstractInferenceServiceTests.java

@@ -351,11 +351,17 @@ public void testParsePersistedConfigWithSecrets_ThrowsUnsupportedModelType() thr
 
             assertThat(
                 exception.getMessage(),
-                containsString(Strings.format("service does not support task type [%s]", parseConfigTestConfig.unsupportedTaskType))
+                containsString(
+                    Strings.format(fetchPersistedConfigTaskTypeParsingErrorMessageFormat(), parseConfigTestConfig.unsupportedTaskType)
+                )
             );
         }
     }
 
+    protected String fetchPersistedConfigTaskTypeParsingErrorMessageFormat() {
+        return "service does not support task type [%s]";
+    }
+
     public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExistsInConfig() throws IOException {
         var parseConfigTestConfig = testConfiguration.commonConfig;
 
@@ -374,7 +380,7 @@ public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExists
                 persistedConfigMap.secrets()
             );
 
-            parseConfigTestConfig.assertModel(model, TaskType.TEXT_EMBEDDING);
+            parseConfigTestConfig.assertModel(model, parseConfigTestConfig.taskType);
         }
     }
 
@@ -396,7 +402,7 @@ public void testParsePersistedConfigWithSecrets_ThrowsWhenAnExtraKeyExistsInServ
                 persistedConfigMap.secrets()
             );
 
-            parseConfigTestConfig.assertModel(model, TaskType.TEXT_EMBEDDING);
+            parseConfigTestConfig.assertModel(model, parseConfigTestConfig.taskType);
         }
     }
 
@@ -413,7 +419,7 @@ public void testParsePersistedConfigWithSecrets_ThrowsWhenAnExtraKeyExistsInTask
 
             var model = service.parsePersistedConfigWithSecrets("id", parseConfigTestConfig.taskType, config.config(), config.secrets());
 
-            parseConfigTestConfig.assertModel(model, TaskType.TEXT_EMBEDDING);
+            parseConfigTestConfig.assertModel(model, parseConfigTestConfig.taskType);
         }
     }
 
@@ -430,7 +436,7 @@ public void testParsePersistedConfigWithSecrets_ThrowsWhenAnExtraKeyExistsInSecr
 
             var model = service.parsePersistedConfigWithSecrets("id", parseConfigTestConfig.taskType, config.config(), config.secrets());
 
-            parseConfigTestConfig.assertModel(model, TaskType.TEXT_EMBEDDING);
+            parseConfigTestConfig.assertModel(model, parseConfigTestConfig.taskType);
         }
     }