[Semantic Text] Integration Test

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

@@ -0,0 +1,180 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.integration;
+
+import org.elasticsearch.action.DocWriteResponse;
+import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.index.IndexVersion;
+import org.elasticsearch.index.mapper.InferenceMetadataFieldsMapper;
+import org.elasticsearch.license.LicenseSettings;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
+import org.elasticsearch.test.ESIntegTestCase;
+import org.elasticsearch.test.index.IndexVersionUtils;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xpack.inference.LocalStateInferencePlugin;
+import org.elasticsearch.xpack.inference.Utils;
+import org.elasticsearch.xpack.inference.mock.TestSparseInferenceServiceExtension;
+import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder;
+import org.elasticsearch.xpack.inference.registry.ModelRegistry;
+import org.junit.Before;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHighlight;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
+import static org.hamcrest.Matchers.equalTo;
+
+public class SemanticTextIndexVersionIT extends ESIntegTestCase {
+    private static final IndexVersion SEMANTIC_TEXT_INTRODUCED_VERSION = IndexVersion.fromId(8512000);
+    private static final double PERCENTAGE_TO_TEST = 0.5;
+    private Set<IndexVersion> availableVersions;
+
+    @Before
+    public void setup() throws Exception {
+        ModelRegistry modelRegistry = internalCluster().getCurrentMasterNodeInstance(ModelRegistry.class);
+        Utils.storeSparseModel(modelRegistry);
+        availableVersions = IndexVersionUtils.allReleasedVersions()
+            .stream()
+            .filter(indexVersion -> indexVersion.after(SEMANTIC_TEXT_INTRODUCED_VERSION))
+            .collect(Collectors.toSet());
+    }
+
+    @Override
+    protected boolean forbidPrivateIndexSettings() {
+        return false;
+    }
+
+    @Override
+    protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
+        return Settings.builder().put(otherSettings).put(LicenseSettings.SELF_GENERATED_LICENSE_TYPE.getKey(), "trial").build();
+    }
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return List.of(LocalStateInferencePlugin.class);
+    }
+
+    /**
+     * Generate settings for an index with a specific version.
+     */
+    private Settings getIndexSettingsWithVersion(IndexVersion version) {
+        return Settings.builder().put(indexSettings()).put("index.version.created", version).build();
+    }
+
+    /**
+     * Creates a subset of indices with different versions for testing.
+     *
+     * @return Map of created indices with their versions
+     */
+    protected Map<String, IndexVersion> createRandomVersionIndices() throws IOException {
+        int versionsCount = (int) Math.ceil(availableVersions.size() * PERCENTAGE_TO_TEST);
+        List<IndexVersion> selectedVersions = randomSubsetOf(versionsCount, availableVersions);
+        Map<String, IndexVersion> result = new HashMap<>();
+
+        for (int i = 0; i < selectedVersions.size(); i++) {
+            String indexName = "test_semantic" + "_" + i;
+            IndexVersion version = selectedVersions.get(i);
+            createIndex(indexName, getIndexSettingsWithVersion(version));
+            result.put(indexName, version);
+        }
+
+        return result;
+    }
+
+    /**
+     * This test creates an index, ingests data, and performs searches (including highlighting when applicable)
+     * for a selected subset of index versions.
+     */
+    public void testSemanticText() throws Exception {
+        Map<String, IndexVersion> indices = createRandomVersionIndices();
+        for (String indexName : indices.keySet()) {
+            IndexVersion version = indices.get(indexName);
+
+            // Test index creation
+            assertTrue("Index " + indexName + " should exist", indexExists(indexName));
+            assertEquals(
+                "Index version should match",
+                version.id(),
+                client().admin()
+                    .indices()
+                    .prepareGetSettings(TimeValue.THIRTY_SECONDS, indexName)
+                    .get()
+                    .getIndexToSettings()
+                    .get(indexName)
+                    .getAsVersionId("index.version.created", IndexVersion::fromId)
+                    .id()
+            );
+
+            // Test update mapping
+            XContentBuilder mapping = XContentFactory.jsonBuilder()
+                .startObject()
+                .startObject("properties")
+                .startObject("semantic_field")
+                .field("type", "semantic_text")
+                .field("inference_id", TestSparseInferenceServiceExtension.TestInferenceService.NAME)
+                .endObject()
+                .endObject()
+                .endObject();
+
+            assertAcked(client().admin().indices().preparePutMapping(indexName).setSource(mapping).get());
+
+            // Test data ingestion
+            String[] text = new String[] { "inference test", "another inference test" };
+            DocWriteResponse docWriteResponse = client().prepareIndex(indexName).setSource(Map.of("semantic_field", text)).get();
+
+            assertEquals("Document should be created", "created", docWriteResponse.getResult().toString().toLowerCase(Locale.ROOT));
+
+            // Ensure index is ready
+            client().admin().indices().refresh(new RefreshRequest(indexName)).get();
+            ensureGreen(indexName);
+
+            // Semantic Search
+            SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().query(new SemanticQueryBuilder("semantic_field", "inference"))
+                .trackTotalHits(true);
+
+            assertResponse(
+                client().search(new SearchRequest(indexName).source(sourceBuilder)),
+                response -> { assertHitCount(response, 1L); }
+            );
+
+            Settings settings = client().admin()
+                .indices()
+                .prepareGetSettings(TimeValue.THIRTY_SECONDS, indexName)
+                .get()
+                .getIndexToSettings()
+                .get(indexName);
+
+            // Semantic Search with highlighter only available from 8.18 and 9.0
+            if (InferenceMetadataFieldsMapper.isEnabled(settings)) {
+                SearchSourceBuilder sourceHighlighterBuilder = new SearchSourceBuilder().query(
+                    new SemanticQueryBuilder("semantic_field", "inference")
+                ).highlighter(new HighlightBuilder().field("semantic_field")).trackTotalHits(true);
+
+                assertResponse(client().search(new SearchRequest(indexName).source(sourceHighlighterBuilder)), response -> {
+                    assertHighlight(response, 0, "semantic_field", 0, 2, equalTo("inference test"));
+                    assertHighlight(response, 0, "semantic_field", 1, 2, equalTo("another inference test"));
+                });
+            }
+        }
+    }
+}

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/LocalStateInferencePlugin.java

@@ -8,17 +8,24 @@
 package org.elasticsearch.xpack.inference;
 
 import org.elasticsearch.action.support.MappedActionFilter;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.mapper.Mapper;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.inference.InferenceServiceExtension;
 import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.plugins.SearchPlugin;
+import org.elasticsearch.search.fetch.subphase.highlight.Highlighter;
 import org.elasticsearch.xpack.core.LocalStateCompositeXPackPlugin;
+import org.elasticsearch.xpack.core.ml.inference.MlInferenceNamedXContentProvider;
+import org.elasticsearch.xpack.core.ml.search.SparseVectorQueryBuilder;
 import org.elasticsearch.xpack.core.ssl.SSLService;
+import org.elasticsearch.xpack.inference.highlight.SemanticTextHighlighter;
 import org.elasticsearch.xpack.inference.mock.TestDenseInferenceServiceExtension;
 import org.elasticsearch.xpack.inference.mock.TestSparseInferenceServiceExtension;
 
 import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -63,8 +70,29 @@ public Map<String, Mapper.TypeParser> getMappers() {
         return inferencePlugin.getMappers();
     }
 
+    @Override
+    public Map<String, Highlighter> getHighlighters() {
+        return Map.of(SemanticTextHighlighter.NAME, new SemanticTextHighlighter());
+    }
+
     @Override
     public Collection<MappedActionFilter> getMappedActionFilters() {
         return inferencePlugin.getMappedActionFilters();
     }
+
+    @Override
+    public List<QuerySpec<?>> getQueries() {
+        return inferencePlugin.getQueries();
+    }
+
+    @Override
+    public List<NamedWriteableRegistry.Entry> getNamedWriteables() {
+        List<NamedWriteableRegistry.Entry> namedWriteables = new ArrayList<>(super.getNamedWriteables());
+        namedWriteables.addAll(new MlInferenceNamedXContentProvider().getNamedWriteables());
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(QueryBuilder.class, SparseVectorQueryBuilder.NAME, SparseVectorQueryBuilder::new)
+        );
+
+        return namedWriteables;
+    }
 }