Skip to content

Commit 104ab7c

Browse files
author
Max Hniebergall
committed
Refactor EndpointVersion into enum instead of string
1 parent 388c475 commit 104ab7c

File tree

111 files changed

+765
-667
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

111 files changed

+765
-667
lines changed

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
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", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.inference;
11+
12+
public enum EndpointVersions {
13+
FIRST_ENDPOINT_VERSION("2023-09-29"),
14+
PARAMETERS_INTRODUCED_ENDPOINT_VERSION("2024-10-17");
15+
16+
private final String name;
17+
18+
EndpointVersions(String s) {
19+
name = s;
20+
}
21+
22+
public String toString() {
23+
return this.name;
24+
}
25+
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ void parseRequestConfig(
4646
String modelId,
4747
TaskType taskType,
4848
Map<String, Object> config,
49-
String endpointVersion,
49+
EndpointVersions endpointVersion,
5050
ActionListener<Model> parsedModelListener
5151
);
5252

@@ -69,7 +69,7 @@ Model parsePersistedConfigWithSecrets(
6969
TaskType taskType,
7070
Map<String, Object> config,
7171
Map<String, Object> secrets,
72-
String endpointVersion
72+
EndpointVersions endpointVersion
7373
);
7474

7575
/**
@@ -84,7 +84,7 @@ Model parsePersistedConfigWithSecrets(
8484
* @param endpointVersion
8585
* @return The parsed {@link Model}
8686
*/
87-
Model parsePersistedConfig(String modelId, TaskType taskType, Map<String, Object> config, String endpointVersion);
87+
Model parsePersistedConfig(String modelId, TaskType taskType, Map<String, Object> config, EndpointVersions endpointVersion);
8888

8989
/**
9090
* Perform inference on the model.

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

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,6 @@ public class ModelConfigurations implements ToFilteredXContentObject, VersionedN
3333
public static final String CHUNKING_SETTINGS = "chunking_settings";
3434
public static final String INCLUDE_PARAMETERS = "include_parameters";
3535
public static final String ENDPOINT_VERSION_FIELD_NAME = "endpoint_version";
36-
public static final String FIRST_ENDPOINT_VERSION = "2023-09-29";
37-
public static final String PARAMETERS_INTRODUCED_ENDPOINT_VERSION = "2024-10-17";
3836
private static final String NAME = "inference_model";
3937

4038
public static ModelConfigurations of(Model model, TaskSettings taskSettings) {
@@ -73,7 +71,7 @@ public static ModelConfigurations of(Model model, ServiceSettings serviceSetting
7371
private final ServiceSettings serviceSettings;
7472
private final TaskSettings taskSettings;
7573
private final ChunkingSettings chunkingSettings;
76-
private final String endpointVersion;
74+
private final EndpointVersions endpointVersion;
7775

7876
/**
7977
* Allows no task settings to be defined. This will default to the {@link EmptyTaskSettings} object.
@@ -83,7 +81,7 @@ public ModelConfigurations(
8381
TaskType taskType,
8482
String service,
8583
ServiceSettings serviceSettings,
86-
String endpointVersion
84+
EndpointVersions endpointVersion
8785
) {
8886
this(inferenceEntityId, taskType, service, serviceSettings, EmptyTaskSettings.INSTANCE, endpointVersion);
8987
}
@@ -94,7 +92,7 @@ public ModelConfigurations(
9492
String service,
9593
ServiceSettings serviceSettings,
9694
ChunkingSettings chunkingSettings,
97-
String endpointVersion
95+
EndpointVersions endpointVersion
9896
) {
9997
this(inferenceEntityId, taskType, service, serviceSettings, EmptyTaskSettings.INSTANCE, chunkingSettings, endpointVersion);
10098
}
@@ -105,7 +103,7 @@ public ModelConfigurations(
105103
String service,
106104
ServiceSettings serviceSettings,
107105
TaskSettings taskSettings,
108-
String endpointVersion
106+
EndpointVersions endpointVersion
109107
) {
110108
this.inferenceEntityId = Objects.requireNonNull(inferenceEntityId);
111109
this.taskType = Objects.requireNonNull(taskType);
@@ -123,7 +121,7 @@ public ModelConfigurations(
123121
ServiceSettings serviceSettings,
124122
TaskSettings taskSettings,
125123
ChunkingSettings chunkingSettings,
126-
String endpointVersion
124+
EndpointVersions endpointVersion
127125
) {
128126
this.inferenceEntityId = Objects.requireNonNull(inferenceEntityId);
129127
this.taskType = Objects.requireNonNull(taskType);
@@ -144,8 +142,8 @@ public ModelConfigurations(StreamInput in) throws IOException {
144142
? in.readOptionalNamedWriteable(ChunkingSettings.class)
145143
: null;
146144
this.endpointVersion = in.getTransportVersion().onOrAfter(TransportVersions.INFERENCE_API_PARAMATERS_INTRODUCED)
147-
? Objects.requireNonNullElse(in.readOptionalString(), FIRST_ENDPOINT_VERSION)
148-
: FIRST_ENDPOINT_VERSION;
145+
? Objects.requireNonNullElse(in.readEnum(EndpointVersions.class), EndpointVersions.FIRST_ENDPOINT_VERSION)
146+
: EndpointVersions.FIRST_ENDPOINT_VERSION;
149147
}
150148

151149
@Override
@@ -159,7 +157,7 @@ public void writeTo(StreamOutput out) throws IOException {
159157
out.writeOptionalNamedWriteable(chunkingSettings);
160158
}
161159
if (out.getTransportVersion().onOrAfter(TransportVersions.INFERENCE_API_PARAMATERS_INTRODUCED)) {
162-
out.writeOptionalString(endpointVersion); // not nullable after 9.0
160+
out.writeEnum(endpointVersion); // not nullable after 9.0
163161
}
164162
}
165163

@@ -187,7 +185,7 @@ public ChunkingSettings getChunkingSettings() {
187185
return chunkingSettings;
188186
}
189187

190-
public String getEndpointVersion() {
188+
public EndpointVersions getEndpointVersion() {
191189
return endpointVersion;
192190
}
193191

@@ -235,7 +233,7 @@ public XContentBuilder toFilteredXContent(XContentBuilder builder, Params params
235233
builder.field(PARAMETERS, taskSettings);
236234
}
237235
if (params.paramAsBoolean(FOR_INDEX, false)) {
238-
builder.field(ENDPOINT_VERSION_FIELD_NAME, endpointVersion);
236+
builder.field(ENDPOINT_VERSION_FIELD_NAME, endpointVersion);
239237
}
240238
builder.endObject();
241239
return builder;

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,5 +21,5 @@ public record UnparsedModel(
2121
String service,
2222
Map<String, Object> settings,
2323
Map<String, Object> secrets,
24-
String endpointVersion
24+
EndpointVersions endpointVersion
2525
) {}

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/PutInferenceModelAction.java‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.elasticsearch.common.io.stream.StreamInput;
1717
import org.elasticsearch.common.io.stream.StreamOutput;
1818
import org.elasticsearch.common.xcontent.XContentHelper;
19+
import org.elasticsearch.inference.EndpointVersions;
1920
import org.elasticsearch.inference.ModelConfigurations;
2021
import org.elasticsearch.inference.TaskType;
2122
import org.elasticsearch.rest.RestStatus;
@@ -31,10 +32,8 @@
3132
import java.util.Objects;
3233

3334
import static org.elasticsearch.inference.ModelConfigurations.ENDPOINT_VERSION_FIELD_NAME;
34-
import static org.elasticsearch.inference.ModelConfigurations.FIRST_ENDPOINT_VERSION;
3535
import static org.elasticsearch.inference.ModelConfigurations.OLD_TASK_SETTINGS;
3636
import static org.elasticsearch.inference.ModelConfigurations.PARAMETERS;
37-
import static org.elasticsearch.inference.ModelConfigurations.PARAMETERS_INTRODUCED_ENDPOINT_VERSION;
3837

3938
public class PutInferenceModelAction extends ActionType<PutInferenceModelAction.Response> {
4039

@@ -91,12 +90,12 @@ public BytesReference getRewrittenContent() {
9190
);
9291
} else if (newContent.containsKey(PARAMETERS)) {
9392
newContent.put(OLD_TASK_SETTINGS, newContent.get(PARAMETERS));
94-
newContent.put(ENDPOINT_VERSION_FIELD_NAME, PARAMETERS_INTRODUCED_ENDPOINT_VERSION);
93+
newContent.put(ENDPOINT_VERSION_FIELD_NAME, EndpointVersions.PARAMETERS_INTRODUCED_ENDPOINT_VERSION);
9594
newContent.remove(PARAMETERS);
9695
} else if (newContent.containsKey(OLD_TASK_SETTINGS)) {
97-
newContent.put(ENDPOINT_VERSION_FIELD_NAME, FIRST_ENDPOINT_VERSION);
96+
newContent.put(ENDPOINT_VERSION_FIELD_NAME, EndpointVersions.FIRST_ENDPOINT_VERSION);
9897
} else {
99-
newContent.put(ENDPOINT_VERSION_FIELD_NAME, FIRST_ENDPOINT_VERSION);
98+
newContent.put(ENDPOINT_VERSION_FIELD_NAME, EndpointVersions.FIRST_ENDPOINT_VERSION);
10099
}
101100
try (XContentBuilder builder = XContentFactory.contentBuilder(this.contentType)) {
102101
builder.map(newContent);

‎x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/action/PutInferenceModelActionTests.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import org.elasticsearch.common.bytes.BytesArray;
1313
import org.elasticsearch.common.bytes.BytesReference;
1414
import org.elasticsearch.common.xcontent.XContentHelper;
15+
import org.elasticsearch.inference.EndpointVersions;
1516
import org.elasticsearch.inference.TaskType;
1617
import org.elasticsearch.test.ESTestCase;
1718
import org.elasticsearch.xcontent.XContentBuilder;
@@ -25,7 +26,6 @@
2526
import java.util.Map;
2627

2728
import static org.elasticsearch.inference.ModelConfigurations.ENDPOINT_VERSION_FIELD_NAME;
28-
import static org.elasticsearch.inference.ModelConfigurations.FIRST_ENDPOINT_VERSION;
2929
import static org.elasticsearch.inference.ModelConfigurations.OLD_TASK_SETTINGS;
3030
import static org.elasticsearch.inference.ModelConfigurations.PARAMETERS;
3131

@@ -117,6 +117,6 @@ public void testWithTaskSettings() throws IOException {
117117
assertNull(map.get(PARAMETERS));
118118
assertEquals("elasticsearch", map.get("service"));
119119
assertEquals(serviceSettingsValues, map.get("service_settings"));
120-
assertEquals(ENDPOINT_VERSION_FIELD_NAME, FIRST_ENDPOINT_VERSION);
120+
assertEquals(ENDPOINT_VERSION_FIELD_NAME, EndpointVersions.FIRST_ENDPOINT_VERSION);
121121
}
122122
}

‎x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/AbstractTestInferenceService.java‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import org.elasticsearch.common.ValidationException;
1313
import org.elasticsearch.common.io.stream.StreamInput;
1414
import org.elasticsearch.common.io.stream.StreamOutput;
15+
import org.elasticsearch.inference.EndpointVersions;
1516
import org.elasticsearch.inference.InferenceService;
1617
import org.elasticsearch.inference.Model;
1718
import org.elasticsearch.inference.ModelConfigurations;
@@ -60,7 +61,7 @@ public TestServiceModel parsePersistedConfigWithSecrets(
6061
TaskType taskType,
6162
Map<String, Object> config,
6263
Map<String, Object> secrets,
63-
String endpointVersion
64+
EndpointVersions endpointVersion
6465
) {
6566
var serviceSettingsMap = (Map<String, Object>) config.remove(ModelConfigurations.SERVICE_SETTINGS);
6667
var secretSettingsMap = (Map<String, Object>) secrets.remove(ModelSecrets.SECRET_SETTINGS);
@@ -76,7 +77,7 @@ public TestServiceModel parsePersistedConfigWithSecrets(
7677

7778
@Override
7879
@SuppressWarnings("unchecked")
79-
public Model parsePersistedConfig(String modelId, TaskType taskType, Map<String, Object> config, String endpointVersion) {
80+
public Model parsePersistedConfig(String modelId, TaskType taskType, Map<String, Object> config, EndpointVersions endpointVersion) {
8081
var serviceSettingsMap = (Map<String, Object>) config.remove(ModelConfigurations.SERVICE_SETTINGS);
8182

8283
var serviceSettings = getServiceSettingsFromMap(serviceSettingsMap);
@@ -106,7 +107,7 @@ public TestServiceModel(
106107
ServiceSettings serviceSettings,
107108
TestTaskSettings taskSettings,
108109
TestSecretSettings secretSettings,
109-
String endpointVersion
110+
EndpointVersions endpointVersion
110111
) {
111112
super(
112113
new ModelConfigurations(modelId, taskType, service, serviceSettings, taskSettings, endpointVersion),

‎x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestDenseInferenceServiceExtension.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
1919
import org.elasticsearch.inference.ChunkedInferenceServiceResults;
2020
import org.elasticsearch.inference.ChunkingOptions;
21+
import org.elasticsearch.inference.EndpointVersions;
2122
import org.elasticsearch.inference.InferenceServiceExtension;
2223
import org.elasticsearch.inference.InferenceServiceResults;
2324
import org.elasticsearch.inference.InputType;
@@ -46,7 +47,7 @@ public List<Factory> getInferenceServiceFactories() {
4647
}
4748

4849
public static class TestDenseModel extends Model {
49-
public TestDenseModel(String inferenceEntityId, TestServiceSettings serviceSettings, String endpointVersion) {
50+
public TestDenseModel(String inferenceEntityId, TestServiceSettings serviceSettings, EndpointVersions endpointVersion) {
5051
super(
5152
new ModelConfigurations(
5253
inferenceEntityId,
@@ -76,7 +77,7 @@ public void parseRequestConfig(
7677
String modelId,
7778
TaskType taskType,
7879
Map<String, Object> config,
79-
String endpointVersion,
80+
EndpointVersions endpointVersion,
8081
ActionListener<Model> parsedModelListener
8182
) {
8283
var serviceSettingsMap = (Map<String, Object>) config.remove(ModelConfigurations.SERVICE_SETTINGS);

‎x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestRerankingServiceExtension.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import org.elasticsearch.core.TimeValue;
1818
import org.elasticsearch.inference.ChunkedInferenceServiceResults;
1919
import org.elasticsearch.inference.ChunkingOptions;
20+
import org.elasticsearch.inference.EndpointVersions;
2021
import org.elasticsearch.inference.InferenceServiceExtension;
2122
import org.elasticsearch.inference.InferenceServiceResults;
2223
import org.elasticsearch.inference.InputType;
@@ -42,7 +43,7 @@ public List<Factory> getInferenceServiceFactories() {
4243
}
4344

4445
public static class TestRerankingModel extends Model {
45-
public TestRerankingModel(String inferenceEntityId, TestServiceSettings serviceSettings, String endpointVersion) {
46+
public TestRerankingModel(String inferenceEntityId, TestServiceSettings serviceSettings, EndpointVersions endpointVersion) {
4647
super(
4748
new ModelConfigurations(inferenceEntityId, TaskType.RERANK, TestInferenceService.NAME, serviceSettings, endpointVersion),
4849
new ModelSecrets(new AbstractTestInferenceService.TestSecretSettings("api_key"))
@@ -66,7 +67,7 @@ public void parseRequestConfig(
6667
String modelId,
6768
TaskType taskType,
6869
Map<String, Object> config,
69-
String endpointVersion,
70+
EndpointVersions endpointVersion,
7071
ActionListener<Model> parsedModelListener
7172
) {
7273
var serviceSettingsMap = (Map<String, Object>) config.remove(ModelConfigurations.SERVICE_SETTINGS);

‎x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestSparseInferenceServiceExtension.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import org.elasticsearch.core.TimeValue;
1818
import org.elasticsearch.inference.ChunkedInferenceServiceResults;
1919
import org.elasticsearch.inference.ChunkingOptions;
20+
import org.elasticsearch.inference.EndpointVersions;
2021
import org.elasticsearch.inference.InferenceServiceExtension;
2122
import org.elasticsearch.inference.InferenceServiceResults;
2223
import org.elasticsearch.inference.InputType;
@@ -45,7 +46,7 @@ public List<Factory> getInferenceServiceFactories() {
4546
}
4647

4748
public static class TestSparseModel extends Model {
48-
public TestSparseModel(String inferenceEntityId, TestServiceSettings serviceSettings, String endpointVersion) {
49+
public TestSparseModel(String inferenceEntityId, TestServiceSettings serviceSettings, EndpointVersions endpointVersion) {
4950
super(
5051
new ModelConfigurations(
5152
inferenceEntityId,
@@ -75,7 +76,7 @@ public void parseRequestConfig(
7576
String modelId,
7677
TaskType taskType,
7778
Map<String, Object> config,
78-
String endpointVersion,
79+
EndpointVersions endpointVersion,
7980
ActionListener<Model> parsedModelListener
8081
) {
8182
var serviceSettingsMap = (Map<String, Object>) config.remove(ModelConfigurations.SERVICE_SETTINGS);

0 commit comments

Comments
 (0)