Adding endpoint creation validation for all task types to remaining services (#115020)

* Adding endpoint creation validation for all task types to remaining services

* Update Cohere IT tests for rerank validation

* Adding missing import

* Update docs/changelog/115020.yaml

* Fixing GoogleVertex tests after merge from upstream

---------

Co-authored-by: Elastic Machine <[email protected]>

Author: dan-rubinstein

docs/changelog/115020.yaml

@@ -0,0 +1,5 @@
+pr: 115020
+summary: Adding endpoint creation validation for all task types to remaining services
+area: Machine Learning
+type: enhancement
+issues: []

x-pack/plugin/inference/qa/mixed-cluster/src/javaRestTest/java/org/elasticsearch/xpack/inference/qa/mixed/CohereServiceMixedIT.java

@@ -135,6 +135,7 @@ public void testRerank() throws IOException {
 
         final String inferenceId = "mixed-cluster-rerank";
 
+        cohereRerankServer.enqueue(new MockResponse().setResponseCode(200).setBody(rerankResponse()));
         put(inferenceId, rerankConfig(getUrl(cohereRerankServer)), TaskType.RERANK);
         assertRerank(inferenceId);
 

x-pack/plugin/inference/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/xpack/application/CohereServiceUpgradeIT.java

@@ -201,6 +201,7 @@ public void testRerank() throws IOException {
         var testTaskType = TaskType.RERANK;
 
         if (isOldCluster()) {
+            cohereRerankServer.enqueue(new MockResponse().setResponseCode(200).setBody(rerankResponse()));
             put(oldClusterId, rerankConfig(getUrl(cohereRerankServer)), testTaskType);
             var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get(old_cluster_endpoint_identifier);
             assertThat(configs, hasSize(1));
@@ -229,6 +230,7 @@ public void testRerank() throws IOException {
             assertRerank(oldClusterId);
 
             // New endpoint
+            cohereRerankServer.enqueue(new MockResponse().setResponseCode(200).setBody(rerankResponse()));
             put(upgradedClusterId, rerankConfig(getUrl(cohereRerankServer)), testTaskType);
             configs = (List<Map<String, Object>>) get(upgradedClusterId).get("endpoints");
             assertThat(configs, hasSize(1));

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/alibabacloudsearch/AlibabaCloudSearchService.java

@@ -18,7 +18,6 @@
 import org.elasticsearch.inference.ChunkingOptions;
 import org.elasticsearch.inference.ChunkingSettings;
 import org.elasticsearch.inference.EmptySettingsConfiguration;
-import org.elasticsearch.inference.InferenceService;
 import org.elasticsearch.inference.InferenceServiceConfiguration;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.inference.InputType;
@@ -51,6 +50,7 @@
 import org.elasticsearch.xpack.inference.services.alibabacloudsearch.sparse.AlibabaCloudSearchSparseModel;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
+import org.elasticsearch.xpack.inference.services.validation.ModelValidatorBuilder;
 
 import java.util.EnumSet;
 import java.util.HashMap;
@@ -60,7 +60,6 @@
 
 import static org.elasticsearch.inference.TaskType.SPARSE_EMBEDDING;
 import static org.elasticsearch.inference.TaskType.TEXT_EMBEDDING;
-import static org.elasticsearch.xpack.core.inference.action.InferenceAction.Request.DEFAULT_TIMEOUT;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.parsePersistedConfigErrorMsg;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMap;
@@ -332,68 +331,39 @@ private EmbeddingRequestChunker.EmbeddingType getEmbeddingTypeFromTaskType(TaskT
      */
     @Override
     public void checkModelConfig(Model model, ActionListener<Model> listener) {
+        // TODO: Remove this function once all services have been updated to use the new model validators
+        ModelValidatorBuilder.buildModelValidator(model.getTaskType()).validate(this, model, listener);
+    }
+
+    @Override
+    public Model updateModelWithEmbeddingDetails(Model model, int embeddingSize) {
         if (model instanceof AlibabaCloudSearchEmbeddingsModel embeddingsModel) {
-            ServiceUtils.getEmbeddingSize(
-                model,
-                this,
-                listener.delegateFailureAndWrap((l, size) -> l.onResponse(updateModelWithEmbeddingDetails(embeddingsModel, size)))
+            var serviceSettings = embeddingsModel.getServiceSettings();
+
+            var updatedServiceSettings = new AlibabaCloudSearchEmbeddingsServiceSettings(
+                new AlibabaCloudSearchServiceSettings(
+                    serviceSettings.getCommonSettings().modelId(),
+                    serviceSettings.getCommonSettings().getHost(),
+                    serviceSettings.getCommonSettings().getWorkspaceName(),
+                    serviceSettings.getCommonSettings().getHttpSchema(),
+                    serviceSettings.getCommonSettings().rateLimitSettings()
+                ),
+                SimilarityMeasure.DOT_PRODUCT,
+                embeddingSize,
+                serviceSettings.getMaxInputTokens()
             );
+
+            return new AlibabaCloudSearchEmbeddingsModel(embeddingsModel, updatedServiceSettings);
         } else {
-            checkAlibabaCloudSearchServiceConfig(model, this, listener);
+            throw ServiceUtils.invalidModelTypeForUpdateModelWithEmbeddingDetails(model.getClass());
         }
     }
 
-    private AlibabaCloudSearchEmbeddingsModel updateModelWithEmbeddingDetails(AlibabaCloudSearchEmbeddingsModel model, int embeddingSize) {
-        AlibabaCloudSearchEmbeddingsServiceSettings serviceSettings = new AlibabaCloudSearchEmbeddingsServiceSettings(
-            new AlibabaCloudSearchServiceSettings(
-                model.getServiceSettings().getCommonSettings().modelId(),
-                model.getServiceSettings().getCommonSettings().getHost(),
-                model.getServiceSettings().getCommonSettings().getWorkspaceName(),
-                model.getServiceSettings().getCommonSettings().getHttpSchema(),
-                model.getServiceSettings().getCommonSettings().rateLimitSettings()
-            ),
-            SimilarityMeasure.DOT_PRODUCT,
-            embeddingSize,
-            model.getServiceSettings().getMaxInputTokens()
-        );
-
-        return new AlibabaCloudSearchEmbeddingsModel(model, serviceSettings);
-    }
-
     @Override
     public TransportVersion getMinimalSupportedVersion() {
         return TransportVersions.ML_INFERENCE_ALIBABACLOUD_SEARCH_ADDED;
     }
 
-    /**
-     * For other models except of text embedding
-     * check the model's service settings and task settings
-     *
-     * @param model The new model
-     * @param service The inferenceService
-     * @param listener The listener
-     */
-    private void checkAlibabaCloudSearchServiceConfig(Model model, InferenceService service, ActionListener<Model> listener) {
-        String input = ALIBABA_CLOUD_SEARCH_SERVICE_CONFIG_INPUT;
-        String query = model.getTaskType().equals(TaskType.RERANK) ? ALIBABA_CLOUD_SEARCH_SERVICE_CONFIG_QUERY : null;
-
-        service.infer(
-            model,
-            query,
-            List.of(input),
-            false,
-            Map.of(),
-            InputType.INGEST,
-            DEFAULT_TIMEOUT,
-            listener.delegateFailureAndWrap((delegate, r) -> {
-                listener.onResponse(model);
-            })
-        );
-    }
-
-    private static final String ALIBABA_CLOUD_SEARCH_SERVICE_CONFIG_INPUT = "input";
-    private static final String ALIBABA_CLOUD_SEARCH_SERVICE_CONFIG_QUERY = "query";
-
     public static class Configuration {
         public static InferenceServiceConfiguration get() {
             return configuration.getOrCompute();

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockService.java

@@ -49,6 +49,7 @@
 import org.elasticsearch.xpack.inference.services.amazonbedrock.embeddings.AmazonBedrockEmbeddingsModel;
 import org.elasticsearch.xpack.inference.services.amazonbedrock.embeddings.AmazonBedrockEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
+import org.elasticsearch.xpack.inference.services.validation.ModelValidatorBuilder;
 
 import java.io.IOException;
 import java.util.EnumSet;
@@ -303,49 +304,34 @@ public Set<TaskType> supportedStreamingTasks() {
      */
     @Override
     public void checkModelConfig(Model model, ActionListener<Model> listener) {
-        if (model instanceof AmazonBedrockEmbeddingsModel embeddingsModel) {
-            ServiceUtils.getEmbeddingSize(
-                model,
-                this,
-                listener.delegateFailureAndWrap((l, size) -> l.onResponse(updateModelWithEmbeddingDetails(embeddingsModel, size)))
-            );
-        } else {
-            listener.onResponse(model);
-        }
+        // TODO: Remove this function once all services have been updated to use the new model validators
+        ModelValidatorBuilder.buildModelValidator(model.getTaskType()).validate(this, model, listener);
     }
 
-    private AmazonBedrockEmbeddingsModel updateModelWithEmbeddingDetails(AmazonBedrockEmbeddingsModel model, int embeddingSize) {
-        AmazonBedrockEmbeddingsServiceSettings serviceSettings = model.getServiceSettings();
-        if (serviceSettings.dimensionsSetByUser()
-            && serviceSettings.dimensions() != null
-            && serviceSettings.dimensions() != embeddingSize) {
-            throw new ElasticsearchStatusException(
-                Strings.format(
-                    "The retrieved embeddings size [%s] does not match the size specified in the settings [%s]. "
-                        + "Please recreate the [%s] configuration with the correct dimensions",
-                    embeddingSize,
-                    serviceSettings.dimensions(),
-                    model.getConfigurations().getInferenceEntityId()
-                ),
-                RestStatus.BAD_REQUEST
+    @Override
+    public Model updateModelWithEmbeddingDetails(Model model, int embeddingSize) {
+        if (model instanceof AmazonBedrockEmbeddingsModel embeddingsModel) {
+            var serviceSettings = embeddingsModel.getServiceSettings();
+            var similarityFromModel = serviceSettings.similarity();
+            var similarityToUse = similarityFromModel == null
+                ? getProviderDefaultSimilarityMeasure(embeddingsModel.provider())
+                : similarityFromModel;
+
+            var updatedServiceSettings = new AmazonBedrockEmbeddingsServiceSettings(
+                serviceSettings.region(),
+                serviceSettings.modelId(),
+                serviceSettings.provider(),
+                embeddingSize,
+                serviceSettings.dimensionsSetByUser(),
+                serviceSettings.maxInputTokens(),
+                similarityToUse,
+                serviceSettings.rateLimitSettings()
             );
-        }
-
-        var similarityFromModel = serviceSettings.similarity();
-        var similarityToUse = similarityFromModel == null ? getProviderDefaultSimilarityMeasure(model.provider()) : similarityFromModel;
-
-        AmazonBedrockEmbeddingsServiceSettings settingsToUse = new AmazonBedrockEmbeddingsServiceSettings(
-            serviceSettings.region(),
-            serviceSettings.modelId(),
-            serviceSettings.provider(),
-            embeddingSize,
-            serviceSettings.dimensionsSetByUser(),
-            serviceSettings.maxInputTokens(),
-            similarityToUse,
-            serviceSettings.rateLimitSettings()
-        );
 
-        return new AmazonBedrockEmbeddingsModel(model, settingsToUse);
+            return new AmazonBedrockEmbeddingsModel(embeddingsModel, updatedServiceSettings);
+        } else {
+            throw ServiceUtils.invalidModelTypeForUpdateModelWithEmbeddingDetails(model.getClass());
+        }
     }
 
     private static void checkProviderForTask(TaskType taskType, AmazonBedrockProvider provider) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/anthropic/AnthropicService.java

@@ -39,6 +39,7 @@
 import org.elasticsearch.xpack.inference.services.anthropic.completion.AnthropicChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
+import org.elasticsearch.xpack.inference.services.validation.ModelValidatorBuilder;
 
 import java.util.EnumSet;
 import java.util.HashMap;
@@ -176,6 +177,12 @@ public AnthropicModel parsePersistedConfig(String inferenceEntityId, TaskType ta
         );
     }
 
+    @Override
+    public void checkModelConfig(Model model, ActionListener<Model> listener) {
+        // TODO: Remove this function once all services have been updated to use the new model validators
+        ModelValidatorBuilder.buildModelValidator(model.getTaskType()).validate(this, model, listener);
+    }
+
     @Override
     public InferenceServiceConfiguration getConfiguration() {
         return Configuration.get();

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiService.java

@@ -11,7 +11,6 @@
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.util.LazyInitializable;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
@@ -46,6 +45,7 @@
 import org.elasticsearch.xpack.inference.services.azureopenai.embeddings.AzureOpenAiEmbeddingsModel;
 import org.elasticsearch.xpack.inference.services.azureopenai.embeddings.AzureOpenAiEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
+import org.elasticsearch.xpack.inference.services.validation.ModelValidatorBuilder;
 
 import java.util.EnumSet;
 import java.util.HashMap;
@@ -294,48 +294,32 @@ protected void doChunkedInfer(
      */
     @Override
     public void checkModelConfig(Model model, ActionListener<Model> listener) {
-        if (model instanceof AzureOpenAiEmbeddingsModel embeddingsModel) {
-            ServiceUtils.getEmbeddingSize(
-                model,
-                this,
-                listener.delegateFailureAndWrap((l, size) -> l.onResponse(updateModelWithEmbeddingDetails(embeddingsModel, size)))
-            );
-        } else {
-            listener.onResponse(model);
-        }
+        // TODO: Remove this function once all services have been updated to use the new model validators
+        ModelValidatorBuilder.buildModelValidator(model.getTaskType()).validate(this, model, listener);
     }
 
-    private AzureOpenAiEmbeddingsModel updateModelWithEmbeddingDetails(AzureOpenAiEmbeddingsModel model, int embeddingSize) {
-        if (model.getServiceSettings().dimensionsSetByUser()
-            && model.getServiceSettings().dimensions() != null
-            && model.getServiceSettings().dimensions() != embeddingSize) {
-            throw new ElasticsearchStatusException(
-                Strings.format(
-                    "The retrieved embeddings size [%s] does not match the size specified in the settings [%s]. "
-                        + "Please recreate the [%s] configuration with the correct dimensions",
-                    embeddingSize,
-                    model.getServiceSettings().dimensions(),
-                    model.getConfigurations().getInferenceEntityId()
-                ),
-                RestStatus.BAD_REQUEST
+    @Override
+    public Model updateModelWithEmbeddingDetails(Model model, int embeddingSize) {
+        if (model instanceof AzureOpenAiEmbeddingsModel embeddingsModel) {
+            var serviceSettings = embeddingsModel.getServiceSettings();
+            var similarityFromModel = serviceSettings.similarity();
+            var similarityToUse = similarityFromModel == null ? SimilarityMeasure.DOT_PRODUCT : similarityFromModel;
+
+            var updatedServiceSettings = new AzureOpenAiEmbeddingsServiceSettings(
+                serviceSettings.resourceName(),
+                serviceSettings.deploymentId(),
+                serviceSettings.apiVersion(),
+                embeddingSize,
+                serviceSettings.dimensionsSetByUser(),
+                serviceSettings.maxInputTokens(),
+                similarityToUse,
+                serviceSettings.rateLimitSettings()
             );
-        }
-
-        var similarityFromModel = model.getServiceSettings().similarity();
-        var similarityToUse = similarityFromModel == null ? SimilarityMeasure.DOT_PRODUCT : similarityFromModel;
-
-        AzureOpenAiEmbeddingsServiceSettings serviceSettings = new AzureOpenAiEmbeddingsServiceSettings(
-            model.getServiceSettings().resourceName(),
-            model.getServiceSettings().deploymentId(),
-            model.getServiceSettings().apiVersion(),
-            embeddingSize,
-            model.getServiceSettings().dimensionsSetByUser(),
-            model.getServiceSettings().maxInputTokens(),
-            similarityToUse,
-            model.getServiceSettings().rateLimitSettings()
-        );
 
-        return new AzureOpenAiEmbeddingsModel(model, serviceSettings);
+            return new AzureOpenAiEmbeddingsModel(embeddingsModel, updatedServiceSettings);
+        } else {
+            throw ServiceUtils.invalidModelTypeForUpdateModelWithEmbeddingDetails(model.getClass());
+        }
     }
 
     @Override