Passing request to the error response handler

jonathan-buttner · jonathan-buttner · commit 5c28ee899c66 · 2025-05-07T16:32:54.000-04:00
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/BaseResponseHandler.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/BaseResponseHandler.java
@@ -17,7 +17,7 @@
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 
 import java.util.Objects;
-import java.util.function.Function;
+import java.util.function.BiFunction;
 
 import static org.elasticsearch.core.Strings.format;
 import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
@@ -37,17 +37,21 @@ public abstract class BaseResponseHandler implements ResponseHandler {
 
     protected final String requestType;
     private final ResponseParser parseFunction;
-    private final Function<HttpResult, ErrorResponse> errorParseFunction;
+    private final BiFunction<Request, HttpResult, ErrorResponse> errorParseFunction;
     private final boolean canHandleStreamingResponses;
 
-    public BaseResponseHandler(String requestType, ResponseParser parseFunction, Function<HttpResult, ErrorResponse> errorParseFunction) {
+    public BaseResponseHandler(
+        String requestType,
+        ResponseParser parseFunction,
+        BiFunction<Request, HttpResult, ErrorResponse> errorParseFunction
+    ) {
         this(requestType, parseFunction, errorParseFunction, false);
     }
 
     public BaseResponseHandler(
         String requestType,
         ResponseParser parseFunction,
-        Function<HttpResult, ErrorResponse> errorParseFunction,
+        BiFunction<Request, HttpResult, ErrorResponse> errorParseFunction,
         boolean canHandleStreamingResponses
     ) {
         this.requestType = Objects.requireNonNull(requestType);
@@ -96,7 +100,7 @@ public void validateResponse(
     protected abstract void checkForFailureStatusCode(Request request, HttpResult result);
 
     private void checkForErrorObject(Request request, HttpResult result) {
-        var errorEntity = errorParseFunction.apply(result);
+        var errorEntity = errorParseFunction.apply(request, result);
 
         if (errorEntity.errorStructureFound()) {
             // We don't really know what happened because the status code was 200 so we'll return a failure and let the
@@ -109,7 +113,7 @@ private void checkForErrorObject(Request request, HttpResult result) {
     }
 
     protected Exception buildError(String message, Request request, HttpResult result) {
-        var errorEntityMsg = errorParseFunction.apply(result);
+        var errorEntityMsg = errorParseFunction.apply(request, result);
         return buildError(message, request, result, errorEntityMsg);
     }
 
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/ErrorMessageResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/ErrorMessageResponseEntity.java
@@ -13,6 +13,7 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 import java.util.Map;
 import java.util.Objects;
@@ -68,4 +69,8 @@ public static ErrorResponse fromResponse(HttpResult response, String defaultMess
     public static ErrorResponse fromResponse(HttpResult response) {
         return fromResponse(response, "");
     }
+
+    public static ErrorResponse fromResponse(Request request, HttpResult response) {
+        return fromResponse(response);
+    }
 }
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/alibabacloudsearch/response/AlibabaCloudSearchErrorResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/alibabacloudsearch/response/AlibabaCloudSearchErrorResponseEntity.java
@@ -15,6 +15,7 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 public class AlibabaCloudSearchErrorResponseEntity extends ErrorResponse {
     private static final Logger logger = LogManager.getLogger(AlibabaCloudSearchErrorResponseEntity.class);
@@ -23,6 +24,10 @@ private AlibabaCloudSearchErrorResponseEntity(String errorMessage) {
         super(errorMessage);
     }
 
+    public static ErrorResponse fromResponse(Request request, HttpResult response) {
+        return fromResponse(response);
+    }
+
     /**
      * An example error response for invalid auth would look like
      * <code>
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/azureopenai/response/AzureMistralOpenAiExternalResponseHandler.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/azureopenai/response/AzureMistralOpenAiExternalResponseHandler.java
@@ -25,7 +25,7 @@
 import org.elasticsearch.xpack.inference.services.openai.OpenAiStreamingProcessor;
 
 import java.util.concurrent.Flow;
-import java.util.function.Function;
+import java.util.function.BiFunction;
 
 import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
 import static org.elasticsearch.xpack.inference.external.http.retry.ResponseHandlerUtils.getFirstHeaderOrUnknown;
@@ -55,7 +55,7 @@ public class AzureMistralOpenAiExternalResponseHandler extends BaseResponseHandl
     public AzureMistralOpenAiExternalResponseHandler(
         String requestType,
         ResponseParser parseFunction,
-        Function<HttpResult, ErrorResponse> errorParseFunction,
+        BiFunction<Request, HttpResult, ErrorResponse> errorParseFunction,
         boolean canHandleStreamingResponses
     ) {
         super(requestType, parseFunction, errorParseFunction);
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/response/CohereErrorResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/response/CohereErrorResponseEntity.java
@@ -13,13 +13,18 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 public class CohereErrorResponseEntity extends ErrorResponse {
 
     private CohereErrorResponseEntity(String errorMessage) {
         super(errorMessage);
     }
 
+    public static ErrorResponse fromResponse(Request request, HttpResult response) {
+        return fromResponse(response);
+    }
+
     /**
      * An example error response for invalid auth would look like
      * <code>
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/response/ElasticInferenceServiceErrorResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/response/ElasticInferenceServiceErrorResponseEntity.java
@@ -16,6 +16,7 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 import java.io.IOException;
 
@@ -37,6 +38,10 @@ private ElasticInferenceServiceErrorResponseEntity(String errorMessage) {
         super(errorMessage);
     }
 
+    public static ErrorResponse fromResponse(Request request, HttpResult response) {
+        return fromResponse(response);
+    }
+
     public static ErrorResponse fromResponse(HttpResult response) {
         return fromParser(
             () -> XContentFactory.xContent(XContentType.JSON).createParser(XContentParserConfiguration.EMPTY, response.body())
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googleaistudio/response/GoogleAiStudioErrorResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googleaistudio/response/GoogleAiStudioErrorResponseEntity.java
@@ -13,6 +13,7 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 import java.util.Map;
 import java.util.Objects;
@@ -23,6 +24,10 @@ private GoogleAiStudioErrorResponseEntity(String errorMessage) {
         super(errorMessage);
     }
 
+    public static ErrorResponse fromResponse(Request request, HttpResult response) {
+        return fromResponse(response);
+    }
+
     /**
      * An example error response for invalid auth would look like
      * <code>
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/response/GoogleVertexAiErrorResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/response/GoogleVertexAiErrorResponseEntity.java
@@ -13,6 +13,7 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 import java.util.Map;
 import java.util.Objects;
@@ -23,6 +24,10 @@ private GoogleVertexAiErrorResponseEntity(String errorMessage) {
         super(errorMessage);
     }
 
+    public static ErrorResponse fromResponse(Request request, HttpResult response) {
+        return fromResponse(response);
+    }
+
     /**
      * An example error response for invalid auth would look like
      * <code>
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/response/HuggingFaceErrorResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/response/HuggingFaceErrorResponseEntity.java
@@ -13,13 +13,18 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 public class HuggingFaceErrorResponseEntity extends ErrorResponse {
 
     public HuggingFaceErrorResponseEntity(String message) {
         super(message);
     }
 
+    public static ErrorResponse fromResponse(Request request, HttpResult response) {
+        return fromResponse(response);
+    }
+
     /**
      * An example error response for invalid auth would look like
      * <code>
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ibmwatsonx/response/IbmWatsonxErrorResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ibmwatsonx/response/IbmWatsonxErrorResponseEntity.java
@@ -13,6 +13,7 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 import java.util.Map;
 import java.util.Objects;
@@ -23,6 +24,10 @@ private IbmWatsonxErrorResponseEntity(String errorMessage) {
         super(errorMessage);
     }
 
+    public static ErrorResponse fromResponse(Request request, HttpResult response) {
+        return fromResponse(response);
+    }
+
     @SuppressWarnings("unchecked")
     public static ErrorResponse fromResponse(HttpResult response) {
         try (
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/jinaai/response/JinaAIErrorResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/jinaai/response/JinaAIErrorResponseEntity.java
@@ -13,13 +13,18 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 public class JinaAIErrorResponseEntity extends ErrorResponse {
 
     private JinaAIErrorResponseEntity(String errorMessage) {
         super(errorMessage);
     }
 
+    public static ErrorResponse fromResponse(Request request, HttpResult response) {
+        return fromResponse(response);
+    }
+
     /**
      * Parse an HTTP response into a JinaAIErrorResponseEntity
      *
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiChatCompletionResponseHandler.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiChatCompletionResponseHandler.java
@@ -13,7 +13,7 @@
 import org.elasticsearch.xpack.inference.external.http.retry.RetryException;
 import org.elasticsearch.xpack.inference.external.request.Request;
 
-import java.util.function.Function;
+import java.util.function.BiFunction;
 
 public class OpenAiChatCompletionResponseHandler extends OpenAiResponseHandler {
     public OpenAiChatCompletionResponseHandler(String requestType, ResponseParser parseFunction) {
@@ -23,7 +23,7 @@ public OpenAiChatCompletionResponseHandler(String requestType, ResponseParser pa
     protected OpenAiChatCompletionResponseHandler(
         String requestType,
         ResponseParser parseFunction,
-        Function<HttpResult, ErrorResponse> errorParseFunction
+        BiFunction<Request, HttpResult, ErrorResponse> errorParseFunction
     ) {
         super(requestType, parseFunction, errorParseFunction, true);
     }
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiResponseHandler.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiResponseHandler.java
@@ -22,7 +22,7 @@
 import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventProcessor;
 
 import java.util.concurrent.Flow;
-import java.util.function.Function;
+import java.util.function.BiFunction;
 
 import static org.elasticsearch.xpack.inference.external.http.retry.ResponseHandlerUtils.getFirstHeaderOrUnknown;
 
@@ -50,7 +50,7 @@ public OpenAiResponseHandler(String requestType, ResponseParser parseFunction, b
     protected OpenAiResponseHandler(
         String requestType,
         ResponseParser parseFunction,
-        Function<HttpResult, ErrorResponse> errorParseFunction,
+        BiFunction<Request, HttpResult, ErrorResponse> errorParseFunction,
         boolean canHandleStreamingResponses
     ) {
         super(requestType, parseFunction, errorParseFunction, canHandleStreamingResponses);
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiUnifiedChatCompletionResponseHandler.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiUnifiedChatCompletionResponseHandler.java
@@ -120,6 +120,10 @@ private static class OpenAiErrorResponse extends ErrorResponse {
             );
         }
 
+        private static ErrorResponse fromResponse(Request request, HttpResult response) {
+            return fromResponse(response);
+        }
+
         private static ErrorResponse fromResponse(HttpResult response) {
             try (
                 XContentParser parser = XContentFactory.xContent(XContentType.JSON)
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/voyageai/response/VoyageAIErrorResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/voyageai/response/VoyageAIErrorResponseEntity.java
@@ -13,13 +13,18 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 public class VoyageAIErrorResponseEntity extends ErrorResponse {
 
     private VoyageAIErrorResponseEntity(String errorMessage) {
         super(errorMessage);
     }
 
+    public static ErrorResponse fromResponse(Request request, HttpResult response) {
+        return fromResponse(response);
+    }
+
     /**
      * Parse an HTTP response into a VoyageAIErrorResponseEntity
      *
