Inference endpoint validation for OpenAIService (elastic#113137) (elastic#113546)

* Adding service integration and model validators for inference services.

* Adding ModelValidators to OpenAiService

* Cleaning up tests

* Cleaning up tests

* Adding mock response for completion model OpenAIService integration tests

---------

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

Author: dan-rubinstein

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

@@ -175,6 +175,17 @@ default void checkModelConfig(Model model, ActionListener<Model> listener) {
         listener.onResponse(model);
     };
 
+    /**
+     * Update a text embedding model's dimensions based on a provided embedding
+     * size and set the default similarity if required. The default behaviour is to just return the model.
+     * @param model The original model without updated embedding details
+     * @param embeddingSize The embedding size to update the model with
+     * @return The model with updated embedding details
+     */
+    default Model updateModelWithEmbeddingDetails(Model model, int embeddingSize) {
+        return model;
+    }
+
     /**
      * Return true if this model is hosted in the local Elasticsearch cluster
      * @return True if in cluster

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

@@ -34,6 +34,15 @@ default Integer dimensions() {
         return null;
     }
 
+    /**
+     * Boolean signifying whether the dimensions were set by the user
+     *
+     * @return boolean signifying whether the dimensions were set by the user
+     */
+    default Boolean dimensionsSetByUser() {
+        return null;
+    }
+
     /**
      * The data type for the embeddings this service works with. Defaults to null,
      * Text Embedding models should return a non-null value

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

@@ -95,6 +95,8 @@ public void testOpenAiCompletions() throws IOException {
         final String inferenceId = "mixed-cluster-completions";
         final String upgradedClusterId = "upgraded-cluster-completions";
 
+        // queue a response as PUT will call the service
+        openAiChatCompletionsServer.enqueue(new MockResponse().setResponseCode(200).setBody(chatCompletionsResponse()));
         put(inferenceId, chatCompletionsConfig(getUrl(openAiChatCompletionsServer)), TaskType.COMPLETION);
 
         var configsMap = get(TaskType.COMPLETION, inferenceId);

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

@@ -128,6 +128,7 @@ public void testOpenAiCompletions() throws IOException {
         var testTaskType = TaskType.COMPLETION;
 
         if (isOldCluster()) {
+            openAiChatCompletionsServer.enqueue(new MockResponse().setResponseCode(200).setBody(chatCompletionsResponse()));
             put(oldClusterId, chatCompletionsConfig(getUrl(openAiChatCompletionsServer)), testTaskType);
 
             var configs = (List<Map<String, Object>>) get(testTaskType, oldClusterId).get(old_cluster_endpoint_identifier);
@@ -157,6 +158,7 @@ public void testOpenAiCompletions() throws IOException {
 
             assertCompletionInference(oldClusterId);
 
+            openAiChatCompletionsServer.enqueue(new MockResponse().setResponseCode(200).setBody(chatCompletionsResponse()));
             put(upgradedClusterId, chatCompletionsConfig(getUrl(openAiChatCompletionsServer)), testTaskType);
             configs = (List<Map<String, Object>>) get(testTaskType, upgradedClusterId).get("endpoints");
             assertThat(configs, hasSize(1));

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

@@ -183,7 +183,8 @@ public Integer dimensions() {
         return dimensions;
     }
 
-    public boolean dimensionsSetByUser() {
+    @Override
+    public Boolean dimensionsSetByUser() {
         return this.dimensionsSetByUser;
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/azureaistudio/embeddings/AzureAiStudioEmbeddingsServiceSettings.java

@@ -146,7 +146,8 @@ public SimilarityMeasure similarity() {
         return similarity;
     }
 
-    public boolean dimensionsSetByUser() {
+    @Override
+    public Boolean dimensionsSetByUser() {
         return this.dimensionsSetByUser;
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiService.java

@@ -11,8 +11,8 @@
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.common.Strings;
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.Strings;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.inference.ChunkedInferenceServiceResults;
 import org.elasticsearch.inference.ChunkingOptions;
@@ -31,10 +31,10 @@
 import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
 import org.elasticsearch.xpack.inference.services.SenderService;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
-import org.elasticsearch.xpack.inference.services.ServiceUtils;
 import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsModel;
 import org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsServiceSettings;
+import org.elasticsearch.xpack.inference.services.validation.ModelValidatorBuilder;
 
 import java.util.List;
 import java.util.Map;
@@ -255,48 +255,35 @@ protected void doChunkedInfer(
      */
     @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 OpenAiEmbeddingsModel embeddingsModel) {
-            ServiceUtils.getEmbeddingSize(
-                model,
-                this,
-                listener.delegateFailureAndWrap((l, size) -> l.onResponse(updateModelWithEmbeddingDetails(embeddingsModel, size)))
+            var serviceSettings = embeddingsModel.getServiceSettings();
+            var similarityFromModel = serviceSettings.similarity();
+            var similarityToUse = similarityFromModel == null ? SimilarityMeasure.DOT_PRODUCT : similarityFromModel;
+
+            var updatedServiceSettings = new OpenAiEmbeddingsServiceSettings(
+                serviceSettings.modelId(),
+                serviceSettings.uri(),
+                serviceSettings.organizationId(),
+                similarityToUse,
+                embeddingSize,
+                serviceSettings.maxInputTokens(),
+                serviceSettings.dimensionsSetByUser(),
+                serviceSettings.rateLimitSettings()
             );
-        } else {
-            listener.onResponse(model);
-        }
-    }
 
-    private OpenAiEmbeddingsModel updateModelWithEmbeddingDetails(OpenAiEmbeddingsModel model, int embeddingSize) {
-        if (model.getServiceSettings().dimensionsSetByUser()
-            && model.getServiceSettings().dimensions() != null
-            && model.getServiceSettings().dimensions() != embeddingSize) {
+            return new OpenAiEmbeddingsModel(embeddingsModel, updatedServiceSettings);
+        } else {
             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()
-                ),
+                Strings.format("Can't update embedding details for model with unexpected type %s", model.getClass()),
                 RestStatus.BAD_REQUEST
             );
         }
-
-        var similarityFromModel = model.getServiceSettings().similarity();
-        var similarityToUse = similarityFromModel == null ? SimilarityMeasure.DOT_PRODUCT : similarityFromModel;
-
-        OpenAiEmbeddingsServiceSettings serviceSettings = new OpenAiEmbeddingsServiceSettings(
-            model.getServiceSettings().modelId(),
-            model.getServiceSettings().uri(),
-            model.getServiceSettings().organizationId(),
-            similarityToUse,
-            embeddingSize,
-            model.getServiceSettings().maxInputTokens(),
-            model.getServiceSettings().dimensionsSetByUser(),
-            model.getServiceSettings().rateLimitSettings()
-        );
-
-        return new OpenAiEmbeddingsModel(model, serviceSettings);
     }
 
     @Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsServiceSettings.java

@@ -247,6 +247,7 @@ public Integer dimensions() {
         return dimensions;
     }
 
+    @Override
     public Boolean dimensionsSetByUser() {
         return dimensionsSetByUser;
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/validation/ModelValidator.java

@@ -0,0 +1,16 @@
+/*
+ * 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.validation;
+
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.inference.InferenceService;
+import org.elasticsearch.inference.Model;
+
+public interface ModelValidator {
+    void validate(InferenceService service, Model model, ActionListener<Model> listener);
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/validation/ModelValidatorBuilder.java

@@ -0,0 +1,29 @@
+/*
+ * 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.validation;
+
+import org.elasticsearch.core.Strings;
+import org.elasticsearch.inference.TaskType;
+
+public class ModelValidatorBuilder {
+    public static ModelValidator buildModelValidator(TaskType taskType) {
+        if (taskType == null) {
+            throw new IllegalArgumentException("Task type can't be null");
+        }
+
+        switch (taskType) {
+            case TEXT_EMBEDDING -> {
+                return new TextEmbeddingModelValidator(new SimpleServiceIntegrationValidator());
+            }
+            case SPARSE_EMBEDDING, RERANK, COMPLETION, ANY -> {
+                return new SimpleModelValidator(new SimpleServiceIntegrationValidator());
+            }
+            default -> throw new IllegalArgumentException(Strings.format("Can't validate inference model of for task type %s ", taskType));
+        }
+    }
+}