Merge branch 'main' into remove-yaml-tests

Author: joshua-adams-1

docs/changelog/129848.yaml

@@ -0,0 +1,5 @@
+pr: 129848
+summary: "[ML] Add Azure AI Rerank support to the Inference Plugin"
+area: Machine Learning
+type: enhancement
+issues: []

modules/repository-s3/src/main/resources/org/elasticsearch/repositories/s3/regions_by_endpoint.txt

@@ -6,6 +6,10 @@ ap-east-1          s3-fips.ap-east-1.amazonaws.com
 ap-east-1          s3-fips.dualstack.ap-east-1.amazonaws.com
 ap-east-1          s3.ap-east-1.amazonaws.com
 ap-east-1          s3.dualstack.ap-east-1.amazonaws.com
+ap-east-2          s3-fips.ap-east-2.amazonaws.com
+ap-east-2          s3-fips.dualstack.ap-east-2.amazonaws.com
+ap-east-2          s3.ap-east-2.amazonaws.com
+ap-east-2          s3.dualstack.ap-east-2.amazonaws.com
 ap-northeast-1     s3-fips.ap-northeast-1.amazonaws.com
 ap-northeast-1     s3-fips.dualstack.ap-northeast-1.amazonaws.com
 ap-northeast-1     s3.ap-northeast-1.amazonaws.com
@@ -56,6 +60,14 @@ aws-iso-b-global   s3-fips.aws-iso-b-global.sc2s.sgov.gov
 aws-iso-b-global   s3-fips.dualstack.aws-iso-b-global.sc2s.sgov.gov
 aws-iso-b-global   s3.aws-iso-b-global.sc2s.sgov.gov
 aws-iso-b-global   s3.dualstack.aws-iso-b-global.sc2s.sgov.gov
+aws-iso-e-global   s3-fips.aws-iso-e-global.cloud.adc-e.uk
+aws-iso-e-global   s3-fips.dualstack.aws-iso-e-global.cloud.adc-e.uk
+aws-iso-e-global   s3.aws-iso-e-global.cloud.adc-e.uk
+aws-iso-e-global   s3.dualstack.aws-iso-e-global.cloud.adc-e.uk
+aws-iso-f-global   s3-fips.aws-iso-f-global.csp.hci.ic.gov
+aws-iso-f-global   s3-fips.dualstack.aws-iso-f-global.csp.hci.ic.gov
+aws-iso-f-global   s3.aws-iso-f-global.csp.hci.ic.gov
+aws-iso-f-global   s3.dualstack.aws-iso-f-global.csp.hci.ic.gov
 aws-iso-global     s3-fips.aws-iso-global.c2s.ic.gov
 aws-iso-global     s3-fips.dualstack.aws-iso-global.c2s.ic.gov
 aws-iso-global     s3.aws-iso-global.c2s.ic.gov
@@ -76,6 +88,10 @@ cn-north-1         s3.cn-north-1.amazonaws.com.cn
 cn-north-1         s3.dualstack.cn-north-1.amazonaws.com.cn
 cn-northwest-1     s3.cn-northwest-1.amazonaws.com.cn
 cn-northwest-1     s3.dualstack.cn-northwest-1.amazonaws.com.cn
+eusc-de-east-1     s3-fips.eusc-de-east-1.amazonaws.eu
+eusc-de-east-1     s3-fips.dualstack.eusc-de-east-1.amazonaws.eu
+eusc-de-east-1     s3.eusc-de-east-1.amazonaws.eu
+eusc-de-east-1     s3.dualstack.eusc-de-east-1.amazonaws.eu
 eu-central-1       s3-fips.dualstack.eu-central-1.amazonaws.com
 eu-central-1       s3-fips.eu-central-1.amazonaws.com
 eu-central-1       s3.dualstack.eu-central-1.amazonaws.com

modules/repository-s3/src/test/java/org/elasticsearch/repositories/s3/RegionFromEndpointGuesserTests.java

@@ -9,6 +9,11 @@
 
 package org.elasticsearch.repositories.s3;
 
+import software.amazon.awssdk.endpoints.Endpoint;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.endpoints.S3EndpointParams;
+import software.amazon.awssdk.services.s3.endpoints.internal.DefaultS3EndpointProvider;
+
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.test.ESTestCase;
 
@@ -23,6 +28,14 @@ public void testRegionGuessing() {
         assertRegionGuess("random.endpoint.internal.net", null);
     }
 
+    public void testHasEntryForEachRegion() {
+        final var defaultS3EndpointProvider = new DefaultS3EndpointProvider();
+        for (var region : Region.regions()) {
+            final Endpoint endpoint = safeGet(defaultS3EndpointProvider.resolveEndpoint(S3EndpointParams.builder().region(region).build()));
+            assertNotNull(region.id(), RegionFromEndpointGuesser.guessRegion(endpoint.url().toString()));
+        }
+    }
+
     private static void assertRegionGuess(String endpoint, @Nullable String expectedRegion) {
         assertEquals(endpoint, expectedRegion, RegionFromEndpointGuesser.guessRegion(endpoint));
     }

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -341,6 +341,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion LOOKUP_JOIN_CCS = def(9_120_0_00);
     public static final TransportVersion NODE_USAGE_STATS_FOR_THREAD_POOLS_IN_CLUSTER_INFO = def(9_121_0_00);
     public static final TransportVersion ESQL_CATEGORIZE_OPTIONS = def(9_122_0_00);
+    public static final TransportVersion ML_INFERENCE_AZURE_AI_STUDIO_RERANK_ADDED = def(9_123_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/action/admin/cluster/allocation/ClusterAllocationExplainRequest.java

@@ -236,6 +236,7 @@ public String toString() {
             }
         }
         sb.append(",").append(INCLUDE_YES_DECISIONS_PARAMETER_NAME).append("?=").append(includeYesDecisions);
+        sb.append(",").append(INCLUDE_DISK_INFO_PARAMETER_NAME).append("?=").append(includeDiskInfo);
         return sb.toString();
     }
 

server/src/test/java/org/elasticsearch/action/admin/cluster/allocation/ClusterAllocationExplainRequestTests.java

@@ -37,4 +37,66 @@ public void testSerialization() throws Exception {
         assertEquals(request.getCurrentNode(), actual.getCurrentNode());
     }
 
+    public void testToStringWithEmptyBody() {
+        ClusterAllocationExplainRequest clusterAllocationExplainRequest = new ClusterAllocationExplainRequest(randomTimeValue());
+        clusterAllocationExplainRequest.includeYesDecisions(true);
+        clusterAllocationExplainRequest.includeDiskInfo(false);
+
+        String expected = "ClusterAllocationExplainRequest[useAnyUnassignedShard=true,"
+            + "include_yes_decisions?=true,include_disk_info?=false";
+        assertEquals(expected, clusterAllocationExplainRequest.toString());
+    }
+
+    public void testToStringWithValidBodyButCurrentNodeIsNull() {
+        String index = "test-index";
+        int shard = randomInt();
+        boolean primary = randomBoolean();
+        ClusterAllocationExplainRequest clusterAllocationExplainRequest = new ClusterAllocationExplainRequest(
+            randomTimeValue(),
+            index,
+            shard,
+            primary,
+            null
+        );
+        clusterAllocationExplainRequest.includeYesDecisions(false);
+        clusterAllocationExplainRequest.includeDiskInfo(true);
+
+        String expected = "ClusterAllocationExplainRequest[index="
+            + index
+            + ",shard="
+            + shard
+            + ",primary?="
+            + primary
+            + ",include_yes_decisions?=false"
+            + ",include_disk_info?=true";
+        assertEquals(expected, clusterAllocationExplainRequest.toString());
+    }
+
+    public void testToStringWithAllBodyParameters() {
+        String index = "test-index";
+        int shard = randomInt();
+        boolean primary = randomBoolean();
+        String currentNode = "current_node";
+        ClusterAllocationExplainRequest clusterAllocationExplainRequest = new ClusterAllocationExplainRequest(
+            randomTimeValue(),
+            index,
+            shard,
+            primary,
+            currentNode
+        );
+        clusterAllocationExplainRequest.includeYesDecisions(false);
+        clusterAllocationExplainRequest.includeDiskInfo(true);
+
+        String expected = "ClusterAllocationExplainRequest[index="
+            + index
+            + ",shard="
+            + shard
+            + ",primary?="
+            + primary
+            + ",current_node="
+            + currentNode
+            + ",include_yes_decisions?=false"
+            + ",include_disk_info?=true";
+        assertEquals(expected, clusterAllocationExplainRequest.toString());
+    }
 }

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java

@@ -111,6 +111,7 @@ public void testGetServicesWithRerankTaskType() throws IOException {
             containsInAnyOrder(
                 List.of(
                     "alibabacloud-ai-search",
+                    "azureaistudio",
                     "cohere",
                     "elasticsearch",
                     "googlevertexai",

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceNamedWriteablesProvider.java

@@ -50,6 +50,8 @@
 import org.elasticsearch.xpack.inference.services.azureaistudio.completion.AzureAiStudioChatCompletionTaskSettings;
 import org.elasticsearch.xpack.inference.services.azureaistudio.embeddings.AzureAiStudioEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.azureaistudio.embeddings.AzureAiStudioEmbeddingsTaskSettings;
+import org.elasticsearch.xpack.inference.services.azureaistudio.rerank.AzureAiStudioRerankServiceSettings;
+import org.elasticsearch.xpack.inference.services.azureaistudio.rerank.AzureAiStudioRerankTaskSettings;
 import org.elasticsearch.xpack.inference.services.azureopenai.AzureOpenAiSecretSettings;
 import org.elasticsearch.xpack.inference.services.azureopenai.completion.AzureOpenAiCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.azureopenai.completion.AzureOpenAiCompletionTaskSettings;
@@ -306,6 +308,17 @@ private static void addAzureAiStudioNamedWriteables(List<NamedWriteableRegistry.
                 AzureAiStudioChatCompletionTaskSettings::new
             )
         );
+
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                AzureAiStudioRerankServiceSettings.NAME,
+                AzureAiStudioRerankServiceSettings::new
+            )
+        );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(TaskSettings.class, AzureAiStudioRerankTaskSettings.NAME, AzureAiStudioRerankTaskSettings::new)
+        );
     }
 
     private static void addAzureOpenAiNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {

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

@@ -95,7 +95,7 @@ public void validateResponse(
 
     protected abstract void checkForFailureStatusCode(Request request, HttpResult result);
 
-    private void checkForErrorObject(Request request, HttpResult result) {
+    protected void checkForErrorObject(Request request, HttpResult result) {
         var errorEntity = errorParseFunction.apply(result);
 
         if (errorEntity.errorStructureFound()) {
@@ -116,12 +116,12 @@ protected Exception buildError(String message, Request request, HttpResult resul
     protected Exception buildError(String message, Request request, HttpResult result, ErrorResponse errorResponse) {
         var responseStatusCode = result.response().getStatusLine().getStatusCode();
         return new ElasticsearchStatusException(
-            errorMessage(message, request, result, errorResponse, responseStatusCode),
+            constructErrorMessage(message, request, errorResponse, responseStatusCode),
             toRestStatus(responseStatusCode)
         );
     }
 
-    protected String errorMessage(String message, Request request, HttpResult result, ErrorResponse errorResponse, int statusCode) {
+    public static String constructErrorMessage(String message, Request request, ErrorResponse errorResponse, int statusCode) {
         return (errorResponse == null
             || errorResponse.errorStructureFound() == false
             || Strings.isNullOrEmpty(errorResponse.getErrorMessage()))

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

@@ -0,0 +1,162 @@
+/*
+ * 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.rest.RestStatus;
+import org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
+import org.elasticsearch.xpack.inference.external.http.HttpResult;
+import org.elasticsearch.xpack.inference.external.request.Request;
+
+import java.util.Locale;
+import java.util.Objects;
+
+import static org.elasticsearch.core.Strings.format;
+import static org.elasticsearch.xpack.inference.external.http.retry.BaseResponseHandler.SERVER_ERROR_OBJECT;
+import static org.elasticsearch.xpack.inference.external.http.retry.BaseResponseHandler.toRestStatus;
+
+public class ChatCompletionErrorResponseHandler {
+    private static final String STREAM_ERROR = "stream_error";
+
+    private final UnifiedChatCompletionErrorParser unifiedChatCompletionErrorParser;
+
+    public ChatCompletionErrorResponseHandler(UnifiedChatCompletionErrorParser errorParser) {
+        this.unifiedChatCompletionErrorParser = Objects.requireNonNull(errorParser);
+    }
+
+    public void checkForErrorObject(Request request, HttpResult result) {
+        var errorEntity = unifiedChatCompletionErrorParser.parse(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
+            // client retry if necessary
+            // If we did want to retry here, we'll need to determine if this was a streaming request, if it was
+            // we shouldn't retry because that would replay the entire streaming request and the client would get
+            // duplicate chunks back
+            throw new RetryException(false, buildChatCompletionErrorInternal(SERVER_ERROR_OBJECT, request, result, errorEntity));
+        }
+    }
+
+    public UnifiedChatCompletionException buildChatCompletionError(String message, Request request, HttpResult result) {
+        var errorResponse = unifiedChatCompletionErrorParser.parse(result);
+        return buildChatCompletionErrorInternal(message, request, result, errorResponse);
+    }
+
+    private UnifiedChatCompletionException buildChatCompletionErrorInternal(
+        String message,
+        Request request,
+        HttpResult result,
+        UnifiedChatCompletionErrorResponse errorResponse
+    ) {
+        assert request.isStreaming() : "Only streaming requests support this format";
+        var statusCode = result.response().getStatusLine().getStatusCode();
+        var errorMessage = BaseResponseHandler.constructErrorMessage(message, request, errorResponse, statusCode);
+        var restStatus = toRestStatus(statusCode);
+
+        if (errorResponse.errorStructureFound()) {
+            return new UnifiedChatCompletionException(
+                restStatus,
+                errorMessage,
+                errorResponse.type(),
+                errorResponse.code(),
+                errorResponse.param()
+            );
+        } else {
+            return buildDefaultChatCompletionError(errorResponse, errorMessage, restStatus);
+        }
+    }
+
+    /**
+     * 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.
+     * Only streaming requests should use this method.
+     *
+     * @param errorResponse the error response extracted from the HTTP result
+     * @param errorMessage the error message to include in the exception
+     * @param restStatus the REST status code of the response
+     * @return an instance of {@link UnifiedChatCompletionException} with details from the error response
+     */
+    private static UnifiedChatCompletionException buildDefaultChatCompletionError(
+        ErrorResponse errorResponse,
+        String errorMessage,
+        RestStatus restStatus
+    ) {
+        return new UnifiedChatCompletionException(
+            restStatus,
+            errorMessage,
+            createErrorType(errorResponse),
+            restStatus.name().toLowerCase(Locale.ROOT)
+        );
+    }
+
+    /**
+     * Builds a mid-stream error for a streaming request.
+     * This method is used when an error occurs while processing a streaming response.
+     * 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) {
+        var error = unifiedChatCompletionErrorParser.parse(message);
+
+        if (error.errorStructureFound()) {
+            return new UnifiedChatCompletionException(
+                RestStatus.INTERNAL_SERVER_ERROR,
+                format(
+                    "%s for request from inference entity id [%s]. Error message: [%s]",
+                    SERVER_ERROR_OBJECT,
+                    inferenceEntityId,
+                    error.getErrorMessage()
+                ),
+                error.type(),
+                error.code(),
+                error.param()
+            );
+        } else if (e != null) {
+            // If the error response does not match, we can still return an exception based on the original throwable
+            return UnifiedChatCompletionException.fromThrowable(e);
+        } else {
+            // If no specific error response is found, we return a default mid-stream error
+            return buildDefaultMidStreamChatCompletionError(inferenceEntityId, error);
+        }
+    }
+
+    /**
+     * Builds a default mid-stream error for a streaming request.
+     * This method is used when no specific error response is found in the message.
+     * Only streaming requests should use this method.
+     *
+     * @param inferenceEntityId the ID of the inference entity
+     * @param errorResponse     the error response extracted from the message
+     * @return a {@link UnifiedChatCompletionException} representing the default mid-stream error
+     */
+    private static UnifiedChatCompletionException buildDefaultMidStreamChatCompletionError(
+        String inferenceEntityId,
+        ErrorResponse errorResponse
+    ) {
+        return new UnifiedChatCompletionException(
+            RestStatus.INTERNAL_SERVER_ERROR,
+            format("%s for request from inference entity id [%s]", SERVER_ERROR_OBJECT, inferenceEntityId),
+            createErrorType(errorResponse),
+            STREAM_ERROR
+        );
+    }
+
+    /**
+     * Creates a string representation of the error type based on the provided ErrorResponse.
+     * This method is used to generate a human-readable error type for logging or exception messages.
+     *
+     * @param errorResponse the ErrorResponse object
+     * @return a string representing the error type
+     */
+    private static String createErrorType(ErrorResponse errorResponse) {
+        return errorResponse != null ? errorResponse.getClass().getSimpleName() : "unknown";
+    }
+}