add default inference endpoint for Elastic Inference Service rerank (#129681)

brendan-jugan-elastic · elasticsearchmachine · web-flow · commit cef717c08706 · 2025-06-23T09:48:50.000-04:00
* add Elastic Inference Service rerank default inference endpoint

* [CI] Auto commit changes from spotless

* fix integ tests

* update mock Elastic Inference Service authorization response

* fix rerank service test

---------

Co-authored-by: elasticsearchmachine &lt;infra-root+elasticsearchmachine@elastic.co&gt;
diff --git a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetModelsWithElasticInferenceServiceIT.java b/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetModelsWithElasticInferenceServiceIT.java
@@ -33,7 +33,7 @@ public void testGetDefaultEndpoints() throws IOException {
         var allModels = getAllModels();
         var chatCompletionModels = getModels("_all", TaskType.CHAT_COMPLETION);
 
-        assertThat(allModels, hasSize(5));
+        assertThat(allModels, hasSize(6));
         assertThat(chatCompletionModels, hasSize(1));
 
         for (var model : chatCompletionModels) {
@@ -42,6 +42,7 @@ public void testGetDefaultEndpoints() throws IOException {
 
         assertInferenceIdTaskType(allModels, ".rainbow-sprinkles-elastic", TaskType.CHAT_COMPLETION);
         assertInferenceIdTaskType(allModels, ".elser-v2-elastic", TaskType.SPARSE_EMBEDDING);
+        assertInferenceIdTaskType(allModels, ".rerank-v1-elastic", TaskType.RERANK);
     }
 
     private static void assertInferenceIdTaskType(List<Map<String, Object>> models, String inferenceId, TaskType taskType) {
diff --git a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java b/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java
@@ -111,7 +111,7 @@ public void testGetServicesWithTextEmbeddingTaskType() throws IOException {
 
     public void testGetServicesWithRerankTaskType() throws IOException {
         List<Object> services = getServices(TaskType.RERANK);
-        assertThat(services.size(), equalTo(9));
+        assertThat(services.size(), equalTo(10));
 
         var providers = providers(services);
 
@@ -127,7 +127,8 @@ public void testGetServicesWithRerankTaskType() throws IOException {
                     "jinaai",
                     "test_reranking_service",
                     "voyageai",
-                    "hugging_face"
+                    "hugging_face",
+                    "elastic"
                 ).toArray()
             )
         );
diff --git a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/MockElasticInferenceServiceAuthorizationServer.java b/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/MockElasticInferenceServiceAuthorizationServer.java
@@ -41,6 +41,10 @@ public void enqueueAuthorizeAllModelsResponse() {
                     {
                       "model_name": "elser-v2",
                       "task_types": ["embed/text/sparse"]
+                    },
+                    {
+                      "model_name": "rerank-v1",
+                      "task_types": ["rerank/text/text-similarity"]
                     }
                 ]
             }
diff --git a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/InferenceRevokeDefaultEndpointsIT.java b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/InferenceRevokeDefaultEndpointsIT.java
@@ -197,6 +197,10 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
                         {
                           "model_name": "elser-v2",
                           "task_types": ["embed/text/sparse"]
+                        },
+                        {
+                          "model_name": "rerank-v1",
+                          "task_types": ["rerank/text/text-similarity"]
                         }
                     ]
                 }
@@ -221,16 +225,25 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
                                 ".rainbow-sprinkles-elastic",
                                 MinimalServiceSettings.chatCompletion(ElasticInferenceService.NAME),
                                 service
+                            ),
+                            new InferenceService.DefaultConfigId(
+                                ".rerank-v1-elastic",
+                                MinimalServiceSettings.rerank(ElasticInferenceService.NAME),
+                                service
                             )
                         )
                     )
                 );
-                assertThat(service.supportedTaskTypes(), is(EnumSet.of(TaskType.CHAT_COMPLETION, TaskType.SPARSE_EMBEDDING)));
+                assertThat(
+                    service.supportedTaskTypes(),
+                    is(EnumSet.of(TaskType.CHAT_COMPLETION, TaskType.SPARSE_EMBEDDING, TaskType.RERANK))
+                );
 
                 PlainActionFuture<List<Model>> listener = new PlainActionFuture<>();
                 service.defaultConfigs(listener);
                 assertThat(listener.actionGet(TIMEOUT).get(0).getConfigurations().getInferenceEntityId(), is(".elser-v2-elastic"));
                 assertThat(listener.actionGet(TIMEOUT).get(1).getConfigurations().getInferenceEntityId(), is(".rainbow-sprinkles-elastic"));
+                assertThat(listener.actionGet(TIMEOUT).get(2).getConfigurations().getInferenceEntityId(), is(".rerank-v1-elastic"));
 
                 var getModelListener = new PlainActionFuture<UnparsedModel>();
                 // persists the default endpoints
@@ -248,6 +261,10 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
                         {
                           "model_name": "elser-v2",
                           "task_types": ["embed/text/sparse"]
+                        },
+                        {
+                          "model_name": "rerank-v1",
+                          "task_types": ["rerank/text/text-similarity"]
                         }
                     ]
                 }
@@ -267,11 +284,16 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
                                 ".elser-v2-elastic",
                                 MinimalServiceSettings.sparseEmbedding(ElasticInferenceService.NAME),
                                 service
+                            ),
+                            new InferenceService.DefaultConfigId(
+                                ".rerank-v1-elastic",
+                                MinimalServiceSettings.rerank(ElasticInferenceService.NAME),
+                                service
                             )
                         )
                     )
                 );
-                assertThat(service.supportedTaskTypes(), is(EnumSet.of(TaskType.SPARSE_EMBEDDING)));
+                assertThat(service.supportedTaskTypes(), is(EnumSet.of(TaskType.SPARSE_EMBEDDING, TaskType.RERANK)));
 
                 var getModelListener = new PlainActionFuture<UnparsedModel>();
                 modelRegistry.getModel(".rainbow-sprinkles-elastic", getModelListener);
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java
@@ -52,6 +52,7 @@
 import org.elasticsearch.xpack.inference.services.elastic.completion.ElasticInferenceServiceCompletionModel;
 import org.elasticsearch.xpack.inference.services.elastic.completion.ElasticInferenceServiceCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.elastic.rerank.ElasticInferenceServiceRerankModel;
+import org.elasticsearch.xpack.inference.services.elastic.rerank.ElasticInferenceServiceRerankServiceSettings;
 import org.elasticsearch.xpack.inference.services.elastic.sparseembeddings.ElasticInferenceServiceSparseEmbeddingsModel;
 import org.elasticsearch.xpack.inference.services.elastic.sparseembeddings.ElasticInferenceServiceSparseEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
@@ -95,6 +96,10 @@ public class ElasticInferenceService extends SenderService {
     static final String DEFAULT_ELSER_MODEL_ID_V2 = "elser-v2";
     static final String DEFAULT_ELSER_ENDPOINT_ID_V2 = defaultEndpointId(DEFAULT_ELSER_MODEL_ID_V2);
 
+    // rerank-v1
+    static final String DEFAULT_RERANK_MODEL_ID_V1 = "rerank-v1";
+    static final String DEFAULT_RERANK_ENDPOINT_ID_V1 = defaultEndpointId(DEFAULT_RERANK_MODEL_ID_V1);
+
     /**
      * The task types that the {@link InferenceAction.Request} can accept.
      */
@@ -159,6 +164,19 @@ private static Map<String, DefaultModelConfig> initDefaultEndpoints(
                     elasticInferenceServiceComponents
                 ),
                 MinimalServiceSettings.sparseEmbedding(NAME)
+            ),
+            DEFAULT_RERANK_MODEL_ID_V1,
+            new DefaultModelConfig(
+                new ElasticInferenceServiceRerankModel(
+                    DEFAULT_RERANK_ENDPOINT_ID_V1,
+                    TaskType.RERANK,
+                    NAME,
+                    new ElasticInferenceServiceRerankServiceSettings(DEFAULT_RERANK_MODEL_ID_V1, null),
+                    EmptyTaskSettings.INSTANCE,
+                    EmptySecretSettings.INSTANCE,
+                    elasticInferenceServiceComponents
+                ),
+                MinimalServiceSettings.rerank(NAME)
             )
         );
     }
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/rerank/ElasticInferenceServiceRerankModel.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/rerank/ElasticInferenceServiceRerankModel.java
@@ -87,7 +87,7 @@ public URI uri() {
     private URI createUri() throws ElasticsearchStatusException {
         try {
             // TODO, consider transforming the base URL into a URI for better error handling.
-            return new URI(elasticInferenceServiceComponents().elasticInferenceServiceUrl() + "/api/v1/rerank");
+            return new URI(elasticInferenceServiceComponents().elasticInferenceServiceUrl() + "/api/v1/rerank/text/text-similarity");
         } catch (URISyntaxException e) {
             throw new ElasticsearchStatusException(
                 "Failed to create URI for service ["
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/response/ElasticInferenceServiceAuthorizationResponseEntity.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/response/ElasticInferenceServiceAuthorizationResponseEntity.java
@@ -43,7 +43,9 @@ public class ElasticInferenceServiceAuthorizationResponseEntity implements Infer
         "embed/text/sparse",
         TaskType.SPARSE_EMBEDDING,
         "chat",
-        TaskType.CHAT_COMPLETION
+        TaskType.CHAT_COMPLETION,
+        "rerank/text/text-similarity",
+        TaskType.RERANK
     );
 
     @SuppressWarnings("unchecked")
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceTests.java
@@ -1294,6 +1294,10 @@ public void testDefaultConfigs_Returns_DefaultEndpoints_WhenTaskTypeIsCorrect()
                     {
                       "model_name": "elser-v2",
                       "task_types": ["embed/text/sparse"]
+                    },
+                    {
+                      "model_name": "rerank-v1",
+                      "task_types": ["rerank/text/text-similarity"]
                     }
                 ]
             }
@@ -1319,18 +1323,25 @@ public void testDefaultConfigs_Returns_DefaultEndpoints_WhenTaskTypeIsCorrect()
                             ".rainbow-sprinkles-elastic",
                             MinimalServiceSettings.chatCompletion(ElasticInferenceService.NAME),
                             service
+                        ),
+                        new InferenceService.DefaultConfigId(
+                            ".rerank-v1-elastic",
+                            MinimalServiceSettings.rerank(ElasticInferenceService.NAME),
+                            service
                         )
                     )
                 )
             );
-            assertThat(service.supportedTaskTypes(), is(EnumSet.of(TaskType.CHAT_COMPLETION, TaskType.SPARSE_EMBEDDING)));
+            assertThat(service.supportedTaskTypes(), is(EnumSet.of(TaskType.CHAT_COMPLETION, TaskType.SPARSE_EMBEDDING, TaskType.RERANK)));
 
             PlainActionFuture<List<Model>> listener = new PlainActionFuture<>();
             service.defaultConfigs(listener);
             var models = listener.actionGet(TIMEOUT);
-            assertThat(models.size(), is(2));
+            assertThat(models.size(), is(3));
             assertThat(models.get(0).getConfigurations().getInferenceEntityId(), is(".elser-v2-elastic"));
             assertThat(models.get(1).getConfigurations().getInferenceEntityId(), is(".rainbow-sprinkles-elastic"));
+            assertThat(models.get(2).getConfigurations().getInferenceEntityId(), is(".rerank-v1-elastic"));
+
         }
     }
 

Original file line number	Diff line number	Diff line change
`@@ -41,6 +41,10 @@ public void enqueueAuthorizeAllModelsResponse() {`
`41`	`41`	`{`
`42`	`42`	`"model_name": "elser-v2",`
`43`	`43`	`"task_types": ["embed/text/sparse"]`
	`44`	`+ },`
	`45`	`+ {`
	`46`	`+ "model_name": "rerank-v1",`
	`47`	`+ "task_types": ["rerank/text/text-similarity"]`
`44`	`48`	`}`
`45`	`49`	`]`
`46`	`50`	`}`
Original file line number	Diff line number	Diff line change
`@@ -197,6 +197,10 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA`
`197`	`197`	`{`
`198`	`198`	`"model_name": "elser-v2",`
`199`	`199`	`"task_types": ["embed/text/sparse"]`
	`200`	`+ },`
	`201`	`+ {`
	`202`	`+ "model_name": "rerank-v1",`
	`203`	`+ "task_types": ["rerank/text/text-similarity"]`
`200`	`204`	`}`
`201`	`205`	`]`
`202`	`206`	`}`
`@@ -221,16 +225,25 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA`
`221`	`225`	`".rainbow-sprinkles-elastic",`
`222`	`226`	`MinimalServiceSettings.chatCompletion(ElasticInferenceService.NAME),`
`223`	`227`	`service`
	`228`	`+ ),`
	`229`	`+ new InferenceService.DefaultConfigId(`
	`230`	`+ ".rerank-v1-elastic",`
	`231`	`+ MinimalServiceSettings.rerank(ElasticInferenceService.NAME),`
	`232`	`+ service`
`224`	`233`	`)`
`225`	`234`	`)`
`226`	`235`	`)`
`227`	`236`	`);`
`228`		`- assertThat(service.supportedTaskTypes(), is(EnumSet.of(TaskType.CHAT_COMPLETION, TaskType.SPARSE_EMBEDDING)));`
	`237`	`+ assertThat(`
	`238`	`+ service.supportedTaskTypes(),`
	`239`	`+ is(EnumSet.of(TaskType.CHAT_COMPLETION, TaskType.SPARSE_EMBEDDING, TaskType.RERANK))`
	`240`	`+ );`
`229`	`241`
`230`	`242`	`PlainActionFuture<List<Model>> listener = new PlainActionFuture<>();`
`231`	`243`	`service.defaultConfigs(listener);`
`232`	`244`	`assertThat(listener.actionGet(TIMEOUT).get(0).getConfigurations().getInferenceEntityId(), is(".elser-v2-elastic"));`
`233`	`245`	`assertThat(listener.actionGet(TIMEOUT).get(1).getConfigurations().getInferenceEntityId(), is(".rainbow-sprinkles-elastic"));`
	`246`	`+ assertThat(listener.actionGet(TIMEOUT).get(2).getConfigurations().getInferenceEntityId(), is(".rerank-v1-elastic"));`
`234`	`247`
`235`	`248`	`var getModelListener = new PlainActionFuture<UnparsedModel>();`
`236`	`249`	`// persists the default endpoints`
`@@ -248,6 +261,10 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA`
`248`	`261`	`{`
`249`	`262`	`"model_name": "elser-v2",`
`250`	`263`	`"task_types": ["embed/text/sparse"]`
	`264`	`+ },`
	`265`	`+ {`
	`266`	`+ "model_name": "rerank-v1",`
	`267`	`+ "task_types": ["rerank/text/text-similarity"]`
`251`	`268`	`}`
`252`	`269`	`]`
`253`	`270`	`}`
`@@ -267,11 +284,16 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA`
`267`	`284`	`".elser-v2-elastic",`
`268`	`285`	`MinimalServiceSettings.sparseEmbedding(ElasticInferenceService.NAME),`
`269`	`286`	`service`
	`287`	`+ ),`
	`288`	`+ new InferenceService.DefaultConfigId(`
	`289`	`+ ".rerank-v1-elastic",`
	`290`	`+ MinimalServiceSettings.rerank(ElasticInferenceService.NAME),`
	`291`	`+ service`
`270`	`292`	`)`
`271`	`293`	`)`
`272`	`294`	`)`
`273`	`295`	`);`
`274`		`- assertThat(service.supportedTaskTypes(), is(EnumSet.of(TaskType.SPARSE_EMBEDDING)));`
	`296`	`+ assertThat(service.supportedTaskTypes(), is(EnumSet.of(TaskType.SPARSE_EMBEDDING, TaskType.RERANK)));`
`275`	`297`
`276`	`298`	`var getModelListener = new PlainActionFuture<UnparsedModel>();`
`277`	`299`	`modelRegistry.getModel(".rainbow-sprinkles-elastic", getModelListener);`
Original file line number	Diff line number	Diff line change
`@@ -1294,6 +1294,10 @@ public void testDefaultConfigs_Returns_DefaultEndpoints_WhenTaskTypeIsCorrect()`
`1294`	`1294`	`{`
`1295`	`1295`	`"model_name": "elser-v2",`
`1296`	`1296`	`"task_types": ["embed/text/sparse"]`
	`1297`	`+ },`
	`1298`	`+ {`
	`1299`	`+ "model_name": "rerank-v1",`
	`1300`	`+ "task_types": ["rerank/text/text-similarity"]`
`1297`	`1301`	`}`
`1298`	`1302`	`]`
`1299`	`1303`	`}`
`@@ -1319,18 +1323,25 @@ public void testDefaultConfigs_Returns_DefaultEndpoints_WhenTaskTypeIsCorrect()`
`1319`	`1323`	`".rainbow-sprinkles-elastic",`
`1320`	`1324`	`MinimalServiceSettings.chatCompletion(ElasticInferenceService.NAME),`
`1321`	`1325`	`service`
	`1326`	`+ ),`
	`1327`	`+ new InferenceService.DefaultConfigId(`
	`1328`	`+ ".rerank-v1-elastic",`
	`1329`	`+ MinimalServiceSettings.rerank(ElasticInferenceService.NAME),`
	`1330`	`+ service`
`1322`	`1331`	`)`
`1323`	`1332`	`)`
`1324`	`1333`	`)`
`1325`	`1334`	`);`
`1326`		`- assertThat(service.supportedTaskTypes(), is(EnumSet.of(TaskType.CHAT_COMPLETION, TaskType.SPARSE_EMBEDDING)));`
	`1335`	`+ assertThat(service.supportedTaskTypes(), is(EnumSet.of(TaskType.CHAT_COMPLETION, TaskType.SPARSE_EMBEDDING, TaskType.RERANK)));`
`1327`	`1336`
`1328`	`1337`	`PlainActionFuture<List<Model>> listener = new PlainActionFuture<>();`
`1329`	`1338`	`service.defaultConfigs(listener);`
`1330`	`1339`	`var models = listener.actionGet(TIMEOUT);`
`1331`		`- assertThat(models.size(), is(2));`
	`1340`	`+ assertThat(models.size(), is(3));`
`1332`	`1341`	`assertThat(models.get(0).getConfigurations().getInferenceEntityId(), is(".elser-v2-elastic"));`
`1333`	`1342`	`assertThat(models.get(1).getConfigurations().getInferenceEntityId(), is(".rainbow-sprinkles-elastic"));`
	`1343`	`+ assertThat(models.get(2).getConfigurations().getInferenceEntityId(), is(".rerank-v1-elastic"));`
	`1344`	`+`
`1334`	`1345`	`}`
`1335`	`1346`	`}`
`1336`	`1347`