[ML] Custom service add support for input_type, top_n, and return_documents (elastic#129441) (elastic#129860)

* Making progress on different request parameters

* Working tests

* Adding custom service validator for rerank

* Fixing embedding bug

* Adding transport version check

* Fixing tests

* Fixing license header

* Fixing writeTo

* Moving file and removing commented code

* Fixing test

* Fixing tests

* Refactoring and tests

* Fixing test

(cherry picked from commit d9b34d4)

# Conflicts:
#	server/src/main/java/org/elasticsearch/TransportVersions.java

Author: jonathan-buttner

server/src/main/java/module-info.java

@@ -476,6 +476,7 @@
             org.elasticsearch.serverless.apifiltering;
     exports org.elasticsearch.lucene.spatial;
     exports org.elasticsearch.inference.configuration;
+    exports org.elasticsearch.inference.validation;
     exports org.elasticsearch.monitor.metrics;
     exports org.elasticsearch.plugins.internal.rewriter to org.elasticsearch.inference;
     exports org.elasticsearch.index.codec.perfield;

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

@@ -245,6 +245,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_REMOVE_ERROR_PARSING_8_19 = def(8_841_0_52);
     public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_EMBEDDING_BATCH_SIZE_8_19 = def(8_841_0_53);
     public static final TransportVersion STREAMS_LOGS_SUPPORT_8_19 = def(8_841_0_54);
+    public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_INPUT_TYPE_8_19 = def(8_841_0_55);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -14,6 +14,7 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.inference.validation.ServiceIntegrationValidator;
 
 import java.io.Closeable;
 import java.util.EnumSet;
@@ -248,4 +249,14 @@ default void updateModelsWithDynamicFields(List<Model> model, ActionListener<Lis
      * after ensuring the node's internals are set up (for example if this ensures the internal ES client is ready for use).
      */
     default void onNodeStarted() {}
+
+    /**
+     * Get the service integration validator for the given task type.
+     * This allows services to provide custom validation logic.
+     * @param taskType The task type
+     * @return The service integration validator or null if the default should be used
+     */
+    default ServiceIntegrationValidator getServiceIntegrationValidator(TaskType taskType) {
+        return null;
+    }
 }

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

@@ -10,8 +10,12 @@
 package org.elasticsearch.inference;
 
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.ValidationException;
 
+import java.util.EnumSet;
+import java.util.HashMap;
 import java.util.Locale;
+import java.util.Map;
 
 import static org.elasticsearch.core.Strings.format;
 
@@ -29,6 +33,13 @@ public enum InputType {
     INTERNAL_SEARCH,
     INTERNAL_INGEST;
 
+    private static final EnumSet<InputType> SUPPORTED_REQUEST_VALUES = EnumSet.of(
+        InputType.CLASSIFICATION,
+        InputType.CLUSTERING,
+        InputType.INGEST,
+        InputType.SEARCH
+    );
+
     @Override
     public String toString() {
         return name().toLowerCase(Locale.ROOT);
@@ -57,4 +68,70 @@ public static boolean isSpecified(InputType inputType) {
     public static String invalidInputTypeMessage(InputType inputType) {
         return Strings.format("received invalid input type value [%s]", inputType.toString());
     }
+
+    /**
+     * Ensures that a map used for translating input types is valid. The keys of the map are the external representation,
+     * and the values correspond to the values in this class.
+     * Throws a {@link ValidationException} if any value is not a valid InputType.
+     *
+     * @param inputTypeTranslation the map of input type translations to validate
+     * @param validationException  a ValidationException to which errors will be added
+     */
+    public static Map<InputType, String> validateInputTypeTranslationValues(
+        Map<String, Object> inputTypeTranslation,
+        ValidationException validationException
+    ) {
+        if (inputTypeTranslation == null || inputTypeTranslation.isEmpty()) {
+            return Map.of();
+        }
+
+        var translationMap = new HashMap<InputType, String>();
+
+        for (var entry : inputTypeTranslation.entrySet()) {
+            var key = entry.getKey();
+            var value = entry.getValue();
+
+            if (value instanceof String == false || Strings.isNullOrEmpty((String) value)) {
+                validationException.addValidationError(
+                    Strings.format(
+                        "Input type translation value for key [%s] must be a String that is "
+                            + "not null and not empty, received: [%s], type: [%s].",
+                        key,
+                        value,
+                        value == null ? "null" : value.getClass().getSimpleName()
+                    )
+                );
+
+                throw validationException;
+            }
+
+            try {
+                var inputTypeKey = InputType.fromStringValidateSupportedRequestValue(key);
+                translationMap.put(inputTypeKey, (String) value);
+            } catch (Exception e) {
+                validationException.addValidationError(
+                    Strings.format(
+                        "Invalid input type translation for key: [%s], is not a valid value. Must be one of %s",
+                        key,
+                        SUPPORTED_REQUEST_VALUES
+                    )
+                );
+
+                throw validationException;
+            }
+        }
+
+        return translationMap;
+    }
+
+    private static InputType fromStringValidateSupportedRequestValue(String name) {
+        var inputType = fromRestString(name);
+        if (SUPPORTED_REQUEST_VALUES.contains(inputType) == false) {
+            throw new IllegalArgumentException(
+                format("Unrecognized input_type [%s], must be one of %s", inputType, SUPPORTED_REQUEST_VALUES)
+            );
+        }
+
+        return inputType;
+    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/validation/ServiceIntegrationValidator.java renamed to server/src/main/java/org/elasticsearch/inference/validation/ServiceIntegrationValidator.java

@@ -1,11 +1,13 @@
 /*
  * 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.
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-package org.elasticsearch.xpack.inference.services.validation;
+package org.elasticsearch.inference.validation;
 
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.core.TimeValue;

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

@@ -216,7 +216,7 @@ private void parseAndStoreModel(
             if (skipValidationAndStart) {
                 storeModelListener.onResponse(model);
             } else {
-                ModelValidatorBuilder.buildModelValidator(model.getTaskType(), service instanceof ElasticsearchInternalService)
+                ModelValidatorBuilder.buildModelValidator(model.getTaskType(), service)
                     .validate(service, model, timeout, storeModelListener);
             }
         });

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

@@ -399,6 +399,18 @@ public static String extractRequiredString(
         return requiredField;
     }
 
+    public static String extractOptionalEmptyString(Map<String, Object> map, String settingName, ValidationException validationException) {
+        int initialValidationErrorCount = validationException.validationErrors().size();
+        String optionalField = ServiceUtils.removeAsType(map, settingName, String.class, validationException);
+
+        if (validationException.validationErrors().size() > initialValidationErrorCount) {
+            // new validation error occurred
+            return null;
+        }
+
+        return optionalField;
+    }
+
     public static String extractOptionalString(
         Map<String, Object> map,
         String settingName,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomRequestManager.java

@@ -23,10 +23,13 @@
 import org.elasticsearch.xpack.inference.external.http.sender.ExecutableInferenceRequest;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.QueryAndDocsInputs;
+import org.elasticsearch.xpack.inference.services.custom.request.CompletionParameters;
 import org.elasticsearch.xpack.inference.services.custom.request.CustomRequest;
+import org.elasticsearch.xpack.inference.services.custom.request.EmbeddingParameters;
+import org.elasticsearch.xpack.inference.services.custom.request.RequestParameters;
+import org.elasticsearch.xpack.inference.services.custom.request.RerankParameters;
 import org.elasticsearch.xpack.inference.services.custom.response.CustomResponseEntity;
 
-import java.util.List;
 import java.util.Objects;
 import java.util.function.Supplier;
 
@@ -65,19 +68,16 @@ public void execute(
         Supplier<Boolean> hasRequestCompletedFunction,
         ActionListener<InferenceServiceResults> listener
     ) {
-        String query;
-        List<String> input;
+        RequestParameters requestParameters;
         if (inferenceInputs instanceof QueryAndDocsInputs) {
-            QueryAndDocsInputs queryAndDocsInputs = QueryAndDocsInputs.of(inferenceInputs);
-            query = queryAndDocsInputs.getQuery();
-            input = queryAndDocsInputs.getChunks();
+            requestParameters = RerankParameters.of(QueryAndDocsInputs.of(inferenceInputs));
         } else if (inferenceInputs instanceof ChatCompletionInput chatInputs) {
-            query = null;
-            input = chatInputs.getInputs();
+            requestParameters = CompletionParameters.of(chatInputs);
         } else if (inferenceInputs instanceof EmbeddingsInput) {
-            EmbeddingsInput embeddingsInput = EmbeddingsInput.of(inferenceInputs);
-            query = null;
-            input = embeddingsInput.getStringInputs();
+            requestParameters = EmbeddingParameters.of(
+                EmbeddingsInput.of(inferenceInputs),
+                model.getServiceSettings().getInputTypeTranslator()
+            );
         } else {
             listener.onFailure(
                 new ElasticsearchStatusException(
@@ -89,7 +89,7 @@ public void execute(
         }
 
         try {
-            var request = new CustomRequest(query, input, model);
+            var request = new CustomRequest(requestParameters, model);
             execute(new ExecutableInferenceRequest(requestSender, logger, request, handler, hasRequestCompletedFunction, listener));
         } catch (Exception e) {
             // Intentionally not logging this exception because it could contain sensitive information from the CustomRequest construction

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java

@@ -27,19 +27,26 @@
 import org.elasticsearch.inference.SettingsConfiguration;
 import org.elasticsearch.inference.SimilarityMeasure;
 import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.inference.validation.ServiceIntegrationValidator;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xpack.inference.chunking.ChunkingSettingsBuilder;
 import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
 import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
+import org.elasticsearch.xpack.inference.external.http.sender.ChatCompletionInput;
 import org.elasticsearch.xpack.inference.external.http.sender.EmbeddingsInput;
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
+import org.elasticsearch.xpack.inference.external.http.sender.QueryAndDocsInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
 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.custom.request.CompletionParameters;
 import org.elasticsearch.xpack.inference.services.custom.request.CustomRequest;
+import org.elasticsearch.xpack.inference.services.custom.request.EmbeddingParameters;
+import org.elasticsearch.xpack.inference.services.custom.request.RequestParameters;
+import org.elasticsearch.xpack.inference.services.custom.request.RerankParameters;
+import org.elasticsearch.xpack.inference.services.validation.CustomServiceIntegrationValidator;
 
 import java.util.EnumSet;
 import java.util.HashMap;
@@ -115,20 +122,29 @@ public void parseRequestConfig(
      * This does some initial validation with mock inputs to determine if any templates are missing a field to fill them.
      */
     private static void validateConfiguration(CustomModel model) {
-        String query = null;
-        if (model.getTaskType() == TaskType.RERANK) {
-            query = "test query";
-        }
-
         try {
-            new CustomRequest(query, List.of("test input"), model).createHttpRequest();
+            new CustomRequest(createParameters(model), model).createHttpRequest();
         } catch (IllegalStateException e) {
             var validationException = new ValidationException();
             validationException.addValidationError(Strings.format("Failed to validate model configuration: %s", e.getMessage()));
             throw validationException;
         }
     }
 
+    private static RequestParameters createParameters(CustomModel model) {
+        return switch (model.getTaskType()) {
+            case RERANK -> RerankParameters.of(new QueryAndDocsInputs("test query", List.of("test input")));
+            case COMPLETION -> CompletionParameters.of(new ChatCompletionInput(List.of("test input")));
+            case TEXT_EMBEDDING, SPARSE_EMBEDDING -> EmbeddingParameters.of(
+                new EmbeddingsInput(List.of("test input"), null, null),
+                model.getServiceSettings().getInputTypeTranslator()
+            );
+            default -> throw new IllegalStateException(
+                Strings.format("Unsupported task type [%s] for custom service", model.getTaskType())
+            );
+        };
+    }
+
     private static ChunkingSettings extractChunkingSettings(Map<String, Object> config, TaskType taskType) {
         if (TaskType.TEXT_EMBEDDING.equals(taskType)) {
             return ChunkingSettingsBuilder.fromMap(removeFromMap(config, ModelConfigurations.CHUNKING_SETTINGS));
@@ -257,7 +273,8 @@ public void doInfer(
 
     @Override
     protected void validateInputType(InputType inputType, Model model, ValidationException validationException) {
-        ServiceUtils.validateInputTypeIsUnspecifiedOrInternal(inputType, validationException);
+        // The custom service doesn't do any validation for the input type because if the input type is supported a default
+        // must be supplied within the service settings.
     }
 
     @Override
@@ -327,7 +344,9 @@ private static CustomServiceSettings getCustomServiceSettings(CustomModel custom
             serviceSettings.getQueryParameters(),
             serviceSettings.getRequestContentString(),
             serviceSettings.getResponseJsonParser(),
-            serviceSettings.rateLimitSettings()
+            serviceSettings.rateLimitSettings(),
+            serviceSettings.getBatchSize(),
+            serviceSettings.getInputTypeTranslator()
         );
     }
 
@@ -353,4 +372,13 @@ public static InferenceServiceConfiguration get() {
             }
         );
     }
+
+    @Override
+    public ServiceIntegrationValidator getServiceIntegrationValidator(TaskType taskType) {
+        if (taskType == TaskType.RERANK) {
+            return new CustomServiceIntegrationValidator();
+        }
+
+        return null;
+    }
 }