Skip to content

Commit ea4ad64

Browse files
ymao1ymao1
committed
external request tests
1 parent f1bbcc6 commit ea4ad64

File tree

35 files changed

+852
-611
lines changed

35 files changed

+852
-611
lines changed

‎server/src/main/java/org/elasticsearch/inference/InputType.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,4 +47,8 @@ public static InputType fromRestString(String name) {
4747
public static boolean isInternalType(InputType inputType) {
4848
return inputType == InputType.INTERNAL_INGEST || inputType == InputType.INTERNAL_SEARCH;
4949
}
50+
51+
public static boolean isSpecified(InputType inputType) {
52+
return inputType != null && inputType != InputType.UNSPECIFIED;
53+
}
5054
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/alibabacloudsearch/AlibabaCloudSearchEmbeddingsRequestEntity.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
4242
builder.field(TEXTS_FIELD, input);
4343

4444
// prefer the root level inputType over task settings input type
45-
if (inputType != null) {
45+
if (InputType.isSpecified(inputType)) {
4646
builder.field(INPUT_TYPE_FIELD, convertToString(inputType));
47-
} else if (taskSettings.getInputType() != null) {
47+
} else if (InputType.isSpecified(taskSettings.getInputType())) {
4848
builder.field(INPUT_TYPE_FIELD, convertToString(taskSettings.getInputType()));
4949
}
5050

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/alibabacloudsearch/AlibabaCloudSearchSparseRequestEntity.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
3939
builder.field(TEXTS_FIELD, input);
4040

4141
// prefer the root level inputType over task settings input type
42-
if (inputType != null) {
42+
if (InputType.isSpecified(inputType)) {
4343
builder.field(INPUT_TYPE_FIELD, AlibabaCloudSearchEmbeddingsRequestEntity.convertToString(inputType));
44-
} else if (taskSettings.getInputType() != null) {
44+
} else if (InputType.isSpecified(taskSettings.getInputType())) {
4545
builder.field(INPUT_TYPE_FIELD, AlibabaCloudSearchEmbeddingsRequestEntity.convertToString(taskSettings.getInputType()));
4646
}
4747

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/amazonbedrock/embeddings/AmazonBedrockCohereEmbeddingsRequestEntity.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
3636
builder.startObject();
3737
builder.field(TEXTS_FIELD, input);
3838

39-
if (inputType != null) {
39+
if (InputType.isSpecified(inputType)) {
4040
builder.field(INPUT_TYPE_FIELD, convertToString(inputType));
4141
} else {
4242
// input_type is required so default to document

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/azureaistudio/AzureAiStudioEmbeddingsRequestEntity.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import static org.elasticsearch.xpack.inference.services.azureaistudio.AzureAiStudioConstants.DIMENSIONS_FIELD;
2020
import static org.elasticsearch.xpack.inference.services.azureaistudio.AzureAiStudioConstants.INPUT_FIELD;
2121
import static org.elasticsearch.xpack.inference.services.azureaistudio.AzureAiStudioConstants.USER_FIELD;
22+
import static org.elasticsearch.xpack.inference.services.cohere.embeddings.CohereEmbeddingsTaskSettings.invalidInputTypeMessage;
2223

2324
public record AzureAiStudioEmbeddingsRequestEntity(
2425
List<String> input,
@@ -28,6 +29,10 @@ public record AzureAiStudioEmbeddingsRequestEntity(
2829
boolean dimensionsSetByUser
2930
) implements ToXContentObject {
3031

32+
private static final String DOCUMENT = "document";
33+
private static final String QUERY = "query";
34+
public static final String INPUT_TYPE_FIELD = "input_type";
35+
3136
public AzureAiStudioEmbeddingsRequestEntity {
3237
Objects.requireNonNull(input);
3338
}
@@ -46,8 +51,24 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
4651
builder.field(DIMENSIONS_FIELD, dimensions);
4752
}
4853

54+
if (InputType.isSpecified(inputType)) {
55+
builder.field(INPUT_TYPE_FIELD, convertToString(inputType));
56+
}
57+
4958
builder.endObject();
5059

5160
return builder;
5261
}
62+
63+
// default for testing
64+
public static String convertToString(InputType inputType) {
65+
return switch (inputType) {
66+
case INGEST, INTERNAL_INGEST -> DOCUMENT;
67+
case SEARCH, INTERNAL_SEARCH -> QUERY;
68+
default -> {
69+
assert false : invalidInputTypeMessage(inputType);
70+
yield null;
71+
}
72+
};
73+
}
5374
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/cohere/CohereEmbeddingsRequestEntity.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,9 +52,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
5252
}
5353

5454
// prefer the root level inputType over task settings input type
55-
if (inputType != null) {
55+
if (InputType.isSpecified(inputType)) {
5656
builder.field(INPUT_TYPE_FIELD, convertToString(inputType));
57-
} else if (taskSettings.getInputType() != null) {
57+
} else if (InputType.isSpecified(taskSettings.getInputType())) {
5858
builder.field(INPUT_TYPE_FIELD, convertToString(taskSettings.getInputType()));
5959
}
6060

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googleaistudio/GoogleAiStudioEmbeddingsRequestEntity.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
7575
builder.field(OUTPUT_DIMENSIONALITY_FIELD, dimensions);
7676
}
7777

78-
if (inputType != null) {
78+
if (InputType.isSpecified(inputType)) {
7979
builder.field(TASK_TYPE_FIELD, convertToString(inputType));
8080
}
8181

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestEntity.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
5151
builder.field(CONTENT_FIELD, input);
5252

5353
// prefer the root level inputType over task settings input type
54-
if (inputType != null) {
54+
if (InputType.isSpecified(inputType)) {
5555
builder.field(TASK_TYPE_FIELD, convertToString(inputType));
56-
} else if (taskSettings.getInputType() != null) {
56+
} else if (InputType.isSpecified(taskSettings.getInputType())) {
5757
builder.field(TASK_TYPE_FIELD, convertToString(taskSettings.getInputType()));
5858
}
5959
}
@@ -74,7 +74,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
7474
return builder;
7575
}
7676

77-
static String convertToString(InputType inputType) {
77+
public static String convertToString(InputType inputType) {
7878
return switch (inputType) {
7979
case INGEST, INTERNAL_INGEST -> RETRIEVAL_DOCUMENT_TASK_TYPE;
8080
case SEARCH, INTERNAL_SEARCH -> RETRIEVAL_QUERY_TASK_TYPE;

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/jinaai/JinaAIEmbeddingsRequestEntity.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,9 +54,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
5454
}
5555

5656
// prefer the root level inputType over task settings input type
57-
if (inputType != null) {
57+
if (InputType.isSpecified(inputType)) {
5858
builder.field(TASK_TYPE_FIELD, convertToString(inputType));
59-
} else if (taskSettings.getInputType() != null) {
59+
} else if (InputType.isSpecified(taskSettings.getInputType())) {
6060
builder.field(TASK_TYPE_FIELD, convertToString(taskSettings.getInputType()));
6161
}
6262

@@ -65,7 +65,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
6565
}
6666

6767
// default for testing
68-
static String convertToString(InputType inputType) {
68+
public static String convertToString(InputType inputType) {
6969
return switch (inputType) {
7070
case INGEST, INTERNAL_INGEST -> SEARCH_DOCUMENT;
7171
case SEARCH, INTERNAL_SEARCH -> SEARCH_QUERY;

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
5050
builder.field(MODEL_FIELD, model);
5151

5252
// prefer the root level inputType over task settings input type
53-
if (inputType != null) {
53+
if (InputType.isSpecified(inputType)) {
5454
builder.field(INPUT_TYPE_FIELD, convertToString(inputType));
55-
} else if (taskSettings.getInputType() != null) {
55+
} else if (InputType.isSpecified(taskSettings.getInputType())) {
5656
builder.field(INPUT_TYPE_FIELD, convertToString(taskSettings.getInputType()));
5757
}
5858

0 commit comments

Comments
 (0)