Adding more tests and cleaning up

Author: jonathan-buttner

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

@@ -46,7 +46,7 @@ public CustomModel(
             inferenceId,
             taskType,
             service,
-            CustomServiceSettings.fromMap(serviceSettings, context, taskType, inferenceId),
+            CustomServiceSettings.fromMap(serviceSettings, context, taskType),
             CustomTaskSettings.fromMap(taskSettings),
             CustomSecretSettings.fromMap(secrets)
         );
@@ -66,7 +66,7 @@ public CustomModel(
             inferenceId,
             taskType,
             service,
-            CustomServiceSettings.fromMap(serviceSettings, context, taskType, inferenceId),
+            CustomServiceSettings.fromMap(serviceSettings, context, taskType),
             CustomTaskSettings.fromMap(taskSettings),
             CustomSecretSettings.fromMap(secrets),
             chunkingSettings

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

@@ -333,7 +333,7 @@ private static CustomServiceSettings getCustomServiceSettings(CustomModel custom
         var similarityToUse = similarityFromModel == null ? SimilarityMeasure.DOT_PRODUCT : similarityFromModel;
 
         return new CustomServiceSettings(
-            new CustomServiceSettings.TextEmbeddingSettings(similarityToUse, embeddingSize, serviceSettings.getMaxInputTokens(), null),
+            new CustomServiceSettings.TextEmbeddingSettings(similarityToUse, embeddingSize, serviceSettings.getMaxInputTokens()),
             serviceSettings.getUrl(),
             serviceSettings.getHeaders(),
             serviceSettings.getQueryParameters(),

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

@@ -66,12 +66,7 @@ public class CustomServiceSettings extends FilteredXContentObject implements Ser
     private static final String RESPONSE_SCOPE = String.join(".", ModelConfigurations.SERVICE_SETTINGS, RESPONSE);
     private static final int DEFAULT_EMBEDDING_BATCH_SIZE = 10;
 
-    public static CustomServiceSettings fromMap(
-        Map<String, Object> map,
-        ConfigurationParseContext context,
-        TaskType taskType,
-        String inferenceId
-    ) {
+    public static CustomServiceSettings fromMap(Map<String, Object> map, ConfigurationParseContext context, TaskType taskType) {
         ValidationException validationException = new ValidationException();
 
         var textEmbeddingSettings = TextEmbeddingSettings.fromMap(map, taskType, validationException);
@@ -140,14 +135,9 @@ public static CustomServiceSettings fromMap(
     public static class TextEmbeddingSettings implements ToXContentFragment, Writeable {
 
         // This specifies float for the element type but null for all other settings
-        public static final TextEmbeddingSettings DEFAULT_FLOAT = new TextEmbeddingSettings(
-            null,
-            null,
-            null,
-            CustomServiceEmbeddingType.FLOAT
-        );
+        public static final TextEmbeddingSettings DEFAULT_FLOAT = new TextEmbeddingSettings(null, null, null);
         // This refers to settings that are not related to the text embedding task type (all the settings should be null)
-        public static final TextEmbeddingSettings NON_TEXT_EMBEDDING_TASK_TYPE_SETTINGS = new TextEmbeddingSettings(null, null, null, null);
+        public static final TextEmbeddingSettings NON_TEXT_EMBEDDING_TASK_TYPE_SETTINGS = new TextEmbeddingSettings(null, null, null);
 
         public static TextEmbeddingSettings fromMap(Map<String, Object> map, TaskType taskType, ValidationException validationException) {
             if (taskType != TaskType.TEXT_EMBEDDING) {
@@ -157,7 +147,7 @@ public static TextEmbeddingSettings fromMap(Map<String, Object> map, TaskType ta
             SimilarityMeasure similarity = extractSimilarity(map, ModelConfigurations.SERVICE_SETTINGS, validationException);
             Integer dims = removeAsType(map, DIMENSIONS, Integer.class);
             Integer maxInputTokens = removeAsType(map, MAX_INPUT_TOKENS, Integer.class);
-            return new TextEmbeddingSettings(similarity, dims, maxInputTokens, null);
+            return new TextEmbeddingSettings(similarity, dims, maxInputTokens);
         }
 
         private final SimilarityMeasure similarityMeasure;
@@ -167,8 +157,7 @@ public static TextEmbeddingSettings fromMap(Map<String, Object> map, TaskType ta
         public TextEmbeddingSettings(
             @Nullable SimilarityMeasure similarityMeasure,
             @Nullable Integer dimensions,
-            @Nullable Integer maxInputTokens,
-            @Nullable CustomServiceEmbeddingType embeddingType
+            @Nullable Integer maxInputTokens
         ) {
             this.similarityMeasure = similarityMeasure;
             this.dimensions = dimensions;
@@ -331,7 +320,12 @@ public Integer dimensions() {
 
     @Override
     public DenseVectorFieldMapper.ElementType elementType() {
-        return responseJsonParser.getEmbeddingType().toElementType();
+        var embeddingType = responseJsonParser.getEmbeddingType();
+        if (embeddingType != null) {
+            return embeddingType.toElementType();
+        }
+
+        return null;
     }
 
     public Integer getMaxInputTokens() {

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

@@ -106,11 +106,12 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         return builder;
     }
 
-    // For testing
+    // Default for testing
     String getTextEmbeddingsPath() {
         return textEmbeddingsPath;
     }
 
+    @Override
     public CustomServiceEmbeddingType getEmbeddingType() {
         return embeddingType;
     }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomModelTests.java

@@ -110,12 +110,7 @@ public static CustomModel getTestModel(TaskType taskType, CustomResponseParser r
         String requestContentString = "\"input\":\"${input}\"";
 
         CustomServiceSettings serviceSettings = new CustomServiceSettings(
-            new CustomServiceSettings.TextEmbeddingSettings(
-                SimilarityMeasure.DOT_PRODUCT,
-                dims,
-                maxInputTokens,
-                CustomServiceEmbeddingType.FLOAT
-            ),
+            new CustomServiceSettings.TextEmbeddingSettings(SimilarityMeasure.DOT_PRODUCT, dims, maxInputTokens),
             url,
             headers,
             QueryParameters.EMPTY,