New SemanticFieldInferenceFieldsIT test

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

@@ -0,0 +1,232 @@
+/*
+ * 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.search.SearchRequest;
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.common.document.DocumentField;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.CheckedSupplier;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.index.IndexVersion;
+import org.elasticsearch.index.mapper.InferenceMetadataFieldsMapper;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.license.LicenseSettings;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.reindex.ReindexPlugin;
+import org.elasticsearch.search.SearchHit;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
+import org.elasticsearch.search.lookup.SourceFilter;
+import org.elasticsearch.test.ESIntegTestCase;
+import org.elasticsearch.test.index.IndexVersionUtils;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xpack.inference.FakeMlPlugin;
+import org.elasticsearch.xpack.inference.LocalStateInferencePlugin;
+import org.elasticsearch.xpack.inference.mock.TestInferenceServicePlugin;
+import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Consumer;
+
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.notNullValue;
+import static org.hamcrest.CoreMatchers.nullValue;
+
+abstract class AbstractInferenceFieldsIT extends ESIntegTestCase {
+
+    @Override
+    protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
+        return Settings.builder().put(LicenseSettings.SELF_GENERATED_LICENSE_TYPE.getKey(), "trial").build();
+    }
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return List.of(LocalStateInferencePlugin.class, TestInferenceServicePlugin.class, ReindexPlugin.class, FakeMlPlugin.class);
+    }
+
+    @Override
+    protected boolean forbidPrivateIndexSettings() {
+        return false;
+    }
+
+    protected Settings generateIndexSettings(IndexVersion indexVersion) {
+        int numDataNodes = internalCluster().numDataNodes();
+        return Settings.builder()
+            .put(IndexMetadata.SETTING_VERSION_CREATED, indexVersion)
+            .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, numDataNodes)
+            .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+            .build();
+    }
+
+    protected void testExcludeInferenceFieldsFromSource(
+        CheckedSupplier<XContentBuilder, IOException> mappingFactory,
+        Collection<String> fieldsToValidate,
+        IndexVersion minIndexVersion,
+        IndexVersion maxIndexVersion,
+        int iterationCount
+    ) throws Exception {
+        final String indexName = randomIdentifier();
+
+        for (int i = 0; i < iterationCount; i++) {
+            final IndexVersion indexVersion = IndexVersionUtils.randomVersionBetween(minIndexVersion, maxIndexVersion);
+            final Settings indexSettings = generateIndexSettings(indexVersion);
+            XContentBuilder mappings = mappingFactory.get();
+            assertAcked(prepareCreate(indexName).setSettings(indexSettings).setMapping(mappings));
+
+            final int docCount = randomIntBetween(10, 50);
+            fieldsToValidate.forEach(fieldName -> indexDocuments(indexName, fieldName, docCount));
+
+            for (String fieldName : fieldsToValidate) {
+                // For each field defined in the mapping, fetch all documents with that field and verify that
+                // response only contains the requested fieldName and no inference fields.
+                assertSearchResponse(
+                    indexName,
+                    new SemanticQueryBuilder(fieldName, randomAlphaOfLength(10)),
+                    indexSettings,
+                    docCount,
+                    request -> {
+                        request.source().fetchSource(generateRandomFetchSourceContext()).fetchField(fieldName);
+                    },
+                    response -> {
+                        for (SearchHit hit : response.getHits()) {
+                            Map<String, DocumentField> documentFields = hit.getDocumentFields();
+                            assertThat(documentFields.size(), is(1));
+                            assertThat(documentFields.containsKey(fieldName), is(true));
+                        }
+                    }
+                );
+            }
+
+            IntegrationTestUtils.deleteIndex(client(), indexName);
+        }
+    }
+
+    protected abstract void indexDocuments(String indexName, String field, int count);
+
+    protected void assertSearchResponse(
+        String indexName,
+        QueryBuilder queryBuilder,
+        Settings indexSettings,
+        int expectedHits,
+        @Nullable Consumer<SearchRequest> searchRequestModifier,
+        @Nullable Consumer<SearchResponse> searchResponseValidator
+    ) throws Exception {
+        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder().query(queryBuilder).size(expectedHits);
+        SearchRequest searchRequest = new SearchRequest(new String[] { indexName }, searchSourceBuilder);
+        if (searchRequestModifier != null) {
+            searchRequestModifier.accept(searchRequest);
+        }
+
+        ExpectedSource expectedSource = getExpectedSource(indexSettings, searchRequest.source().fetchSource());
+        assertResponse(client().search(searchRequest), response -> {
+            assertThat(response.getSuccessfulShards(), equalTo(response.getTotalShards()));
+            assertThat(response.getHits().getTotalHits().value(), equalTo((long) expectedHits));
+
+            for (SearchHit hit : response.getHits()) {
+                switch (expectedSource) {
+                    case NONE -> assertThat(hit.getSourceAsMap(), nullValue());
+                    case INFERENCE_FIELDS_EXCLUDED -> {
+                        Map<String, Object> sourceAsMap = hit.getSourceAsMap();
+                        assertThat(sourceAsMap, notNullValue());
+                        assertThat(sourceAsMap.containsKey(InferenceMetadataFieldsMapper.NAME), is(false));
+                    }
+                    case INFERENCE_FIELDS_INCLUDED -> {
+                        Map<String, Object> sourceAsMap = hit.getSourceAsMap();
+                        assertThat(sourceAsMap, notNullValue());
+                        assertThat(sourceAsMap.containsKey(InferenceMetadataFieldsMapper.NAME), is(true));
+                    }
+                }
+            }
+
+            if (searchResponseValidator != null) {
+                searchResponseValidator.accept(response);
+            }
+        });
+    }
+
+    private static ExpectedSource getExpectedSource(Settings indexSettings, FetchSourceContext fetchSourceContext) {
+        if (fetchSourceContext != null && fetchSourceContext.fetchSource() == false) {
+            return ExpectedSource.NONE;
+        } else if (InferenceMetadataFieldsMapper.isEnabled(indexSettings) == false) {
+            return ExpectedSource.INFERENCE_FIELDS_EXCLUDED;
+        }
+
+        if (fetchSourceContext != null) {
+            SourceFilter filter = fetchSourceContext.filter();
+            if (filter != null) {
+                if (Arrays.asList(filter.getExcludes()).contains(InferenceMetadataFieldsMapper.NAME)) {
+                    return ExpectedSource.INFERENCE_FIELDS_EXCLUDED;
+                } else if (filter.getIncludes().length > 0) {
+                    return Arrays.asList(filter.getIncludes()).contains(InferenceMetadataFieldsMapper.NAME)
+                        ? ExpectedSource.INFERENCE_FIELDS_INCLUDED
+                        : ExpectedSource.INFERENCE_FIELDS_EXCLUDED;
+                }
+            }
+
+            Boolean excludeInferenceFieldsExplicit = fetchSourceContext.excludeInferenceFields();
+            if (excludeInferenceFieldsExplicit != null) {
+                return excludeInferenceFieldsExplicit ? ExpectedSource.INFERENCE_FIELDS_EXCLUDED : ExpectedSource.INFERENCE_FIELDS_INCLUDED;
+            }
+        }
+
+        return ExpectedSource.INFERENCE_FIELDS_EXCLUDED;
+    }
+
+    protected static FetchSourceContext generateRandomFetchSourceContext() {
+        FetchSourceContext fetchSourceContext = switch (randomIntBetween(0, 4)) {
+            case 0 -> FetchSourceContext.FETCH_SOURCE;
+            case 1 -> FetchSourceContext.FETCH_ALL_SOURCE;
+            case 2 -> FetchSourceContext.FETCH_ALL_SOURCE_EXCLUDE_INFERENCE_FIELDS;
+            case 3 -> FetchSourceContext.DO_NOT_FETCH_SOURCE;
+            case 4 -> null;
+            default -> throw new IllegalStateException("Unhandled randomized case");
+        };
+
+        if (fetchSourceContext != null && fetchSourceContext.fetchSource()) {
+            String[] includes = null;
+            String[] excludes = null;
+            if (randomBoolean()) {
+                // Randomly include a non-existent field to test explicit inclusion handling
+                String field = randomBoolean() ? InferenceMetadataFieldsMapper.NAME : randomIdentifier();
+                includes = new String[] { field };
+            }
+            if (randomBoolean()) {
+                // Randomly exclude a non-existent field to test implicit inclusion handling
+                String field = randomBoolean() ? InferenceMetadataFieldsMapper.NAME : randomIdentifier();
+                excludes = new String[] { field };
+            }
+
+            if (includes != null || excludes != null) {
+                fetchSourceContext = FetchSourceContext.of(
+                    fetchSourceContext.fetchSource(),
+                    fetchSourceContext.excludeVectors(),
+                    fetchSourceContext.excludeInferenceFields(),
+                    includes,
+                    excludes
+                );
+            }
+        }
+
+        return fetchSourceContext;
+    }
+
+    protected enum ExpectedSource {
+        NONE,
+        INFERENCE_FIELDS_EXCLUDED,
+        INFERENCE_FIELDS_INCLUDED
+    }
+}

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.core.inference.action.DeleteInferenceEndpointAction;
 import org.elasticsearch.xpack.core.inference.action.PutInferenceModelAction;
+import org.elasticsearch.xpack.inference.mapper.SemanticFieldMapper;
 import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper;
 import org.elasticsearch.xpack.inference.mock.TestDenseInferenceServiceExtension;
 import org.elasticsearch.xpack.inference.mock.TestSparseInferenceServiceExtension;
@@ -101,4 +102,17 @@ public static XContentBuilder generateSemanticTextMapping(Map<String, String> se
 
         return mapping;
     }
+
+    public static XContentBuilder generateSemanticFieldMapping(Map<String, String> semanticFields) throws IOException {
+        XContentBuilder mapping = XContentFactory.jsonBuilder().startObject().startObject("properties");
+        for (var entry : semanticFields.entrySet()) {
+            mapping.startObject(entry.getKey());
+            mapping.field("type", SemanticFieldMapper.CONTENT_TYPE);
+            mapping.field("inference_id", entry.getValue());
+            mapping.endObject();
+        }
+        mapping.endObject().endObject();
+
+        return mapping;
+    }
 }

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

@@ -0,0 +1,71 @@
+/*
+ * 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.index.IndexVersion;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.test.ESIntegTestCase;
+
+import java.util.Map;
+
+import static org.hamcrest.CoreMatchers.is;
+
+/**
+ * Integration tests verifying that the {@code _inference_fields} metadata field is correctly
+ * included or excluded from {@code _source} for the {@code semantic} field type.
+ * This test covers text and image modalities.
+ */
+@ESIntegTestCase.ClusterScope(numDataNodes = 3, numClientNodes = 0, supportsDedicatedMasters = false)
+public class SemanticFieldInferenceFieldsIT extends AbstractInferenceFieldsIT {
+    private static final Map<String, Object> EMBEDDING_SERVICE_SETTINGS = Map.of(
+        "model",
+        "my_model",
+        "dimensions",
+        256,
+        "similarity",
+        "cosine",
+        "api_key",
+        "my_api_key"
+    );
+
+    // A small JPEG encoded as a base64 data URI for use as a multimodal input.
+    private static final Map<String, String> IMAGE_INPUT = Map.of(
+        "type",
+        "image",
+        "value",
+        "data:image/jpeg;base64,Y2F0IG9uIGEgd2luZG93c2lsbA=="
+    );
+
+    public void testExcludeInferenceFieldsFromSourceWithMixedInputs() throws Exception {
+        final String semanticFieldName = randomIdentifier();
+        final String embeddingInferenceId = randomIdentifier();
+        IntegrationTestUtils.createInferenceEndpoint(client(), TaskType.EMBEDDING, embeddingInferenceId, EMBEDDING_SERVICE_SETTINGS);
+        final var fieldMap = Map.of(semanticFieldName, embeddingInferenceId);
+
+        testExcludeInferenceFieldsFromSource(
+            () -> IntegrationTestUtils.generateSemanticFieldMapping(fieldMap),
+            fieldMap.keySet(),
+            IndexVersion.current(),
+            IndexVersion.current(),
+            20
+        );
+
+        IntegrationTestUtils.deleteInferenceEndpoint(client(), TaskType.EMBEDDING, embeddingInferenceId);
+    }
+
+    @Override
+    protected void indexDocuments(String indexName, String field, int count) {
+        for (int i = 0; i < count; i++) {
+            Map<String, Object> source = Map.of(field, randomBoolean() ? randomAlphaOfLength(10) : IMAGE_INPUT);
+            DocWriteResponse response = client().prepareIndex(indexName).setSource(source).get(TEST_REQUEST_TIMEOUT);
+            assertThat(response.getResult(), is(DocWriteResponse.Result.CREATED));
+        }
+        client().admin().indices().prepareRefresh(indexName).get();
+    }
+}