[Inference API] Propagate dimensions for dense text embedding generation using EIS (elastic#137328) (elastic#137364)

Author: timgrein

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/request/ElasticInferenceServiceDenseTextEmbeddingsRequest.java

@@ -56,7 +56,12 @@ public HttpRequestBase createHttpRequestBase() {
         var httpPost = new HttpPost(uri);
         var usageContext = ElasticInferenceServiceUsageContext.fromInputType(inputType);
         var requestEntity = Strings.toString(
-            new ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(inputs, model.getServiceSettings().modelId(), usageContext)
+            new ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(
+                inputs,
+                model.getServiceSettings().modelId(),
+                usageContext,
+                model.getServiceSettings().dimensions()
+            )
         );
 
         ByteArrayEntity byteEntity = new ByteArrayEntity(requestEntity.getBytes(StandardCharsets.UTF_8));

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/request/ElasticInferenceServiceDenseTextEmbeddingsRequestEntity.java

@@ -19,12 +19,14 @@
 public record ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(
     List<String> inputs,
     String modelId,
-    @Nullable ElasticInferenceServiceUsageContext usageContext
+    @Nullable ElasticInferenceServiceUsageContext usageContext,
+    @Nullable Integer dimensions
 ) implements ToXContentObject {
 
     private static final String INPUT_FIELD = "input";
     private static final String MODEL_FIELD = "model";
     private static final String USAGE_CONTEXT = "usage_context";
+    private static final String DIMENSIONS = "dimensions";
 
     public ElasticInferenceServiceDenseTextEmbeddingsRequestEntity {
         Objects.requireNonNull(inputs);
@@ -49,6 +51,11 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field(USAGE_CONTEXT, usageContext);
         }
 
+        // optional field
+        if (Objects.nonNull(dimensions)) {
+            builder.field(DIMENSIONS, dimensions);
+        }
+
         builder.endObject();
 
         return builder;

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/request/ElasticInferenceServiceDenseTextEmbeddingsRequestEntityTests.java

@@ -25,7 +25,8 @@ public void testToXContent_SingleInput_UnspecifiedUsageContext() throws IOExcept
         var entity = new ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(
             List.of("abc"),
             "my-model-id",
-            ElasticInferenceServiceUsageContext.UNSPECIFIED
+            ElasticInferenceServiceUsageContext.UNSPECIFIED,
+            null
         );
         String xContentString = xContentEntityToString(entity);
         assertThat(xContentString, equalToIgnoringWhitespaceInJsonString("""
@@ -39,7 +40,8 @@ public void testToXContent_MultipleInputs_UnspecifiedUsageContext() throws IOExc
         var entity = new ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(
             List.of("abc", "def"),
             "my-model-id",
-            ElasticInferenceServiceUsageContext.UNSPECIFIED
+            ElasticInferenceServiceUsageContext.UNSPECIFIED,
+            null
         );
         String xContentString = xContentEntityToString(entity);
         assertThat(xContentString, equalToIgnoringWhitespaceInJsonString("""
@@ -57,7 +59,8 @@ public void testToXContent_SingleInput_SearchUsageContext() throws IOException {
         var entity = new ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(
             List.of("abc"),
             "my-model-id",
-            ElasticInferenceServiceUsageContext.SEARCH
+            ElasticInferenceServiceUsageContext.SEARCH,
+            null
         );
         String xContentString = xContentEntityToString(entity);
         assertThat(xContentString, equalToIgnoringWhitespaceInJsonString("""
@@ -73,7 +76,8 @@ public void testToXContent_SingleInput_IngestUsageContext() throws IOException {
         var entity = new ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(
             List.of("abc"),
             "my-model-id",
-            ElasticInferenceServiceUsageContext.INGEST
+            ElasticInferenceServiceUsageContext.INGEST,
+            null
         );
         String xContentString = xContentEntityToString(entity);
         assertThat(xContentString, equalToIgnoringWhitespaceInJsonString("""
@@ -85,11 +89,29 @@ public void testToXContent_SingleInput_IngestUsageContext() throws IOException {
             """));
     }
 
+    public void testToXContent_SingleInput_DimensionsSpecified() throws IOException {
+        var entity = new ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(
+            List.of("abc"),
+            "my-model-id",
+            ElasticInferenceServiceUsageContext.UNSPECIFIED,
+            100
+        );
+        String xContentString = xContentEntityToString(entity);
+        assertThat(xContentString, equalToIgnoringWhitespaceInJsonString("""
+            {
+                "input": ["abc"],
+                "model": "my-model-id",
+                "dimensions": 100
+            }
+            """));
+    }
+
     public void testToXContent_MultipleInputs_SearchUsageContext() throws IOException {
         var entity = new ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(
             List.of("first input", "second input", "third input"),
             "my-dense-model",
-            ElasticInferenceServiceUsageContext.SEARCH
+            ElasticInferenceServiceUsageContext.SEARCH,
+            null
         );
         String xContentString = xContentEntityToString(entity);
         assertThat(xContentString, equalToIgnoringWhitespaceInJsonString("""
@@ -109,7 +131,8 @@ public void testToXContent_MultipleInputs_IngestUsageContext() throws IOExceptio
         var entity = new ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(
             List.of("document one", "document two"),
             "embedding-model-v2",
-            ElasticInferenceServiceUsageContext.INGEST
+            ElasticInferenceServiceUsageContext.INGEST,
+            null
         );
         String xContentString = xContentEntityToString(entity);
         assertThat(xContentString, equalToIgnoringWhitespaceInJsonString("""
@@ -128,7 +151,8 @@ public void testToXContent_EmptyInput_UnspecifiedUsageContext() throws IOExcepti
         var entity = new ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(
             List.of(""),
             "my-model-id",
-            ElasticInferenceServiceUsageContext.UNSPECIFIED
+            ElasticInferenceServiceUsageContext.UNSPECIFIED,
+            null
         );
         String xContentString = xContentEntityToString(entity);
         assertThat(xContentString, equalToIgnoringWhitespaceInJsonString("""