[Inference API] Use dimensions field in JinaAI text_embedding requests (#139552)

The dimensions field was parsed from the service settings but never
used. This commit includes the dimensions field in text embedding
requests sent to JinaAI when set by the user.

* Add dimensions_set_by_user to JinaAIEmbeddingsServiceSettings
* Move similarity and max_input_tokens to exposed fields

(cherry picked from commit b2ecf87)

Author: DonalEvans

docs/changelog/139413.yaml

@@ -0,0 +1,5 @@
+pr: 139413
+summary: "[Inference API] Use dimensions field in JinaAI `text_embedding` requests"
+area: Inference
+type: bug
+issues: []

server/src/main/resources/transport/definitions/referable/jina_ai_embedding_dimensions_support_added.csv

@@ -0,0 +1 @@
+9244000,9185014,9112017,8841078

server/src/main/resources/transport/upper_bounds/8.19.csv

@@ -1 +1 @@
-downsample_add_multi_field_sources,8841077
+jina_ai_embedding_dimensions_support_added,8841078

server/src/main/resources/transport/upper_bounds/9.1.csv

@@ -1 +1 @@
-downsample_add_multi_field_sources,9112016
+jina_ai_embedding_dimensions_support_added,9112017

server/src/main/resources/transport/upper_bounds/9.2.csv

@@ -1 +1 @@
-downsample_add_multi_field_sources,9185013
+jina_ai_embedding_dimensions_support_added,9185014

server/src/main/resources/transport/upper_bounds/9.3.csv

@@ -1 +1 @@
-downsample_add_multi_field_sources,9232000
+jina_ai_embedding_dimensions_support_added,9244000

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

@@ -14,6 +14,7 @@ public final class ServiceFields {
 
     public static final String SIMILARITY = "similarity";
     public static final String DIMENSIONS = "dimensions";
+    public static final String DIMENSIONS_SET_BY_USER = "dimensions_set_by_user";
     // Typically we use this to define the maximum tokens for the input text (text being sent to an integration)
     public static final String MAX_INPUT_TOKENS = "max_input_tokens";
     public static final String URL = "url";

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

@@ -308,7 +308,8 @@ public Model updateModelWithEmbeddingDetails(Model model, int embeddingSize) {
                 similarityToUse,
                 embeddingSize,
                 maxInputTokens,
-                serviceSettings.getEmbeddingType()
+                serviceSettings.getEmbeddingType(),
+                serviceSettings.dimensionsSetByUser()
             );
 
             return new JinaAIEmbeddingsModel(embeddingsModel, updatedServiceSettings);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/jinaai/embeddings/JinaAIEmbeddingsServiceSettings.java

@@ -27,9 +27,11 @@
 import java.util.Objects;
 
 import static org.elasticsearch.xpack.inference.services.ServiceFields.DIMENSIONS;
+import static org.elasticsearch.xpack.inference.services.ServiceFields.DIMENSIONS_SET_BY_USER;
 import static org.elasticsearch.xpack.inference.services.ServiceFields.MAX_INPUT_TOKENS;
 import static org.elasticsearch.xpack.inference.services.ServiceFields.SIMILARITY;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalEnum;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalPositiveInteger;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractSimilarity;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeAsType;
 
@@ -42,16 +44,33 @@ public static JinaAIEmbeddingsServiceSettings fromMap(Map<String, Object> map, C
         ValidationException validationException = new ValidationException();
         var commonServiceSettings = JinaAIServiceSettings.fromMap(map, context);
         SimilarityMeasure similarity = extractSimilarity(map, ModelConfigurations.SERVICE_SETTINGS, validationException);
-        Integer dims = removeAsType(map, DIMENSIONS, Integer.class);
+        Integer dimensions = extractOptionalPositiveInteger(map, DIMENSIONS, ModelConfigurations.SERVICE_SETTINGS, validationException);
         Integer maxInputTokens = removeAsType(map, MAX_INPUT_TOKENS, Integer.class);
 
         JinaAIEmbeddingType embeddingTypes = parseEmbeddingType(map, validationException);
 
+        Boolean dimensionsSetByUser;
+        if (context == ConfigurationParseContext.PERSISTENT) {
+            dimensionsSetByUser = removeAsType(map, DIMENSIONS_SET_BY_USER, Boolean.class);
+            if (dimensionsSetByUser == null) {
+                dimensionsSetByUser = Boolean.FALSE;
+            }
+        } else {
+            dimensionsSetByUser = dimensions != null;
+        }
+
         if (validationException.validationErrors().isEmpty() == false) {
             throw validationException;
         }
 
-        return new JinaAIEmbeddingsServiceSettings(commonServiceSettings, similarity, dims, maxInputTokens, embeddingTypes);
+        return new JinaAIEmbeddingsServiceSettings(
+            commonServiceSettings,
+            similarity,
+            dimensions,
+            maxInputTokens,
+            embeddingTypes,
+            dimensionsSetByUser
+        );
     }
 
     static JinaAIEmbeddingType parseEmbeddingType(Map<String, Object> map, ValidationException validationException) {
@@ -72,24 +91,31 @@ static JinaAIEmbeddingType parseEmbeddingType(Map<String, Object> map, Validatio
         "jina_ai_embedding_type_support_added"
     );
 
+    static final TransportVersion JINA_AI_EMBEDDING_DIMENSIONS_SUPPORT_ADDED = TransportVersion.fromName(
+        "jina_ai_embedding_dimensions_support_added"
+    );
+
     private final JinaAIServiceSettings commonSettings;
     private final SimilarityMeasure similarity;
     private final Integer dimensions;
     private final Integer maxInputTokens;
     private final JinaAIEmbeddingType embeddingType;
+    private final Boolean dimensionsSetByUser;
 
     public JinaAIEmbeddingsServiceSettings(
         JinaAIServiceSettings commonSettings,
         @Nullable SimilarityMeasure similarity,
         @Nullable Integer dimensions,
         @Nullable Integer maxInputTokens,
-        @Nullable JinaAIEmbeddingType embeddingType
+        @Nullable JinaAIEmbeddingType embeddingType,
+        Boolean dimensionsSetByUser
     ) {
         this.commonSettings = commonSettings;
         this.similarity = similarity;
         this.dimensions = dimensions;
         this.maxInputTokens = maxInputTokens;
         this.embeddingType = embeddingType != null ? embeddingType : JinaAIEmbeddingType.FLOAT;
+        this.dimensionsSetByUser = Objects.requireNonNull(dimensionsSetByUser);
     }
 
     public JinaAIEmbeddingsServiceSettings(StreamInput in) throws IOException {
@@ -101,6 +127,12 @@ public JinaAIEmbeddingsServiceSettings(StreamInput in) throws IOException {
         this.embeddingType = (in.getTransportVersion().supports(JINA_AI_EMBEDDING_TYPE_SUPPORT_ADDED))
             ? Objects.requireNonNullElse(in.readOptionalEnum(JinaAIEmbeddingType.class), JinaAIEmbeddingType.FLOAT)
             : JinaAIEmbeddingType.FLOAT;
+
+        if (in.getTransportVersion().supports(JINA_AI_EMBEDDING_DIMENSIONS_SUPPORT_ADDED)) {
+            this.dimensionsSetByUser = in.readBoolean();
+        } else {
+            this.dimensionsSetByUser = false;
+        }
     }
 
     public JinaAIServiceSettings getCommonSettings() {
@@ -117,6 +149,11 @@ public Integer dimensions() {
         return dimensions;
     }
 
+    @Override
+    public Boolean dimensionsSetByUser() {
+        return dimensionsSetByUser;
+    }
+
     public Integer maxInputTokens() {
         return maxInputTokens;
     }
@@ -144,18 +181,10 @@ public String getWriteableName() {
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
 
-        builder = commonSettings.toXContentFragment(builder, params);
-        if (similarity != null) {
-            builder.field(SIMILARITY, similarity);
-        }
-        if (dimensions != null) {
-            builder.field(DIMENSIONS, dimensions);
-        }
-        if (maxInputTokens != null) {
-            builder.field(MAX_INPUT_TOKENS, maxInputTokens);
-        }
-        if (embeddingType != null) {
-            builder.field(EMBEDDING_TYPE, embeddingType);
+        toXContentFragmentOfExposedFields(builder, params);
+
+        if (dimensionsSetByUser != null) {
+            builder.field(DIMENSIONS_SET_BY_USER, dimensionsSetByUser);
         }
 
         builder.endObject();
@@ -165,9 +194,18 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
     @Override
     protected XContentBuilder toXContentFragmentOfExposedFields(XContentBuilder builder, Params params) throws IOException {
         commonSettings.toXContentFragmentOfExposedFields(builder, params);
+        if (dimensions != null) {
+            builder.field(DIMENSIONS, dimensions);
+        }
         if (embeddingType != null) {
             builder.field(EMBEDDING_TYPE, embeddingType);
         }
+        if (maxInputTokens != null) {
+            builder.field(MAX_INPUT_TOKENS, maxInputTokens);
+        }
+        if (similarity != null) {
+            builder.field(SIMILARITY, similarity);
+        }
         return builder;
     }
 
@@ -186,6 +224,10 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().supports(JINA_AI_EMBEDDING_TYPE_SUPPORT_ADDED)) {
             out.writeOptionalEnum(JinaAIEmbeddingType.translateToVersion(embeddingType, out.getTransportVersion()));
         }
+
+        if (out.getTransportVersion().supports(JINA_AI_EMBEDDING_DIMENSIONS_SUPPORT_ADDED)) {
+            out.writeBoolean(dimensionsSetByUser);
+        }
     }
 
     @Override
@@ -197,11 +239,12 @@ public boolean equals(Object o) {
             && Objects.equals(similarity, that.similarity)
             && Objects.equals(dimensions, that.dimensions)
             && Objects.equals(maxInputTokens, that.maxInputTokens)
-            && Objects.equals(embeddingType, that.embeddingType);
+            && Objects.equals(embeddingType, that.embeddingType)
+            && Objects.equals(dimensionsSetByUser, that.dimensionsSetByUser);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(commonSettings, similarity, dimensions, maxInputTokens, embeddingType);
+        return Objects.hash(commonSettings, similarity, dimensions, maxInputTokens, embeddingType, dimensionsSetByUser);
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/jinaai/request/JinaAIEmbeddingsRequest.java

@@ -34,6 +34,8 @@ public class JinaAIEmbeddingsRequest extends JinaAIRequest {
     private final String model;
     private final String inferenceEntityId;
     private final JinaAIEmbeddingType embeddingType;
+    private final Integer dimensions;
+    private final boolean dimensionsSetByUser;
 
     public JinaAIEmbeddingsRequest(List<String> input, InputType inputType, JinaAIEmbeddingsModel embeddingsModel) {
         Objects.requireNonNull(embeddingsModel);
@@ -45,15 +47,18 @@ public JinaAIEmbeddingsRequest(List<String> input, InputType inputType, JinaAIEm
         model = embeddingsModel.getServiceSettings().getCommonSettings().modelId();
         embeddingType = embeddingsModel.getServiceSettings().getEmbeddingType();
         inferenceEntityId = embeddingsModel.getInferenceEntityId();
+        dimensions = embeddingsModel.getServiceSettings().dimensions();
+        dimensionsSetByUser = embeddingsModel.getServiceSettings().dimensionsSetByUser();
     }
 
     @Override
     public HttpRequest createHttpRequest() {
         HttpPost httpPost = new HttpPost(account.uri());
 
         ByteArrayEntity byteEntity = new ByteArrayEntity(
-            Strings.toString(new JinaAIEmbeddingsRequestEntity(input, inputType, taskSettings, model, embeddingType))
-                .getBytes(StandardCharsets.UTF_8)
+            Strings.toString(
+                new JinaAIEmbeddingsRequestEntity(input, inputType, taskSettings, model, embeddingType, dimensions, dimensionsSetByUser)
+            ).getBytes(StandardCharsets.UTF_8)
         );
         httpPost.setEntity(byteEntity);