Refactoring the request

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

@@ -58,6 +58,7 @@
 import org.elasticsearch.xpack.inference.services.cohere.embeddings.CohereEmbeddingsTaskSettings;
 import org.elasticsearch.xpack.inference.services.cohere.rerank.CohereRerankServiceSettings;
 import org.elasticsearch.xpack.inference.services.cohere.rerank.CohereRerankTaskSettings;
+import org.elasticsearch.xpack.inference.services.custom.CustomSecretSettings;
 import org.elasticsearch.xpack.inference.services.custom.CustomServiceSettings;
 import org.elasticsearch.xpack.inference.services.custom.CustomTaskSettings;
 import org.elasticsearch.xpack.inference.services.deepseek.DeepSeekChatCompletionModel;
@@ -153,7 +154,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         addAlibabaCloudSearchNamedWriteables(namedWriteables);
         addJinaAINamedWriteables(namedWriteables);
         addVoyageAINamedWriteables(namedWriteables);
-        addCustomWriteables(namedWriteables);
+        addCustomNamedWriteables(namedWriteables);
 
         addUnifiedNamedWriteables(namedWriteables);
 
@@ -163,6 +164,32 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         return namedWriteables;
     }
 
+    private static void addCustomNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                CustomServiceSettings.NAME,
+                CustomServiceSettings::new
+            )
+        );
+
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                TaskSettings.class,
+                CustomTaskSettings.NAME,
+                CustomTaskSettings::new
+            )
+        );
+
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                SecretSettings.class,
+                CustomSecretSettings.NAME,
+                CustomSecretSettings::new
+            )
+        );
+    }
+
     private static void addUnifiedNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
         var writeables = UnifiedCompletionRequest.getNamedWriteables();
         namedWriteables.addAll(writeables);
@@ -665,11 +692,4 @@ private static void addEisNamedWriteables(List<NamedWriteableRegistry.Entry> nam
             )
         );
     }
-
-    private static void addCustomWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
-        namedWriteables.add(
-            new NamedWriteableRegistry.Entry(ServiceSettings.class, CustomServiceSettings.NAME, CustomServiceSettings::new)
-        );
-        namedWriteables.add(new NamedWriteableRegistry.Entry(TaskSettings.class, CustomTaskSettings.NAME, CustomTaskSettings::new));
-    }
 }

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

@@ -31,6 +31,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static org.elasticsearch.core.Strings.format;
@@ -479,29 +480,86 @@ public static Map<String, Object> extractOptionalMap(
     }
 
     /**
-     * Ensures that each value in the map is a {@link String}.
-     * @param map a Map to iterate over
-     * @param settingName the setting name that his map corresponds to
-     * @param validationException aggregated validation exceptions
+     * Ensures the values of the map match one of the supplied types.
+     * @param map Map to validate
+     * @param allowedTypes List of {@link Class} to accept
+     * @param settingName the setting name for the field
+     * @param validationException exception to return if one of the values is invalid
+     * @param censorValue if true the key and value will be included in the exception message
      */
-    public static void validateMapValueStrings(
+    public static void validateMapValues(
         Map<String, Object> map,
+        List<Class<?>> allowedTypes,
         String settingName,
-        ValidationException validationException
+        ValidationException validationException,
+        boolean censorValue
     ) {
         if (map == null) {
             return;
         }
 
         for (var entry : map.entrySet()) {
-            var value = entry.getValue();
-            if (value instanceof String == false) {
-                validationException.addValidationError(ServiceUtils.invalidTypeErrorMsg(settingName, value, String.class.getSimpleName()));
+            boolean isAllowed = false;
+
+            for (Class<?> allowedType : allowedTypes) {
+                if (allowedType.isInstance(entry.getValue())) {
+                    isAllowed = true;
+                    break;
+                }
+            }
+
+            Function<String[], String> errorMessage = (String[] validTypesAsStrings) -> {
+                if (censorValue) {
+                    return Strings.format(
+                        "Map field [%s] has an entry that is not valid. Value type is not one of [%s].",
+                        settingName,
+                        String.join(", ", validTypesAsStrings)
+                    );
+                } else {
+                    return Strings.format(
+                        "Map field [%s] has an entry that is not valid, [%s => %s]. Value type is not one of [%s].",
+                        settingName,
+                        entry.getKey(),
+                        entry.getValue(),
+                        String.join(", ", validTypesAsStrings)
+                    );
+                }
+            };
+
+            if (isAllowed == false) {
+                var validTypesAsStrings = allowedTypes.stream().map(Class::toString).toArray(String[]::new);
+                Arrays.sort(validTypesAsStrings);
+
+                validationException.addValidationError(errorMessage.apply(validTypesAsStrings));
                 throw validationException;
             }
         }
     }
 
+    public static void convertMapStringsToSecureString(Map<String, Object> map) {
+        if (map == null) {
+            return;
+        }
+
+        for (var entry : map.entrySet()) {
+            var value = entry.getValue();
+            if (value instanceof String) {
+                map.put(entry.getKey(), new SecureString(((String) value).toCharArray()));
+            }
+        }
+    }
+
+    /**
+     * Removes null values.
+     */
+    public static void removeNullValues(Map<String, Object> map) {
+        if (map == null) {
+            return;
+        }
+
+        map.values().removeIf(Objects::isNull);
+    }
+
     public static Integer extractRequiredPositiveIntegerLessThanOrEqualToMax(
         Map<String, Object> map,
         String settingName,

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

@@ -18,10 +18,14 @@
 
 import java.io.IOException;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.convertMapStringsToSecureString;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalMap;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeNullValues;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.validateMapValues;
 
 public class CustomSecretSettings implements SecretSettings {
     public static final String NAME = "custom_secret_settings";
@@ -36,20 +40,15 @@ public static CustomSecretSettings fromMap(@Nullable Map<String, Object> map) {
         ValidationException validationException = new ValidationException();
 
         Map<String, Object> requestSecretParamsMap = extractOptionalMap(map, SECRET_PARAMETERS, NAME, validationException);
+        removeNullValues(requestSecretParamsMap);
+        validateMapValues(requestSecretParamsMap, List.of(String.class), SECRET_PARAMETERS, validationException, true);
+        convertMapStringsToSecureString(requestSecretParamsMap);
+
         if (validationException.validationErrors().isEmpty() == false) {
             throw validationException;
         }
 
-        if (requestSecretParamsMap == null) {
-            return null;
-        } else {
-            Map<String, Object> secureSecretParameters = new HashMap<>();
-            for (String paramKey : requestSecretParamsMap.keySet()) {
-                Object paramValue = requestSecretParamsMap.get(paramKey);
-                secureSecretParameters.put(paramKey, paramValue);
-            }
-            return new CustomSecretSettings(secureSecretParameters);
-        }
+        return new CustomSecretSettings(Objects.requireNonNullElse(requestSecretParamsMap, new HashMap<>()));
     }
 
     @Override
@@ -58,15 +57,11 @@ public SecretSettings newSecretSettings(Map<String, Object> newSecrets) {
     }
 
     public CustomSecretSettings(@Nullable Map<String, Object> secretParameters) {
-        this.secretParameters = secretParameters;
+        this.secretParameters = Objects.requireNonNullElse(secretParameters, new HashMap<>());
     }
 
     public CustomSecretSettings(StreamInput in) throws IOException {
-        if (in.readBoolean()) {
-            secretParameters = in.readGenericMap();
-        } else {
-            secretParameters = null;
-        }
+        secretParameters = in.readGenericMap();
     }
 
     public Map<String, Object> getSecretParameters() {
@@ -76,7 +71,7 @@ public Map<String, Object> getSecretParameters() {
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
-        if (secretParameters != null) {
+        if (secretParameters.isEmpty() == false) {
             builder.field(SECRET_PARAMETERS, secretParameters);
         }
         builder.endObject();
@@ -95,12 +90,7 @@ public TransportVersion getMinimalSupportedVersion() {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        if (secretParameters != null) {
-            out.writeBoolean(true);
-            out.writeGenericMap(secretParameters);
-        } else {
-            out.writeBoolean(false);
-        }
+        out.writeGenericMap(secretParameters);
     }
 
     @Override

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

@@ -32,6 +32,7 @@
 
 import java.io.IOException;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
@@ -43,8 +44,9 @@
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractRequiredString;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractSimilarity;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeAsType;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeNullValues;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.throwIfNotEmptyMap;
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.validateMapValueStrings;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.validateMapValues;
 
 public class CustomServiceSettings extends FilteredXContentObject implements ServiceSettings, CustomRateLimitServiceSettings {
     public static final String NAME = "custom_service_settings";
@@ -84,7 +86,7 @@ private record Fields(
         RateLimitSettings rateLimitSettings
     ) {
         public void validate(ValidationException validationException) {
-            validateMapValueStrings(headers, HEADERS, validationException);
+            validateMapValues(headers, List.of(String.class), HEADERS, validationException, false);
 
             if (requestBodyMap == null || responseParserMap == null || jsonParserMap == null) {
                 throw validationException;
@@ -113,6 +115,7 @@ private static Fields from(
         String url = extractRequiredString(map, URL, ModelConfigurations.SERVICE_SETTINGS, validationException);
 
         Map<String, Object> headers = extractOptionalMap(map, HEADERS, ModelConfigurations.SERVICE_SETTINGS, validationException);
+        removeNullValues(headers);
 
         Map<String, Object> requestBodyMap = extractRequiredMap(map, REQUEST, ModelConfigurations.SERVICE_SETTINGS, validationException);
 
@@ -165,10 +168,10 @@ private static Fields from(
     private static CustomServiceSettings fromRequestMap(Map<String, Object> map, TaskType taskType) {
         ValidationException validationException = new ValidationException();
 
-        var serviceSettings = from(map, ConfigurationParseContext.REQUEST, taskType, validationException);
+        var serviceSettingsFields = from(map, ConfigurationParseContext.REQUEST, taskType, validationException);
 
-        serviceSettings.validate(validationException);
-        return CustomServiceSettings.of(serviceSettings);
+        serviceSettingsFields.validate(validationException);
+        return CustomServiceSettings.of(serviceSettingsFields);
     }
 
     private static CustomServiceSettings of(Fields fields) {