[ML] Allow a single model name to map to multiple inference endpoints (#138059)

* Switching to list so we can support completion and chat completion

* Adding test for empty list

Author: jonathan-buttner

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/InternalPreconfiguredEndpoints.java

@@ -16,8 +16,10 @@
 import org.elasticsearch.xpack.inference.services.elastic.rerank.ElasticInferenceServiceRerankServiceSettings;
 import org.elasticsearch.xpack.inference.services.elastic.sparseembeddings.ElasticInferenceServiceSparseEmbeddingsServiceSettings;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 
 import static java.util.stream.Collectors.toMap;
 
@@ -63,65 +65,80 @@ public record MinimalModel(
     private static final ElasticInferenceServiceRerankServiceSettings RERANK_SERVICE_SETTINGS =
         new ElasticInferenceServiceRerankServiceSettings(DEFAULT_RERANK_MODEL_ID_V1);
 
-    private static final Map<String, MinimalModel> MODEL_NAME_TO_MINIMAL_MODEL = Map.of(
+    // A single model name can map to multiple inference endpoints, so we need a String to a List
+    private static final Map<String, List<MinimalModel>> MODEL_NAME_TO_MINIMAL_MODELS = Map.of(
         DEFAULT_CHAT_COMPLETION_MODEL_ID_V1,
-        new MinimalModel(
-            new ModelConfigurations(
-                DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1,
-                TaskType.CHAT_COMPLETION,
-                ElasticInferenceService.NAME,
-                COMPLETION_SERVICE_SETTINGS,
-                ChunkingSettingsBuilder.DEFAULT_SETTINGS
-            ),
-            COMPLETION_SERVICE_SETTINGS
+        List.of(
+            new MinimalModel(
+                new ModelConfigurations(
+                    DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1,
+                    TaskType.CHAT_COMPLETION,
+                    ElasticInferenceService.NAME,
+                    COMPLETION_SERVICE_SETTINGS,
+                    ChunkingSettingsBuilder.DEFAULT_SETTINGS
+                ),
+                COMPLETION_SERVICE_SETTINGS
+            )
         ),
         DEFAULT_ELSER_2_MODEL_ID,
-        new MinimalModel(
-            new ModelConfigurations(
-                DEFAULT_ELSER_ENDPOINT_ID_V2,
-                TaskType.SPARSE_EMBEDDING,
-                ElasticInferenceService.NAME,
-                SPARSE_EMBEDDINGS_SERVICE_SETTINGS,
-                ChunkingSettingsBuilder.DEFAULT_SETTINGS
-            ),
-            SPARSE_EMBEDDINGS_SERVICE_SETTINGS
+        List.of(
+            new MinimalModel(
+                new ModelConfigurations(
+                    DEFAULT_ELSER_ENDPOINT_ID_V2,
+                    TaskType.SPARSE_EMBEDDING,
+                    ElasticInferenceService.NAME,
+                    SPARSE_EMBEDDINGS_SERVICE_SETTINGS,
+                    ChunkingSettingsBuilder.DEFAULT_SETTINGS
+                ),
+                SPARSE_EMBEDDINGS_SERVICE_SETTINGS
+            )
         ),
         DEFAULT_MULTILINGUAL_EMBED_MODEL_ID,
-        new MinimalModel(
-            new ModelConfigurations(
-                DEFAULT_MULTILINGUAL_EMBED_ENDPOINT_ID,
-                TaskType.TEXT_EMBEDDING,
-                ElasticInferenceService.NAME,
-                DENSE_TEXT_EMBEDDINGS_SERVICE_SETTINGS,
-                ChunkingSettingsBuilder.DEFAULT_SETTINGS
-            ),
-            DENSE_TEXT_EMBEDDINGS_SERVICE_SETTINGS
+        List.of(
+            new MinimalModel(
+                new ModelConfigurations(
+                    DEFAULT_MULTILINGUAL_EMBED_ENDPOINT_ID,
+                    TaskType.TEXT_EMBEDDING,
+                    ElasticInferenceService.NAME,
+                    DENSE_TEXT_EMBEDDINGS_SERVICE_SETTINGS,
+                    ChunkingSettingsBuilder.DEFAULT_SETTINGS
+                ),
+                DENSE_TEXT_EMBEDDINGS_SERVICE_SETTINGS
+            )
         ),
         DEFAULT_RERANK_MODEL_ID_V1,
-        new MinimalModel(
-            new ModelConfigurations(
-                DEFAULT_RERANK_ENDPOINT_ID_V1,
-                TaskType.RERANK,
-                ElasticInferenceService.NAME,
-                RERANK_SERVICE_SETTINGS,
-                ChunkingSettingsBuilder.DEFAULT_SETTINGS
-            ),
-            RERANK_SERVICE_SETTINGS
+        List.of(
+            new MinimalModel(
+                new ModelConfigurations(
+                    DEFAULT_RERANK_ENDPOINT_ID_V1,
+                    TaskType.RERANK,
+                    ElasticInferenceService.NAME,
+                    RERANK_SERVICE_SETTINGS,
+                    ChunkingSettingsBuilder.DEFAULT_SETTINGS
+                ),
+                RERANK_SERVICE_SETTINGS
+            )
         )
     );
 
-    private static final Map<String, MinimalModel> INFERENCE_ID_TO_MINIMAL_MODEL = MODEL_NAME_TO_MINIMAL_MODEL.entrySet()
+    private static final Map<String, MinimalModel> INFERENCE_ID_TO_MINIMAL_MODEL = MODEL_NAME_TO_MINIMAL_MODELS.entrySet()
         .stream()
-        .collect(toMap(e -> e.getValue().configurations().getInferenceEntityId(), Map.Entry::getValue));
+        .flatMap(entry -> entry.getValue().stream())
+        .collect(toMap(m -> m.configurations().getInferenceEntityId(), Function.identity()));
 
     public static final Set<String> EIS_PRECONFIGURED_ENDPOINT_IDS = Set.copyOf(INFERENCE_ID_TO_MINIMAL_MODEL.keySet());
 
     public static SimilarityMeasure defaultDenseTextEmbeddingsSimilarity() {
         return SimilarityMeasure.COSINE;
     }
 
-    public static MinimalModel getWithModelName(String modelName) {
-        return MODEL_NAME_TO_MINIMAL_MODEL.get(modelName);
+    public static List<MinimalModel> getWithModelName(String modelName) {
+        var minimalModels = MODEL_NAME_TO_MINIMAL_MODELS.get(modelName);
+        if (minimalModels == null) {
+            return List.of();
+        }
+
+        return minimalModels;
     }
 
     public static MinimalModel getWithInferenceId(String inferenceId) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/AuthorizationPoller.java

@@ -31,6 +31,7 @@
 import org.elasticsearch.xpack.inference.services.elastic.InternalPreconfiguredEndpoints;
 
 import java.util.EnumSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -248,7 +249,7 @@ private Set<String> getNewInferenceEndpointsToStore(ElasticInferenceServiceAutho
 
         var newInferenceIds = authorizedModelIds.stream()
             .map(InternalPreconfiguredEndpoints::getWithModelName)
-            .filter(Objects::nonNull)
+            .flatMap(List::stream)
             .map(model -> model.configurations().getInferenceEntityId())
             .collect(Collectors.toSet());
 

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/InternalPreconfiguredEndpointsTests.java

@@ -0,0 +1,23 @@
+/*
+ * 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.services.elastic;
+
+import org.elasticsearch.test.ESTestCase;
+
+import static org.hamcrest.Matchers.hasSize;
+
+public class InternalPreconfiguredEndpointsTests extends ESTestCase {
+    public void testGetWithModelName_ReturnsAnEmptyList_IfNameDoesNotExist() {
+        assertThat(InternalPreconfiguredEndpoints.getWithModelName("non-existent-model"), hasSize(0));
+    }
+
+    public void testGetWithModelName_ReturnsChatCompletionModels() {
+        var models = InternalPreconfiguredEndpoints.getWithModelName(InternalPreconfiguredEndpoints.DEFAULT_CHAT_COMPLETION_MODEL_ID_V1);
+        assertThat(models, hasSize(1));
+    }
+}