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 index 3dac8d849ba6f..668bab50a5f0f 100644 --- 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 errorParseFunction; + private final BiFunction errorParseFunction; private final boolean canHandleStreamingResponses; - public BaseResponseHandler(String requestType, ResponseParser parseFunction, Function errorParseFunction) { + public BaseResponseHandler( + String requestType, + ResponseParser parseFunction, + BiFunction errorParseFunction + ) { this(requestType, parseFunction, errorParseFunction, false); } public BaseResponseHandler( String requestType, ResponseParser parseFunction, - Function errorParseFunction, + BiFunction 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 index 489844f4d14de..20f92f2ab9f54 100644 --- 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 index 9ea687ed52352..28e581991c219 100644 --- 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 * 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 index 003b0f9e87720..27b4d60173a97 100644 --- 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 errorParseFunction, + BiFunction 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 index 9698f577f289d..09b3182ad5b22 100644 --- 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,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; public class CohereErrorResponseEntity extends ErrorResponse { @@ -20,6 +21,10 @@ 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 * 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 index 871840da2f30c..ae982d5e94a80 100644 --- 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 index 37b39680b9351..0060826345ecf 100644 --- 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 * 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 index bf43be2b69d64..1888fbb1d2e93 100644 --- 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 * 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 index 63b5e5622d871..622d8a191b1b4 100644 --- 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,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; public class HuggingFaceErrorResponseEntity extends ErrorResponse { @@ -20,6 +21,10 @@ 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 * 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 index 012283d54be89..fe29d9ff236c2 100644 --- 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 index a1224294de43d..af207e0e71dd7 100644 --- 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,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; public class JinaAIErrorResponseEntity extends ErrorResponse { @@ -20,6 +21,10 @@ 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 index a491cef126971..e7dec494249f3 100644 --- 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 errorParseFunction + BiFunction 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 index e8d0b2096b240..aff775c5123ab 100644 --- 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 errorParseFunction, + BiFunction 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 index efd52a4960f11..ad3b7127f3b46 100644 --- 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 index de1d399112d6e..dd23c8dbfa909 100644 --- 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,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; public class VoyageAIErrorResponseEntity extends ErrorResponse { @@ -20,6 +21,10 @@ private VoyageAIErrorResponseEntity(String errorMessage) { super(errorMessage); } + public static ErrorResponse fromResponse(Request request, HttpResult response) { + return fromResponse(response); + } + /** * Parse an HTTP response into a VoyageAIErrorResponseEntity *