More tests

Author: jonathan-buttner

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

@@ -21,6 +21,7 @@
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
 import org.elasticsearch.xpack.inference.services.settings.ApiKeySecrets;
+import org.elasticsearch.xpack.inference.services.settings.SerializableSecureString;
 
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -557,17 +558,18 @@ public static void validateMapValues(
                     );
                 } else {
                     return Strings.format(
-                        "Map field [%s] has an entry that is not valid, [%s => %s]. Value type is not one of [%s].",
+                        "Map field [%s] has an entry that is not valid, [%s => %s]. Value type of [%s] is not one of [%s].",
                         settingName,
                         entry.getKey(),
                         entry.getValue(),
+                        entry.getValue(),
                         String.join(", ", validTypesAsStrings)
                     );
                 }
             };
 
             if (isAllowed == false) {
-                var validTypesAsStrings = allowedTypes.stream().map(Class::toString).toArray(String[]::new);
+                var validTypesAsStrings = allowedTypes.stream().map(Class::getSimpleName).toArray(String[]::new);
                 Arrays.sort(validTypesAsStrings);
 
                 validationException.addValidationError(errorMessage.apply(validTypesAsStrings));
@@ -576,7 +578,7 @@ public static void validateMapValues(
         }
     }
 
-    public static Map<String, SecureString> convertMapStringsToSecureString(
+    public static Map<String, SerializableSecureString> convertMapStringsToSecureString(
         Map<String, ?> map,
         String settingName,
         ValidationException validationException
@@ -585,22 +587,24 @@ public static Map<String, SecureString> convertMapStringsToSecureString(
             return Map.of();
         }
 
-        validateMapValues(map, List.of(String.class), settingName, validationException, true);
+        validateMapStringValues(map, settingName, validationException, true);
 
         return map.entrySet()
             .stream()
-            .collect(Collectors.toMap(Map.Entry::getKey, e -> new SecureString(((String) e.getValue()).toCharArray())));
+            .collect(Collectors.toMap(Map.Entry::getKey, e -> new SerializableSecureString((String) e.getValue())));
     }
 
     /**
      * Removes null values.
      */
-    public static void removeNullValues(Map<String, Object> map) {
+    public static Map<String, Object> removeNullValues(Map<String, Object> map) {
         if (map == null) {
-            return;
+            return map;
         }
 
         map.values().removeIf(Objects::isNull);
+
+        return map;
     }
 
     public static Integer extractRequiredPositiveIntegerLessThanOrEqualToMax(

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

@@ -12,10 +12,10 @@
 import org.elasticsearch.common.ValidationException;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.common.settings.SecureString;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.inference.SecretSettings;
 import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xpack.inference.services.settings.SerializableSecureString;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -48,22 +48,22 @@ public static CustomSecretSettings fromMap(@Nullable Map<String, Object> map) {
         return new CustomSecretSettings(secureStringMap);
     }
 
-    private final Map<String, SecureString> secretParameters;
+    private final Map<String, SerializableSecureString> secretParameters;
 
     @Override
     public SecretSettings newSecretSettings(Map<String, Object> newSecrets) {
         return fromMap(new HashMap<>(newSecrets));
     }
 
-    public CustomSecretSettings(@Nullable Map<String, SecureString> secretParameters) {
+    public CustomSecretSettings(@Nullable Map<String, SerializableSecureString> secretParameters) {
         this.secretParameters = Objects.requireNonNullElse(secretParameters, Map.of());
     }
 
     public CustomSecretSettings(StreamInput in) throws IOException {
-        secretParameters = in.readImmutableMap(StreamInput::readString, StreamInput::readSecureString);
+        secretParameters = in.readImmutableMap(SerializableSecureString::new);
     }
 
-    public Map<String, SecureString> getSecretParameters() {
+    public Map<String, SerializableSecureString> getSecretParameters() {
         return secretParameters;
     }
 
@@ -74,7 +74,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.startObject(SECRET_PARAMETERS);
             {
                 for (var entry : secretParameters.entrySet()) {
-                    builder.field(entry.getKey(), entry.getValue().toString());
+                    builder.field(entry.getKey(), entry.getValue());
                 }
             }
             builder.endObject();
@@ -95,7 +95,9 @@ public TransportVersion getMinimalSupportedVersion() {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeMap(secretParameters, StreamOutput::writeString, StreamOutput::writeSecureString);
+        out.writeMap(secretParameters, (streamOutput, v) -> {
+            v.writeTo(streamOutput);
+        });
     }
 
     @Override

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

@@ -22,7 +22,6 @@
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
-import org.elasticsearch.xpack.inference.services.custom.response.BaseCustomResponseParser;
 import org.elasticsearch.xpack.inference.services.custom.response.CompletionResponseParser;
 import org.elasticsearch.xpack.inference.services.custom.response.CustomResponseParser;
 import org.elasticsearch.xpack.inference.services.custom.response.ErrorResponseParser;
@@ -183,7 +182,7 @@ public CustomServiceSettings(StreamInput in) throws IOException {
         url = in.readString();
         headers = in.readImmutableMap(StreamInput::readString);
         requestContentString = in.readString();
-        responseJsonParser = in.readNamedWriteable(BaseCustomResponseParser.class);
+        responseJsonParser = in.readNamedWriteable(CustomResponseParser.class);
         rateLimitSettings = new RateLimitSettings(in);
         errorParser = new ErrorResponseParser(in);
     }

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

@@ -64,7 +64,8 @@ public CustomRequest(String query, List<String> input, CustomModel model) {
         uri = buildUri();
     }
 
-    private static <T> String toJson(T value, String field) {
+    // default for testing
+    static <T> String toJson(T value, String field) {
         try {
             XContentBuilder builder = JsonXContent.contentBuilder();
             // TODO test this, I think it'll write the quotes for us so we don't need to include them in the content string

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.inference.services.custom.response;
 
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.ValidationException;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -15,6 +16,7 @@
 import org.elasticsearch.xpack.inference.common.MapPathExtractor;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
@@ -79,11 +81,22 @@ public String getWriteableName() {
 
     @Override
     public ChatCompletionResults transform(Map<String, Object> map) {
-        var completionList = validateAndCastList(
-            validateList(MapPathExtractor.extract(map, completionResultPath)),
-            (obj) -> toType(obj, String.class)
-        );
-
-        return new ChatCompletionResults(completionList.stream().map(ChatCompletionResults.Result::new).toList());
+        var extractedField = MapPathExtractor.extract(map, completionResultPath);
+
+        validateNonNull(extractedField);
+
+        if (extractedField instanceof List<?> extractedList) {
+            var completionList = validateAndCastList(extractedList, (obj) -> toType(obj, String.class));
+            return new ChatCompletionResults(completionList.stream().map(ChatCompletionResults.Result::new).toList());
+        } else if (extractedField instanceof String extractedString) {
+            return new ChatCompletionResults(List.of(new ChatCompletionResults.Result(extractedString)));
+        } else {
+            throw new IllegalArgumentException(
+                Strings.format(
+                    "Extracted field is an invalid type, expected a list or a string but received [%s]",
+                    extractedField.getClass().getSimpleName()
+                )
+            );
+        }
     }
 }