Skip to content

Commit bc09ab3

Browse files
davidkyledavidkyle
authored
[ML] EIS: Add request origin to the X-elastic-product-use-case header (#135513)
Adds a value `X-elastic-product-use-case` to indicate if the request came from ingest, search or unspecified (REST API). `X-elastic-product-use-case` is a multi-value header that may also be set by a client.
1 parent ce5ac7b commit bc09ab3

9 files changed

+147
-38
lines changed

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -663,7 +663,7 @@ public void onNodeStarted() {
663663

664664
@Override
665665
public Collection<RestHeaderDefinition> getRestHeaders() {
666-
return Set.of(new RestHeaderDefinition(X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER, false));
666+
return Set.of(new RestHeaderDefinition(X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER, true));
667667
}
668668

669669
@Override

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceUsageContext.java‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,12 @@
77

88
package org.elasticsearch.xpack.inference.services.elastic;
99

10+
import org.elasticsearch.inference.InputType;
11+
1012
import java.util.Locale;
1113

14+
import static org.elasticsearch.inference.InputType.INTERNAL_INGEST;
15+
1216
/**
1317
* Specifies the usage context for a request to the Elastic Inference Service.
1418
* This helps to determine the type of resources that are allocated in the Elastic Inference Service for the particular request.
@@ -24,4 +28,25 @@ public String toString() {
2428
return name().toLowerCase(Locale.ROOT);
2529
}
2630

31+
public static ElasticInferenceServiceUsageContext fromInputType(InputType inputType) {
32+
switch (inputType) {
33+
case SEARCH, INTERNAL_SEARCH -> {
34+
return ElasticInferenceServiceUsageContext.SEARCH;
35+
}
36+
case INGEST, INTERNAL_INGEST -> {
37+
return ElasticInferenceServiceUsageContext.INGEST;
38+
}
39+
default -> {
40+
return ElasticInferenceServiceUsageContext.UNSPECIFIED;
41+
}
42+
}
43+
}
44+
45+
public String productUseCaseHeaderValue() {
46+
return switch (this) {
47+
case SEARCH -> "internal_search";
48+
case INGEST -> "internal_ingest";
49+
case UNSPECIFIED -> "unspecified";
50+
};
51+
}
2752
}

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.elasticsearch.inference.InputType;
1717
import org.elasticsearch.xcontent.XContentType;
1818
import org.elasticsearch.xpack.inference.external.request.Request;
19+
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceUsageContext;
1920
import org.elasticsearch.xpack.inference.services.elastic.densetextembeddings.ElasticInferenceServiceDenseTextEmbeddingsModel;
2021
import org.elasticsearch.xpack.inference.telemetry.TraceContext;
2122
import org.elasticsearch.xpack.inference.telemetry.TraceContextHandler;
@@ -25,7 +26,7 @@
2526
import java.util.List;
2627
import java.util.Objects;
2728

28-
import static org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceSparseEmbeddingsRequest.inputTypeToUsageContext;
29+
import static org.elasticsearch.xpack.inference.InferencePlugin.X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER;
2930

3031
public class ElasticInferenceServiceDenseTextEmbeddingsRequest extends ElasticInferenceServiceRequest {
3132

@@ -53,7 +54,7 @@ public ElasticInferenceServiceDenseTextEmbeddingsRequest(
5354
@Override
5455
public HttpRequestBase createHttpRequestBase() {
5556
var httpPost = new HttpPost(uri);
56-
var usageContext = inputTypeToUsageContext(inputType);
57+
var usageContext = ElasticInferenceServiceUsageContext.fromInputType(inputType);
5758
var requestEntity = Strings.toString(
5859
new ElasticInferenceServiceDenseTextEmbeddingsRequestEntity(inputs, model.getServiceSettings().modelId(), usageContext)
5960
);
@@ -63,6 +64,7 @@ public HttpRequestBase createHttpRequestBase() {
6364

6465
traceContextHandler.propagateTraceContext(httpPost);
6566
httpPost.setHeader(new BasicHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType()));
67+
httpPost.setHeader(new BasicHeader(X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER, usageContext.productUseCaseHeaderValue()));
6668

6769
return httpPost;
6870
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ public final HttpRequest createHttpRequest() {
4242
}
4343

4444
if (Objects.nonNull(productUseCase) && productUseCase.isEmpty() == false) {
45-
request.setHeader(X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER, metadata.productUseCase());
45+
request.addHeader(X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER, metadata.productUseCase());
4646
}
4747

4848
return new HttpRequest(request, getInferenceEntityId());

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

Lines changed: 4 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
import java.nio.charset.StandardCharsets;
2727
import java.util.Objects;
2828

29+
import static org.elasticsearch.xpack.inference.InferencePlugin.X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER;
30+
2931
public class ElasticInferenceServiceSparseEmbeddingsRequest extends ElasticInferenceServiceRequest {
3032

3133
private final URI uri;
@@ -55,7 +57,7 @@ public ElasticInferenceServiceSparseEmbeddingsRequest(
5557
@Override
5658
public HttpRequestBase createHttpRequestBase() {
5759
var httpPost = new HttpPost(uri);
58-
var usageContext = inputTypeToUsageContext(inputType);
60+
var usageContext = ElasticInferenceServiceUsageContext.fromInputType(inputType);
5961
var requestEntity = Strings.toString(
6062
new ElasticInferenceServiceSparseEmbeddingsRequestEntity(
6163
truncationResult.input(),
@@ -69,6 +71,7 @@ public HttpRequestBase createHttpRequestBase() {
6971

7072
traceContextHandler.propagateTraceContext(httpPost);
7173
httpPost.setHeader(new BasicHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType()));
74+
httpPost.setHeader(new BasicHeader(X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER, usageContext.productUseCaseHeaderValue()));
7275

7376
return httpPost;
7477
}
@@ -104,19 +107,4 @@ public Request truncate() {
104107
public boolean[] getTruncationInfo() {
105108
return truncationResult.truncated().clone();
106109
}
107-
108-
// visible for testing
109-
static ElasticInferenceServiceUsageContext inputTypeToUsageContext(InputType inputType) {
110-
switch (inputType) {
111-
case SEARCH, INTERNAL_SEARCH -> {
112-
return ElasticInferenceServiceUsageContext.SEARCH;
113-
}
114-
case INGEST, INTERNAL_INGEST -> {
115-
return ElasticInferenceServiceUsageContext.INGEST;
116-
}
117-
default -> {
118-
return ElasticInferenceServiceUsageContext.UNSPECIFIED;
119-
}
120-
}
121-
}
122110
}

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceTests.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,7 @@
9393
import static org.elasticsearch.xpack.inference.services.ServiceComponentsTests.createWithEmptySettings;
9494
import static org.hamcrest.CoreMatchers.instanceOf;
9595
import static org.hamcrest.CoreMatchers.is;
96+
import static org.hamcrest.Matchers.contains;
9697
import static org.hamcrest.Matchers.equalTo;
9798
import static org.hamcrest.Matchers.hasSize;
9899
import static org.hamcrest.Matchers.isA;
@@ -703,7 +704,8 @@ public void testInfer_PropagatesProductUseCaseHeader() throws IOException {
703704
assertThat(request.getHeader(HttpHeaders.CONTENT_TYPE), Matchers.equalTo(XContentType.JSON.mediaType()));
704705

705706
// Check that the product use case header was set correctly
706-
assertThat(request.getHeader(InferencePlugin.X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER), is(productUseCase));
707+
var productUseCaseHeaders = request.getHeaders().get(InferencePlugin.X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER);
708+
assertThat(productUseCaseHeaders, contains("internal_search", productUseCase));
707709

708710
// Verify request body
709711
var requestMap = entityAsMap(request.getBody());
@@ -832,7 +834,8 @@ public void testChunkedInfer_PropagatesProductUseCaseHeader() throws IOException
832834
assertThat(request.getHeader(HttpHeaders.CONTENT_TYPE), equalTo(XContentType.JSON.mediaType()));
833835

834836
// Check that the product use case header was set correctly
835-
assertThat(request.getHeader(InferencePlugin.X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER), is(productUseCase));
837+
var productUseCaseHeaders = request.getHeaders().get(InferencePlugin.X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER);
838+
assertThat(productUseCaseHeaders, contains("internal_ingest", productUseCase));
836839

837840
} finally {
838841
// Clean up the thread context

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceUsageContextTests.java‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.inference.services.elastic;
9+
10+
import org.elasticsearch.inference.InputType;
11+
import org.elasticsearch.test.ESTestCase;
12+
import org.hamcrest.Matchers;
13+
14+
import static org.hamcrest.Matchers.equalTo;
15+
16+
public class ElasticInferenceServiceUsageContextTests extends ESTestCase {
17+
18+
public void testInputTypeToUsageContext_Search() {
19+
assertThat(
20+
ElasticInferenceServiceUsageContext.fromInputType(InputType.SEARCH),
21+
equalTo(ElasticInferenceServiceUsageContext.SEARCH)
22+
);
23+
}
24+
25+
public void testInputTypeToUsageContext_Ingest() {
26+
assertThat(
27+
ElasticInferenceServiceUsageContext.fromInputType(InputType.INGEST),
28+
equalTo(ElasticInferenceServiceUsageContext.INGEST)
29+
);
30+
}
31+
32+
public void testInputTypeToUsageContext_Unspecified() {
33+
assertThat(
34+
ElasticInferenceServiceUsageContext.fromInputType(InputType.UNSPECIFIED),
35+
equalTo(ElasticInferenceServiceUsageContext.UNSPECIFIED)
36+
);
37+
}
38+
39+
public void testInputTypeToUsageContext_Unknown_DefaultToUnspecified() {
40+
assertThat(
41+
ElasticInferenceServiceUsageContext.fromInputType(InputType.CLASSIFICATION),
42+
equalTo(ElasticInferenceServiceUsageContext.UNSPECIFIED)
43+
);
44+
assertThat(
45+
ElasticInferenceServiceUsageContext.fromInputType(InputType.CLUSTERING),
46+
equalTo(ElasticInferenceServiceUsageContext.UNSPECIFIED)
47+
);
48+
}
49+
50+
public void testProductUseCase() {
51+
assertThat(ElasticInferenceServiceUsageContext.SEARCH.productUseCaseHeaderValue(), Matchers.is("internal_search"));
52+
assertThat(ElasticInferenceServiceUsageContext.INGEST.productUseCaseHeaderValue(), Matchers.is("internal_ingest"));
53+
assertThat(ElasticInferenceServiceUsageContext.UNSPECIFIED.productUseCaseHeaderValue(), Matchers.is("unspecified"));
54+
}
55+
}

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.io.IOException;
2020
import java.util.List;
2121

22+
import static org.elasticsearch.xpack.inference.InferencePlugin.X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER;
2223
import static org.elasticsearch.xpack.inference.external.http.Utils.entityAsMap;
2324
import static org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceRequestTests.randomElasticInferenceServiceRequestMetadata;
2425
import static org.hamcrest.Matchers.aMapWithSize;
@@ -146,6 +147,32 @@ public void testGetTruncationInfo_ReturnsNull() {
146147
assertThat(request.getTruncationInfo(), is(nullValue()));
147148
}
148149

150+
public void testDecorate_HttpRequest_WithProductUseCase() {
151+
var input = "elastic";
152+
var modelId = "my-model-id";
153+
var url = "http://eis-gateway.com";
154+
155+
for (var inputType : List.of(InputType.INTERNAL_SEARCH, InputType.INTERNAL_INGEST, InputType.UNSPECIFIED)) {
156+
var request = new ElasticInferenceServiceDenseTextEmbeddingsRequest(
157+
ElasticInferenceServiceDenseTextEmbeddingsModelTests.createModel(url, modelId),
158+
List.of(input),
159+
new TraceContext(randomAlphaOfLength(10), randomAlphaOfLength(10)),
160+
new ElasticInferenceServiceRequestMetadata("my-product-origin", "my-product-use-case-from-metadata"),
161+
inputType
162+
);
163+
164+
var httpRequest = request.createHttpRequest();
165+
166+
assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class));
167+
var httpPost = (HttpPost) httpRequest.httpRequestBase();
168+
169+
var headers = httpPost.getHeaders(X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER);
170+
assertThat(headers.length, is(2));
171+
assertThat(headers[0].getValue(), is(inputType.toString()));
172+
assertThat(headers[1].getValue(), is("my-product-use-case-from-metadata"));
173+
}
174+
}
175+
149176
private ElasticInferenceServiceDenseTextEmbeddingsRequest createRequest(
150177
String url,
151178
String modelId,

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

Lines changed: 25 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,15 +16,14 @@
1616
import org.elasticsearch.xpack.inference.common.Truncator;
1717
import org.elasticsearch.xpack.inference.common.TruncatorTests;
1818
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSparseEmbeddingsModelTests;
19-
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceUsageContext;
2019
import org.elasticsearch.xpack.inference.telemetry.TraceContext;
2120

2221
import java.io.IOException;
2322
import java.util.List;
2423

24+
import static org.elasticsearch.xpack.inference.InferencePlugin.X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER;
2525
import static org.elasticsearch.xpack.inference.external.http.Utils.entityAsMap;
2626
import static org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceRequestTests.randomElasticInferenceServiceRequestMetadata;
27-
import static org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceSparseEmbeddingsRequest.inputTypeToUsageContext;
2827
import static org.hamcrest.Matchers.aMapWithSize;
2928
import static org.hamcrest.Matchers.equalTo;
3029
import static org.hamcrest.Matchers.instanceOf;
@@ -99,21 +98,31 @@ public void testIsTruncated_ReturnsTrue() {
9998
assertTrue(truncatedRequest.getTruncationInfo()[0]);
10099
}
101100

102-
public void testInputTypeToUsageContext_Search() {
103-
assertThat(inputTypeToUsageContext(InputType.SEARCH), equalTo(ElasticInferenceServiceUsageContext.SEARCH));
104-
}
105-
106-
public void testInputTypeToUsageContext_Ingest() {
107-
assertThat(inputTypeToUsageContext(InputType.INGEST), equalTo(ElasticInferenceServiceUsageContext.INGEST));
108-
}
109-
110-
public void testInputTypeToUsageContext_Unspecified() {
111-
assertThat(inputTypeToUsageContext(InputType.UNSPECIFIED), equalTo(ElasticInferenceServiceUsageContext.UNSPECIFIED));
112-
}
101+
public void testDecorate_HttpRequest_WithProductUseCase() {
102+
var input = "elastic";
103+
var modelId = "my-model-id";
104+
var url = "http://eis-gateway.com";
113105

114-
public void testInputTypeToUsageContext_Unknown_DefaultToUnspecified() {
115-
assertThat(inputTypeToUsageContext(InputType.CLASSIFICATION), equalTo(ElasticInferenceServiceUsageContext.UNSPECIFIED));
116-
assertThat(inputTypeToUsageContext(InputType.CLUSTERING), equalTo(ElasticInferenceServiceUsageContext.UNSPECIFIED));
106+
for (var inputType : List.of(InputType.INTERNAL_SEARCH, InputType.INTERNAL_INGEST, InputType.UNSPECIFIED)) {
107+
var request = new ElasticInferenceServiceSparseEmbeddingsRequest(
108+
TruncatorTests.createTruncator(),
109+
new Truncator.TruncationResult(List.of(input), new boolean[] { false }),
110+
ElasticInferenceServiceSparseEmbeddingsModelTests.createModel(url, modelId),
111+
new TraceContext(randomAlphaOfLength(10), randomAlphaOfLength(10)),
112+
new ElasticInferenceServiceRequestMetadata("my-product-origin", "my-product-use-case-from-metadata"),
113+
inputType
114+
);
115+
116+
var httpRequest = request.createHttpRequest();
117+
118+
assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class));
119+
var httpPost = (HttpPost) httpRequest.httpRequestBase();
120+
121+
var headers = httpPost.getHeaders(X_ELASTIC_PRODUCT_USE_CASE_HTTP_HEADER);
122+
assertThat(headers.length, is(2));
123+
assertThat(headers[0].getValue(), is(inputType.toString()));
124+
assertThat(headers[1].getValue(), is("my-product-use-case-from-metadata"));
125+
}
117126
}
118127

119128
public ElasticInferenceServiceSparseEmbeddingsRequest createRequest(String url, String modelId, String input, InputType inputType) {

0 commit comments

Comments
 (0)