Skip to content

Commit a1af0bc

Browse files
Mikep86Mikep86
authored
[9.1] Semantic Text Index Options Integration Tests (#130453) (#130581)
* Semantic Text Index Options Integration Tests (#130453) * Remove default BBQ index options test
1 parent 253889e commit a1af0bc

File tree

2 files changed

+278
-1
lines changed

2 files changed

+278
-1
lines changed

‎x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java

Lines changed: 277 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,277 @@
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.integration;
9+
10+
import org.elasticsearch.action.admin.indices.mapping.get.GetFieldMappingsAction;
11+
import org.elasticsearch.action.admin.indices.mapping.get.GetFieldMappingsRequest;
12+
import org.elasticsearch.action.support.IndicesOptions;
13+
import org.elasticsearch.common.bytes.BytesReference;
14+
import org.elasticsearch.common.settings.Settings;
15+
import org.elasticsearch.common.xcontent.XContentHelper;
16+
import org.elasticsearch.core.Nullable;
17+
import org.elasticsearch.core.TimeValue;
18+
import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
19+
import org.elasticsearch.index.mapper.vectors.IndexOptions;
20+
import org.elasticsearch.inference.TaskType;
21+
import org.elasticsearch.license.GetLicenseAction;
22+
import org.elasticsearch.license.License;
23+
import org.elasticsearch.license.LicenseSettings;
24+
import org.elasticsearch.license.PostStartBasicAction;
25+
import org.elasticsearch.license.PostStartBasicRequest;
26+
import org.elasticsearch.license.PutLicenseAction;
27+
import org.elasticsearch.license.PutLicenseRequest;
28+
import org.elasticsearch.license.TestUtils;
29+
import org.elasticsearch.plugins.Plugin;
30+
import org.elasticsearch.protocol.xpack.license.GetLicenseRequest;
31+
import org.elasticsearch.reindex.ReindexPlugin;
32+
import org.elasticsearch.test.ESIntegTestCase;
33+
import org.elasticsearch.test.InternalTestCluster;
34+
import org.elasticsearch.xcontent.ToXContent;
35+
import org.elasticsearch.xcontent.XContentBuilder;
36+
import org.elasticsearch.xcontent.XContentFactory;
37+
import org.elasticsearch.xcontent.XContentType;
38+
import org.elasticsearch.xpack.core.inference.action.DeleteInferenceEndpointAction;
39+
import org.elasticsearch.xpack.core.inference.action.PutInferenceModelAction;
40+
import org.elasticsearch.xpack.inference.InferenceIndex;
41+
import org.elasticsearch.xpack.inference.LocalStateInferencePlugin;
42+
import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper;
43+
import org.elasticsearch.xpack.inference.mock.TestDenseInferenceServiceExtension;
44+
import org.elasticsearch.xpack.inference.mock.TestInferenceServicePlugin;
45+
import org.elasticsearch.xpack.inference.mock.TestSparseInferenceServiceExtension;
46+
import org.junit.After;
47+
import org.junit.Before;
48+
49+
import java.io.IOException;
50+
import java.util.Collection;
51+
import java.util.HashMap;
52+
import java.util.List;
53+
import java.util.Map;
54+
55+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
56+
import static org.hamcrest.CoreMatchers.equalTo;
57+
58+
public class SemanticTextIndexOptionsIT extends ESIntegTestCase {
59+
private static final String INDEX_NAME = "test-index";
60+
private static final Map<String, Object> BBQ_COMPATIBLE_SERVICE_SETTINGS = Map.of(
61+
"model",
62+
"my_model",
63+
"dimensions",
64+
256,
65+
"similarity",
66+
"cosine",
67+
"api_key",
68+
"my_api_key"
69+
);
70+
71+
private final Map<String, TaskType> inferenceIds = new HashMap<>();
72+
73+
@Override
74+
protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
75+
return Settings.builder().put(LicenseSettings.SELF_GENERATED_LICENSE_TYPE.getKey(), "trial").build();
76+
}
77+
78+
@Override
79+
protected Collection<Class<? extends Plugin>> nodePlugins() {
80+
return List.of(LocalStateInferencePlugin.class, TestInferenceServicePlugin.class, ReindexPlugin.class);
81+
}
82+
83+
@Before
84+
public void resetLicense() throws Exception {
85+
setLicense(License.LicenseType.TRIAL);
86+
}
87+
88+
@After
89+
public void cleanUp() {
90+
assertAcked(
91+
safeGet(
92+
client().admin()
93+
.indices()
94+
.prepareDelete(INDEX_NAME)
95+
.setIndicesOptions(
96+
IndicesOptions.builder().concreteTargetOptions(new IndicesOptions.ConcreteTargetOptions(true)).build()
97+
)
98+
.execute()
99+
)
100+
);
101+
102+
for (var entry : inferenceIds.entrySet()) {
103+
assertAcked(
104+
safeGet(
105+
client().execute(
106+
DeleteInferenceEndpointAction.INSTANCE,
107+
new DeleteInferenceEndpointAction.Request(entry.getKey(), entry.getValue(), true, false)
108+
)
109+
)
110+
);
111+
}
112+
}
113+
114+
public void testValidateIndexOptionsWithBasicLicense() throws Exception {
115+
final String inferenceId = "test-inference-id-1";
116+
final String inferenceFieldName = "inference_field";
117+
createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS);
118+
downgradeLicenseAndRestartCluster();
119+
120+
IndexOptions indexOptions = new DenseVectorFieldMapper.Int8HnswIndexOptions(
121+
randomIntBetween(1, 100),
122+
randomIntBetween(1, 10_000),
123+
null,
124+
null
125+
);
126+
assertAcked(
127+
safeGet(prepareCreate(INDEX_NAME).setMapping(generateMapping(inferenceFieldName, inferenceId, indexOptions)).execute())
128+
);
129+
130+
final Map<String, Object> expectedFieldMapping = generateExpectedFieldMapping(inferenceFieldName, inferenceId, indexOptions);
131+
assertThat(getFieldMappings(inferenceFieldName, false), equalTo(expectedFieldMapping));
132+
}
133+
134+
private void createInferenceEndpoint(TaskType taskType, String inferenceId, Map<String, Object> serviceSettings) throws IOException {
135+
final String service = switch (taskType) {
136+
case TEXT_EMBEDDING -> TestDenseInferenceServiceExtension.TestInferenceService.NAME;
137+
case SPARSE_EMBEDDING -> TestSparseInferenceServiceExtension.TestInferenceService.NAME;
138+
default -> throw new IllegalArgumentException("Unhandled task type [" + taskType + "]");
139+
};
140+
141+
final BytesReference content;
142+
try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
143+
builder.startObject();
144+
builder.field("service", service);
145+
builder.field("service_settings", serviceSettings);
146+
builder.endObject();
147+
148+
content = BytesReference.bytes(builder);
149+
}
150+
151+
PutInferenceModelAction.Request request = new PutInferenceModelAction.Request(
152+
taskType,
153+
inferenceId,
154+
content,
155+
XContentType.JSON,
156+
TEST_REQUEST_TIMEOUT
157+
);
158+
var responseFuture = client().execute(PutInferenceModelAction.INSTANCE, request);
159+
assertThat(responseFuture.actionGet(TEST_REQUEST_TIMEOUT).getModel().getInferenceEntityId(), equalTo(inferenceId));
160+
161+
inferenceIds.put(inferenceId, taskType);
162+
}
163+
164+
private static XContentBuilder generateMapping(String inferenceFieldName, String inferenceId, @Nullable IndexOptions indexOptions)
165+
throws IOException {
166+
XContentBuilder mapping = XContentFactory.jsonBuilder();
167+
mapping.startObject();
168+
mapping.field("properties");
169+
generateFieldMapping(mapping, inferenceFieldName, inferenceId, indexOptions);
170+
mapping.endObject();
171+
172+
return mapping;
173+
}
174+
175+
private static void generateFieldMapping(
176+
XContentBuilder builder,
177+
String inferenceFieldName,
178+
String inferenceId,
179+
@Nullable IndexOptions indexOptions
180+
) throws IOException {
181+
builder.startObject();
182+
builder.startObject(inferenceFieldName);
183+
builder.field("type", SemanticTextFieldMapper.CONTENT_TYPE);
184+
builder.field("inference_id", inferenceId);
185+
if (indexOptions != null) {
186+
builder.startObject("index_options");
187+
if (indexOptions instanceof DenseVectorFieldMapper.DenseVectorIndexOptions) {
188+
builder.field("dense_vector");
189+
indexOptions.toXContent(builder, ToXContent.EMPTY_PARAMS);
190+
}
191+
builder.endObject();
192+
}
193+
builder.endObject();
194+
builder.endObject();
195+
}
196+
197+
private static Map<String, Object> generateExpectedFieldMapping(
198+
String inferenceFieldName,
199+
String inferenceId,
200+
@Nullable IndexOptions indexOptions
201+
) throws IOException {
202+
Map<String, Object> expectedFieldMapping;
203+
try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
204+
generateFieldMapping(builder, inferenceFieldName, inferenceId, indexOptions);
205+
expectedFieldMapping = XContentHelper.convertToMap(BytesReference.bytes(builder), false, XContentType.JSON).v2();
206+
}
207+
208+
return expectedFieldMapping;
209+
}
210+
211+
@SuppressWarnings("unchecked")
212+
private static Map<String, Object> filterNullOrEmptyValues(Map<String, Object> map) {
213+
Map<String, Object> filteredMap = new HashMap<>();
214+
for (var entry : map.entrySet()) {
215+
Object value = entry.getValue();
216+
if (entry.getValue() instanceof Map<?, ?> mapValue) {
217+
if (mapValue.isEmpty()) {
218+
continue;
219+
}
220+
221+
value = filterNullOrEmptyValues((Map<String, Object>) mapValue);
222+
}
223+
224+
if (value != null) {
225+
filteredMap.put(entry.getKey(), value);
226+
}
227+
}
228+
229+
return filteredMap;
230+
}
231+
232+
private static Map<String, Object> getFieldMappings(String fieldName, boolean includeDefaults) {
233+
var request = new GetFieldMappingsRequest().indices(INDEX_NAME).fields(fieldName).includeDefaults(includeDefaults);
234+
return safeGet(client().execute(GetFieldMappingsAction.INSTANCE, request)).fieldMappings(INDEX_NAME, fieldName).sourceAsMap();
235+
}
236+
237+
private static void setLicense(License.LicenseType type) throws Exception {
238+
if (type == License.LicenseType.BASIC) {
239+
assertAcked(
240+
safeGet(
241+
client().execute(
242+
PostStartBasicAction.INSTANCE,
243+
new PostStartBasicRequest(TEST_REQUEST_TIMEOUT, TEST_REQUEST_TIMEOUT).acknowledge(true)
244+
)
245+
)
246+
);
247+
} else {
248+
License license = TestUtils.generateSignedLicense(
249+
type.getTypeName(),
250+
License.VERSION_CURRENT,
251+
-1,
252+
TimeValue.timeValueHours(24)
253+
);
254+
assertAcked(
255+
safeGet(
256+
client().execute(
257+
PutLicenseAction.INSTANCE,
258+
new PutLicenseRequest(TEST_REQUEST_TIMEOUT, TEST_REQUEST_TIMEOUT).license(license)
259+
)
260+
)
261+
);
262+
}
263+
}
264+
265+
private static void assertLicense(License.LicenseType type) {
266+
var getLicenseResponse = safeGet(client().execute(GetLicenseAction.INSTANCE, new GetLicenseRequest(TEST_REQUEST_TIMEOUT)));
267+
assertThat(getLicenseResponse.license().type(), equalTo(type.getTypeName()));
268+
}
269+
270+
private void downgradeLicenseAndRestartCluster() throws Exception {
271+
// Downgrade the license and restart the cluster to force the model registry to rebuild
272+
setLicense(License.LicenseType.BASIC);
273+
internalCluster().fullRestart(new InternalTestCluster.RestartCallback());
274+
ensureGreen(InferenceIndex.INDEX_NAME);
275+
assertLicense(License.LicenseType.BASIC);
276+
}
277+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1223,7 +1223,7 @@ static boolean indexVersionDefaultsToBbqHnsw(IndexVersion indexVersion) {
12231223
|| indexVersion.between(SEMANTIC_TEXT_DEFAULTS_TO_BBQ_BACKPORT_8_X, IndexVersions.UPGRADE_TO_LUCENE_10_0_0);
12241224
}
12251225

1226-
static DenseVectorFieldMapper.DenseVectorIndexOptions defaultBbqHnswDenseVectorIndexOptions() {
1226+
public static DenseVectorFieldMapper.DenseVectorIndexOptions defaultBbqHnswDenseVectorIndexOptions() {
12271227
int m = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
12281228
int efConstruction = Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
12291229
DenseVectorFieldMapper.RescoreVector rescoreVector = new DenseVectorFieldMapper.RescoreVector(DEFAULT_RESCORE_OVERSAMPLE);

0 commit comments

Comments
 (0)