Merge branch 'main' into fix_sort_tests

Author: mayya-sharipova

docs/changelog/128538.yaml

@@ -0,0 +1,5 @@
+pr: 128538
+summary: "Added Mistral Chat Completion support to the Inference Plugin"
+area: Machine Learning
+type: enhancement
+issues: []

muted-tests.yml

@@ -204,6 +204,9 @@ tests:
 - class: org.elasticsearch.packaging.test.DockerTests
   method: test151MachineDependentHeapWithSizeOverride
   issue: https://github.com/elastic/elasticsearch/issues/123437
+- class: org.elasticsearch.action.admin.cluster.node.tasks.CancellableTasksIT
+  method: testChildrenTasksCancelledOnTimeout
+  issue: https://github.com/elastic/elasticsearch/issues/123568
 - class: org.elasticsearch.xpack.searchablesnapshots.FrozenSearchableSnapshotsIntegTests
   method: testCreateAndRestorePartialSearchableSnapshot
   issue: https://github.com/elastic/elasticsearch/issues/123773

server/src/internalClusterTest/java/org/elasticsearch/action/admin/cluster/node/tasks/CancellableTasksIT.java

@@ -17,7 +17,6 @@
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.action.ActionType;
-import org.elasticsearch.action.DelegatingActionListener;
 import org.elasticsearch.action.LatchedActionListener;
 import org.elasticsearch.action.LegacyActionRequest;
 import org.elasticsearch.action.admin.cluster.node.tasks.list.ListTasksResponse;
@@ -27,7 +26,6 @@
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.node.DiscoveryNode;
-import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.util.CollectionUtils;
@@ -36,8 +34,6 @@
 import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.injection.guice.Inject;
-import org.elasticsearch.logging.LogManager;
-import org.elasticsearch.logging.Logger;
 import org.elasticsearch.plugins.ActionPlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.tasks.CancellableTask;
@@ -47,7 +43,6 @@
 import org.elasticsearch.tasks.TaskInfo;
 import org.elasticsearch.tasks.TaskManager;
 import org.elasticsearch.test.ESIntegTestCase;
-import org.elasticsearch.test.junit.annotations.TestIssueLogging;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.ReceiveTimeoutTransportException;
 import org.elasticsearch.transport.SendRequestTransportException;
@@ -79,9 +74,6 @@
 @ESIntegTestCase.ClusterScope(scope = ESIntegTestCase.Scope.TEST)
 public class CancellableTasksIT extends ESIntegTestCase {
 
-    // Temporary addition for investigation into https://github.com/elastic/elasticsearch/issues/123568
-    private static final Logger logger = LogManager.getLogger(CancellableTasksIT.class);
-
     static int idGenerator = 0;
     static final Map<TestRequest, CountDownLatch> beforeSendLatches = ConcurrentCollections.newConcurrentMap();
     static final Map<TestRequest, CountDownLatch> arrivedLatches = ConcurrentCollections.newConcurrentMap();
@@ -374,42 +366,18 @@ public void testRemoveBanParentsOnDisconnect() throws Exception {
         }
     }
 
-    @TestIssueLogging(
-        issueUrl = "https://github.com/elastic/elasticsearch/issues/123568",
-        value = "org.elasticsearch.transport.TransportService.tracer:TRACE"
-            + ",org.elasticsearch.tasks.TaskManager:TRACE"
-            + ",org.elasticsearch.action.admin.cluster.node.tasks.CancellableTasksIT:DEBUG"
-    )
     public void testChildrenTasksCancelledOnTimeout() throws Exception {
         Set<DiscoveryNode> nodes = clusterService().state().nodes().stream().collect(Collectors.toSet());
         final TestRequest rootRequest = generateTestRequest(nodes, 0, between(1, 4), true);
-        logger.info("generated request\n{}", buildTestRequestDescription(rootRequest, "", new StringBuilder()).toString());
         ActionFuture<TestResponse> rootTaskFuture = client().execute(TransportTestAction.ACTION, rootRequest);
-        logger.info("action executed");
         allowEntireRequest(rootRequest);
-        logger.info("execution released");
         waitForRootTask(rootTaskFuture, true);
-        logger.info("root task completed");
         ensureBansAndCancellationsConsistency();
-        logger.info("ensureBansAndCancellationsConsistency completed");
 
         // Make sure all descendent requests have completed
         for (TestRequest subRequest : rootRequest.descendants()) {
-            logger.info("awaiting completion of request {}", subRequest.getDescription());
             safeAwait(completedLatches.get(subRequest));
         }
-        logger.info("all requests completed");
-    }
-
-    // Temporary addition for investigation into https://github.com/elastic/elasticsearch/issues/123568
-    static StringBuilder buildTestRequestDescription(TestRequest request, String prefix, StringBuilder stringBuilder) {
-        stringBuilder.append(prefix)
-            .append(Strings.format("id=%d [timeout=%s] %s", request.id, request.timeout, request.node.descriptionWithoutAttributes()))
-            .append('\n');
-        for (TestRequest subRequest : request.subRequests) {
-            buildTestRequestDescription(subRequest, prefix + "  ", stringBuilder);
-        }
-        return stringBuilder;
     }
 
     static TaskId getRootTaskId(TestRequest request) throws Exception {
@@ -538,8 +506,6 @@ public void writeTo(StreamOutput out) throws IOException {
 
     public static class TransportTestAction extends HandledTransportAction<TestRequest, TestResponse> {
 
-        private static final Logger logger = CancellableTasksIT.logger;
-
         public static ActionType<TestResponse> ACTION = new ActionType<>("internal::test_action");
         private final TransportService transportService;
         private final NodeClient client;
@@ -599,22 +565,7 @@ protected void doExecute(Task task, TestRequest request, ActionListener<TestResp
         protected void startSubTask(TaskId parentTaskId, TestRequest subRequest, ActionListener<TestResponse> listener) {
             subRequest.setParentTask(parentTaskId);
             CountDownLatch completeLatch = completedLatches.get(subRequest);
-            ActionListener<TestResponse> latchedListener = new DelegatingActionListener<>(
-                new LatchedActionListener<>(listener, completeLatch)
-            ) {
-                // Temporary logging addition for investigation into https://github.com/elastic/elasticsearch/issues/123568
-                @Override
-                public void onResponse(TestResponse testResponse) {
-                    logger.debug("processing successful response to request [{}]", subRequest.getDescription());
-                    delegate.onResponse(testResponse);
-                }
-
-                @Override
-                public void onFailure(Exception e) {
-                    logger.debug("processing exceptional response to request [{}]: {}", subRequest.getDescription(), e.getMessage());
-                    super.onFailure(e);
-                }
-            };
+            LatchedActionListener<TestResponse> latchedListener = new LatchedActionListener<>(listener, completeLatch);
             transportService.getThreadPool().generic().execute(new AbstractRunnable() {
                 @Override
                 public void onFailure(Exception e) {
@@ -645,13 +596,6 @@ protected void doRun() throws Exception {
                                 TransportResponseHandler.TRANSPORT_WORKER
                             )
                         );
-                        // Temporary addition for investigation into https://github.com/elastic/elasticsearch/issues/123568
-                        logger.debug(
-                            "sent test request [{}] from [{}] to [{}]",
-                            subRequest.getDescription(),
-                            client.getLocalNodeId(),
-                            subRequest.node.descriptionWithoutAttributes()
-                        );
                     }
                 }
             });

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

@@ -189,6 +189,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion DATA_STREAM_OPTIONS_API_REMOVE_INCLUDE_DEFAULTS_8_19 = def(8_841_0_41);
     public static final TransportVersion JOIN_ON_ALIASES_8_19 = def(8_841_0_42);
     public static final TransportVersion ILM_ADD_SKIP_SETTING_8_19 = def(8_841_0_43);
+    public static final TransportVersion ML_INFERENCE_MISTRAL_CHAT_COMPLETION_ADDED_8_19 = def(8_841_0_44);
     public static final TransportVersion V_9_0_0 = def(9_000_0_09);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_1 = def(9_000_0_10);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_2 = def(9_000_0_11);
@@ -282,6 +283,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion IDP_CUSTOM_SAML_ATTRIBUTES = def(9_087_0_00);
     public static final TransportVersion JOIN_ON_ALIASES = def(9_088_0_00);
     public static final TransportVersion ILM_ADD_SKIP_SETTING = def(9_089_0_00);
+    public static final TransportVersion ML_INFERENCE_MISTRAL_CHAT_COMPLETION_ADDED = def(9_090_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/inference/UnifiedCompletionRequest.java

@@ -78,6 +78,14 @@ public record UnifiedCompletionRequest(
      * {@link #MAX_COMPLETION_TOKENS_FIELD}. Providers are expected to pass in their supported field name.
      */
     private static final String MAX_TOKENS_PARAM = "max_tokens_field";
+    /**
+     * Indicates whether to include the `stream_options` field in the JSON output.
+     * Some providers do not support this field. In such cases, this parameter should be set to "false",
+     * and the `stream_options` field will be excluded from the output.
+     * For providers that do support stream options, this parameter is left unset (default behavior),
+     * which implicitly includes the `stream_options` field in the output.
+     */
+    public static final String INCLUDE_STREAM_OPTIONS_PARAM = "include_stream_options";
 
     /**
      * Creates a {@link org.elasticsearch.xcontent.ToXContent.Params} that causes ToXContent to include the key values:
@@ -91,6 +99,23 @@ public static Params withMaxTokens(String modelId, Params params) {
         );
     }
 
+    /**
+     * Creates a {@link org.elasticsearch.xcontent.ToXContent.Params} that causes ToXContent to include the key values:
+     * - Key: {@link #MODEL_FIELD}, Value: modelId
+     * - Key: {@link #MAX_TOKENS_FIELD}, Value: {@link #MAX_TOKENS_FIELD}
+     * - Key: {@link #INCLUDE_STREAM_OPTIONS_PARAM}, Value: "false"
+     */
+    public static Params withMaxTokensAndSkipStreamOptionsField(String modelId, Params params) {
+        return new DelegatingMapParams(
+            Map.ofEntries(
+                Map.entry(MODEL_ID_PARAM, modelId),
+                Map.entry(MAX_TOKENS_PARAM, MAX_TOKENS_FIELD),
+                Map.entry(INCLUDE_STREAM_OPTIONS_PARAM, Boolean.FALSE.toString())
+            ),
+            params
+        );
+    }
+
     /**
      * Creates a {@link org.elasticsearch.xcontent.ToXContent.Params} that causes ToXContent to include the key values:
      * - Key: {@link #MODEL_FIELD}, Value: modelId

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushTopNToSourceTests.java

@@ -193,7 +193,6 @@ public void testSimpleSortEvalSumLiteralAndField() {
         assertNoPushdownSort(query.asTimeSeries(), "for time series index mode");
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/114515")
     public void testPartiallyPushableSort() {
         // FROM index | EVAL sum = 1 + integer | SORT integer, sum, field | LIMIT 10
         var query = from("index").eval("sum", b -> b.add(b.i(1), b.field("integer"))).sort("integer").sort("sum").sort("field").limit(10);

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

@@ -134,7 +134,7 @@ public void testGetServicesWithRerankTaskType() throws IOException {
 
     public void testGetServicesWithCompletionTaskType() throws IOException {
         List<Object> services = getServices(TaskType.COMPLETION);
-        assertThat(services.size(), equalTo(13));
+        assertThat(services.size(), equalTo(14));
 
         var providers = providers(services);
 
@@ -154,15 +154,16 @@ public void testGetServicesWithCompletionTaskType() throws IOException {
                     "openai",
                     "streaming_completion_test_service",
                     "hugging_face",
-                    "amazon_sagemaker"
+                    "amazon_sagemaker",
+                    "mistral"
                 ).toArray()
             )
         );
     }
 
     public void testGetServicesWithChatCompletionTaskType() throws IOException {
         List<Object> services = getServices(TaskType.CHAT_COMPLETION);
-        assertThat(services.size(), equalTo(7));
+        assertThat(services.size(), equalTo(8));
 
         var providers = providers(services);
 
@@ -176,7 +177,8 @@ public void testGetServicesWithChatCompletionTaskType() throws IOException {
                     "streaming_completion_test_service",
                     "hugging_face",
                     "amazon_sagemaker",
-                    "googlevertexai"
+                    "googlevertexai",
+                    "mistral"
                 ).toArray()
             )
         );

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

@@ -100,6 +100,7 @@
 import org.elasticsearch.xpack.inference.services.jinaai.embeddings.JinaAIEmbeddingsTaskSettings;
 import org.elasticsearch.xpack.inference.services.jinaai.rerank.JinaAIRerankServiceSettings;
 import org.elasticsearch.xpack.inference.services.jinaai.rerank.JinaAIRerankTaskSettings;
+import org.elasticsearch.xpack.inference.services.mistral.completion.MistralChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.mistral.embeddings.MistralEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionTaskSettings;
@@ -266,6 +267,13 @@ private static void addMistralNamedWriteables(List<NamedWriteableRegistry.Entry>
                 MistralEmbeddingsServiceSettings::new
             )
         );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                MistralChatCompletionServiceSettings.NAME,
+                MistralChatCompletionServiceSettings::new
+            )
+        );
 
         // note - no task settings for Mistral embeddings...
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/ErrorMessageResponseEntity.java

@@ -21,12 +21,13 @@
  * A pattern is emerging in how external providers provide error responses.
  *
  * At a minimum, these return:
+ * <pre><code>
  * {
  *     "error: {
  *         "message": "(error message)"
  *     }
  * }
- *
+ * </code></pre>
  * Others may return additional information such as error codes specific to the service.
  *
  * This currently covers error handling for Azure AI Studio, however this pattern