voyage ai

Author: ymao1

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/voyageai/VoyageAIEmbeddingsRequestEntity.java

@@ -72,8 +72,8 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
     static String convertToString(InputType inputType) {
         return switch (inputType) {
             case null -> null;
-            case INGEST -> DOCUMENT;
-            case SEARCH -> QUERY;
+            case INGEST, INTERNAL_INGEST -> DOCUMENT;
+            case SEARCH, INTERNAL_SEARCH -> QUERY;
             default -> {
                 assert false : invalidInputTypeMessage(inputType);
                 yield null;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/voyageai/embeddings/VoyageAIEmbeddingsTaskSettings.java

@@ -41,7 +41,12 @@ public class VoyageAIEmbeddingsTaskSettings implements TaskSettings {
     public static final String NAME = "voyageai_embeddings_task_settings";
     public static final VoyageAIEmbeddingsTaskSettings EMPTY_SETTINGS = new VoyageAIEmbeddingsTaskSettings(null, null);
     static final String INPUT_TYPE = "input_type";
-    static final EnumSet<InputType> VALID_REQUEST_VALUES = EnumSet.of(InputType.INGEST, InputType.SEARCH);
+    static final EnumSet<InputType> VALID_REQUEST_VALUES = EnumSet.of(
+        InputType.INGEST,
+        InputType.SEARCH,
+        InputType.INTERNAL_INGEST,
+        InputType.INTERNAL_SEARCH
+    );
 
     public static VoyageAIEmbeddingsTaskSettings fromMap(Map<String, Object> map) {
         if (map == null || map.isEmpty()) {

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/voyageai/VoyageAIEmbeddingsRequestEntityTests.java

@@ -163,6 +163,38 @@ public void testXContent_WritesNoOptionalFields_WhenTheyAreNotDefined() throws I
             {"input":["abc"],"model":"model"}"""));
     }
 
+    public void testXContent_InputTypeInternalSearch() throws IOException {
+        var entity = new VoyageAIEmbeddingsRequestEntity(
+            List.of("abc"),
+            VoyageAIEmbeddingsServiceSettings.EMPTY_SETTINGS,
+            new VoyageAIEmbeddingsTaskSettings(InputType.INTERNAL_SEARCH, Boolean.FALSE),
+            "model"
+        );
+
+        XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
+        entity.toXContent(builder, null);
+        String xContentResult = Strings.toString(builder);
+
+        MatcherAssert.assertThat(xContentResult, is("""
+            {"input":["abc"],"model":"model","input_type":"query"}"""));
+    }
+
+    public void testXContent_InputTypeInternalIngest() throws IOException {
+        var entity = new VoyageAIEmbeddingsRequestEntity(
+            List.of("abc"),
+            VoyageAIEmbeddingsServiceSettings.EMPTY_SETTINGS,
+            new VoyageAIEmbeddingsTaskSettings(InputType.INTERNAL_INGEST, Boolean.FALSE),
+            "model"
+        );
+
+        XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
+        entity.toXContent(builder, null);
+        String xContentResult = Strings.toString(builder);
+
+        MatcherAssert.assertThat(xContentResult, is("""
+            {"input":["abc"],"model":"model","input_type":"document"}"""));
+    }
+
     public void testConvertToString_ThrowsAssertionFailure_WhenInputTypeIsUnspecified() {
         var thrownException = expectThrows(
             AssertionError.class,

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/voyageai/embeddings/VoyageAIEmbeddingsTaskSettingsTests.java

@@ -181,10 +181,10 @@ public void testOf_UsesRequestTaskSettings_AndRequestInputType() {
         var overriddenTaskSettings = VoyageAIEmbeddingsTaskSettings.of(
             taskSettings,
             new VoyageAIEmbeddingsTaskSettings((InputType) null, null),
-            InputType.INGEST
+            InputType.INTERNAL_INGEST
         );
 
-        MatcherAssert.assertThat(overriddenTaskSettings, is(new VoyageAIEmbeddingsTaskSettings(InputType.INGEST, true)));
+        MatcherAssert.assertThat(overriddenTaskSettings, is(new VoyageAIEmbeddingsTaskSettings(InputType.INTERNAL_INGEST, true)));
     }
 
     @Override