[ML] Refactor Inference API chat completion error handling

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.inference.external.http.retry;
 
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
@@ -22,9 +23,9 @@
 public class ChatCompletionErrorResponseHandler {
     private static final String STREAM_ERROR = "stream_error";
 
-    private final UnifiedChatCompletionErrorParser unifiedChatCompletionErrorParser;
+    private final UnifiedChatCompletionErrorParserContract unifiedChatCompletionErrorParser;
 
-    public ChatCompletionErrorResponseHandler(UnifiedChatCompletionErrorParser errorParser) {
+    public ChatCompletionErrorResponseHandler(UnifiedChatCompletionErrorParserContract errorParser) {
         this.unifiedChatCompletionErrorParser = Objects.requireNonNull(errorParser);
     }
 
@@ -49,14 +50,18 @@ private UnifiedChatCompletionException buildChatCompletionErrorInternal(
                 restStatus,
                 errorMessage,
                 errorResponse.type(),
-                errorResponse.code(),
+                code(errorResponse.code(), restStatus),
                 errorResponse.param()
             );
         } else {
             return buildDefaultChatCompletionError(errorResponse, errorMessage, restStatus);
         }
     }
 
+    private static String code(@Nullable String code, RestStatus status) {
+        return code != null ? code : status.name().toLowerCase(Locale.ROOT);
+    }
+
     /**
      * Builds a default {@link UnifiedChatCompletionException} for a streaming request.
      * This method is used when an error response is received we were unable to parse it in the format we were expecting.

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

@@ -9,7 +9,7 @@
 
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 
-public interface UnifiedChatCompletionErrorParser {
+public interface UnifiedChatCompletionErrorParserContract {
     UnifiedChatCompletionErrorResponse parse(HttpResult result);
 
     UnifiedChatCompletionErrorResponse parse(String result);

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

@@ -8,12 +8,51 @@
 package org.elasticsearch.xpack.inference.external.http.retry;
 
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.xcontent.ConstructingObjectParser;
+import org.elasticsearch.xcontent.ParseField;
+import org.elasticsearch.xpack.inference.external.http.HttpResult;
 
 import java.util.Objects;
+import java.util.Optional;
 
 public class UnifiedChatCompletionErrorResponse extends ErrorResponse {
+
+    private static final ConstructingObjectParser<Optional<UnifiedChatCompletionErrorResponse>, Void> CONSTRUCTING_OBJECT_PARSER =
+        new ConstructingObjectParser<>("streaming_error", true, args -> Optional.ofNullable((UnifiedChatCompletionErrorResponse) args[0]));
+    private static final ConstructingObjectParser<UnifiedChatCompletionErrorResponse, Void> ERROR_BODY_PARSER =
+        new ConstructingObjectParser<>(
+            "streaming_error",
+            true,
+            args -> new UnifiedChatCompletionErrorResponse((String) args[0], (String) args[1], (String) args[2], (String) args[3])
+        );
+
+    static {
+        ERROR_BODY_PARSER.declareString(ConstructingObjectParser.constructorArg(), new ParseField("message"));
+        ERROR_BODY_PARSER.declareString(ConstructingObjectParser.constructorArg(), new ParseField("type"));
+        ERROR_BODY_PARSER.declareStringOrNull(ConstructingObjectParser.optionalConstructorArg(), new ParseField("code"));
+        ERROR_BODY_PARSER.declareStringOrNull(ConstructingObjectParser.optionalConstructorArg(), new ParseField("param"));
+
+        CONSTRUCTING_OBJECT_PARSER.declareObjectOrNull(
+            ConstructingObjectParser.optionalConstructorArg(),
+            ERROR_BODY_PARSER,
+            null,
+            new ParseField("error")
+        );
+    }
+
+    public static final UnifiedChatCompletionErrorParserContract ERROR_PARSER = UnifiedChatCompletionErrorResponseUtils
+        .createErrorParserWithObjectParser(CONSTRUCTING_OBJECT_PARSER);
     public static final UnifiedChatCompletionErrorResponse UNDEFINED_ERROR = new UnifiedChatCompletionErrorResponse();
 
+    /**
+     * Standard error response parser. This can be overridden for those subclasses that
+     * have a different error response structure.
+     * @param response The error response as an HttpResult
+     */
+    public static UnifiedChatCompletionErrorResponse fromHttpResult(HttpResult response) {
+        return ERROR_PARSER.parse(response);
+    }
+
     @Nullable
     private final String code;
     @Nullable

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

@@ -0,0 +1,118 @@
+/*
+ * 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.core.CheckedFunction;
+import org.elasticsearch.xcontent.ConstructingObjectParser;
+import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xcontent.XContentParserConfiguration;
+import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.inference.external.http.HttpResult;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Optional;
+import java.util.concurrent.Callable;
+
+public class UnifiedChatCompletionErrorResponseUtils {
+
+    /**
+     * Creates a {@link UnifiedChatCompletionErrorParserContract} that parses the error response as a string.
+     * This is useful for cases where the error response is too complicated to parse.
+     *
+     * @param type The type of the error, used for categorization.
+     * @return A {@link UnifiedChatCompletionErrorParserContract} instance.
+     */
+    public static UnifiedChatCompletionErrorParserContract createErrorParserWithStringify(String type) {
+        return new UnifiedChatCompletionErrorParserContract() {
+            @Override
+            public UnifiedChatCompletionErrorResponse parse(HttpResult result) {
+                try {
+                    String errorMessage = new String(result.body(), StandardCharsets.UTF_8);
+                    return new UnifiedChatCompletionErrorResponse(errorMessage, type, null, null);
+                } catch (Exception e) {
+                    // swallow the error
+                }
+
+                return UnifiedChatCompletionErrorResponse.UNDEFINED_ERROR;
+            }
+
+            @Override
+            public UnifiedChatCompletionErrorResponse parse(String result) {
+                return new UnifiedChatCompletionErrorResponse(result, type, null, null);
+            }
+        };
+    }
+
+    /**
+     * Creates a {@link UnifiedChatCompletionErrorParserContract} that uses a {@link ConstructingObjectParser} to parse the error response.
+     * This is useful for cases where the error response can be parsed into an object.
+     *
+     * @param objectParser The {@link ConstructingObjectParser} to use for parsing the error response.
+     * @return A {@link UnifiedChatCompletionErrorParserContract} instance.
+     */
+    public static UnifiedChatCompletionErrorParserContract createErrorParserWithObjectParser(
+        ConstructingObjectParser<Optional<UnifiedChatCompletionErrorResponse>, Void> objectParser
+    ) {
+        return new UnifiedChatCompletionErrorParser<>((parser) -> objectParser.apply(parser, null));
+    }
+
+    /**
+     * Creates a {@link UnifiedChatCompletionErrorParserContract} that uses a generic parser function to parse the error response.
+     * This is useful for cases where the error response can be parsed using custom logic, typically when parsing from a map.
+     *
+     * @param genericParser The function that takes an {@link XContentParser} and returns an
+     * {@link Optional<UnifiedChatCompletionErrorResponse>}.
+     * @param <E> The type of exception that the parser can throw.
+     * @return A {@link UnifiedChatCompletionErrorParserContract} instance.
+     */
+    public static <E extends Exception> UnifiedChatCompletionErrorParserContract createErrorParserWithGenericParser(
+        CheckedFunction<XContentParser, Optional<UnifiedChatCompletionErrorResponse>, E> genericParser
+    ) {
+        return new UnifiedChatCompletionErrorParser<>(genericParser);
+    }
+
+    private record UnifiedChatCompletionErrorParser<E extends Exception>(
+        CheckedFunction<XContentParser, Optional<UnifiedChatCompletionErrorResponse>, E> genericParser
+    ) implements UnifiedChatCompletionErrorParserContract {
+
+        @Override
+        public UnifiedChatCompletionErrorResponse parse(HttpResult result) {
+            return executeGenericParser(genericParser, createHttpResultXContentParserFunction(result));
+        }
+
+        @Override
+        public UnifiedChatCompletionErrorResponse parse(String result) {
+            return executeGenericParser(genericParser, createStringXContentParserFunction(result));
+        }
+
+    }
+
+    private static Callable<XContentParser> createHttpResultXContentParserFunction(HttpResult response) {
+        return () -> XContentFactory.xContent(XContentType.JSON).createParser(XContentParserConfiguration.EMPTY, response.body());
+    }
+
+    private static Callable<XContentParser> createStringXContentParserFunction(String response) {
+        return () -> XContentFactory.xContent(XContentType.JSON).createParser(XContentParserConfiguration.EMPTY, response);
+    }
+
+    private static <E extends Exception> UnifiedChatCompletionErrorResponse executeGenericParser(
+        CheckedFunction<XContentParser, Optional<UnifiedChatCompletionErrorResponse>, E> genericParser,
+        Callable<XContentParser> createXContentParser
+    ) {
+        try (XContentParser parser = createXContentParser.call()) {
+            return genericParser.apply(parser).orElse(UnifiedChatCompletionErrorResponse.UNDEFINED_ERROR);
+        } catch (Exception e) {
+            // swallow the error
+        }
+
+        return UnifiedChatCompletionErrorResponse.UNDEFINED_ERROR;
+    }
+
+    private UnifiedChatCompletionErrorResponseUtils() {}
+}