From 584cbea22d60cd805aafb42b3f43a4c207fec83b Mon Sep 17 00:00:00 2001 From: David Causse Date: Tue, 28 Jul 2026 00:07:26 +0200 Subject: [PATCH] Add SemanticHighlighterQueryEnricherProcessor A new search request processor to enrich the semantic highlighter model id. Works the same way as the neural_query_enricher allowing to define a global or a per field model id. The processor inspects the top level highlight config but also any highlight definitions added to nested query inner hits. Limitations: it does not inspect aggregations (which could possibly be useful for top hits aggregation) and this is left as a possible followup. Main reason for not doing it now is that the batch inference does not support aggregations yet. Resolves #1916 Signed-off-by: David Causse --- CHANGELOG.md | 1 + .../SemanticHighlightingConstants.java | 1 + .../neuralsearch/plugin/NeuralSearch.java | 5 +- ...nticHighlighterQueryEnricherProcessor.java | 181 ++++++++ .../stats/events/EventStatName.java | 7 + ...icHighlighterQueryEnricherProcessorIT.java | 203 +++++++++ ...ighlighterQueryEnricherProcessorTests.java | 415 ++++++++++++++++++ 7 files changed, 812 insertions(+), 1 deletion(-) create mode 100644 src/main/java/org/opensearch/neuralsearch/processor/SemanticHighlighterQueryEnricherProcessor.java create mode 100644 src/test/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterQueryEnricherProcessorIT.java create mode 100644 src/test/java/org/opensearch/neuralsearch/processor/SemanticHighlighterQueryEnricherProcessorTests.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 3aa4f6c2b..19aed7891 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ## [Unreleased 3.x](https://github.com/opensearch-project/neural-search/compare/main...HEAD) ### Features +* [SemanticHighlighter] add `semantic_highlighter_query_enricher` a new processor similar to `neural_query_enricher` that allows to enrich the semantic highlighter model id from a search pipeline ([#1916](https://github.com/opensearch-project/neural-search/pull/1917)) ### Enhancements diff --git a/src/main/java/org/opensearch/neuralsearch/highlight/SemanticHighlightingConstants.java b/src/main/java/org/opensearch/neuralsearch/highlight/SemanticHighlightingConstants.java index 5c7b25780..001f7ef15 100644 --- a/src/main/java/org/opensearch/neuralsearch/highlight/SemanticHighlightingConstants.java +++ b/src/main/java/org/opensearch/neuralsearch/highlight/SemanticHighlightingConstants.java @@ -11,6 +11,7 @@ public final class SemanticHighlightingConstants { // System-generated factory and processor types public static final String SYSTEM_FACTORY_TYPE = "semantic-highlighter"; public static final String PROCESSOR_TYPE = "semantic_highlighting"; + public static final String QUERY_ENRICHER_TYPE = "semantic_highlighter_query_enricher"; // Default processor tags and descriptions public static final String DEFAULT_PROCESSOR_TAG = "semantic-highlighter"; diff --git a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java index e386d77bc..fd64494b5 100644 --- a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java +++ b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java @@ -32,6 +32,7 @@ import com.google.common.collect.ImmutableList; import lombok.extern.log4j.Log4j2; import org.opensearch.action.ActionRequest; +import org.opensearch.neuralsearch.processor.SemanticHighlighterQueryEnricherProcessor; import org.opensearch.neuralsearch.query.NeuralQueryBuilder; import org.opensearch.neuralsearch.query.HybridQueryBuilder; import org.opensearch.neuralsearch.query.NeuralSparseQueryBuilder; @@ -389,7 +390,9 @@ public Map fieldDefaultIdMap; + + private SemanticHighlighterQueryEnricherProcessor( + String tag, + String description, + boolean ignoreFailure, + @Nullable String modelId, + @Nullable Map fieldDefaultIdMap + ) { + super(tag, description, ignoreFailure); + this.modelId = modelId; + this.fieldDefaultIdMap = fieldDefaultIdMap; + } + + @Override + public SearchRequest processRequest(SearchRequest searchRequest) { + EventStatsManager.increment(EventStatName.SEMANTIC_HIGHLIGHTING_QUERY_ENRICHER_EXECUTIONS); + Optional source = Optional.ofNullable(searchRequest.source()); + source.map(SearchSourceBuilder::highlighter).ifPresent(this::enrichHighlight); + source.map(SearchSourceBuilder::query).ifPresent(qb -> qb.visit(new NestedQueryHighlightVisitor())); + // NOTE: we explicitly do not enrich TopHitsAggregationBuilder highlighters because it is not useful yet — the batch path ignores + // aggregations entirely (HighlightConfigResolver only walks source.highlighter() and inner_hits, HighlightContextBuilder only reads + // response.getHits()), so under ext.semantic_highlighting_batch the model_id would be set but never used, and highlights would go + // missing silently. Needs the highlighting feature to support aggregations first. + return searchRequest; + } + + private void enrichHighlight(HighlightBuilder hlBuilder) { + Map globalOptions = hlBuilder.options(); + boolean userSuppliedGlobalModelId = globalOptions != null && globalOptions.containsKey(SemanticHighlightingConstants.MODEL_ID); + if (userSuppliedGlobalModelId) { + // if the user provided a global model_id there's no need to enrich anything. + return; + } + + boolean globalIsSemantic = false; + if (SemanticHighlightingConstants.HIGHLIGHTER_TYPE.equals(hlBuilder.highlighterType())) { + globalIsSemantic = true; + if (modelId != null) { + hlBuilder.options(enrichWithModelId(globalOptions, modelId)); + } + } + for (HighlightBuilder.Field field : Optional.ofNullable(hlBuilder.fields()).orElseGet(Collections::emptyList)) { + // Enrich if either: + // - the global type is semantic and the field specific type is unset + // - the field specific type is set to semantic + if ((globalIsSemantic && field.highlighterType() == null) + || SemanticHighlightingConstants.HIGHLIGHTER_TYPE.equals(field.highlighterType())) { + String fieldModelId = (String) Optional.ofNullable(this.fieldDefaultIdMap) + .orElseGet(Collections::emptyMap) + .getOrDefault(field.name(), modelId); + if (fieldModelId != null) { + field.options(enrichWithModelId(field.options(), fieldModelId)); + } + // else: no default model_id and no per-field override for this field, nothing to enrich + } + } + } + + private Map enrichWithModelId(@Nullable Map options, String modelId) { + if (options != null && options.containsKey(SemanticHighlightingConstants.MODEL_ID)) { + return options; + } + Map enrichedOptions = options != null ? new HashMap<>(options) : new HashMap<>(); + enrichedOptions.put(SemanticHighlightingConstants.MODEL_ID, modelId); + return enrichedOptions; + } + + private class NestedQueryHighlightVisitor implements QueryBuilderVisitor { + @Override + public void accept(QueryBuilder qb) { + if (qb instanceof NestedQueryBuilder nested) { + Optional.ofNullable(nested.innerHit()) + .map(InnerHitBuilder::getHighlightBuilder) + .ifPresent(SemanticHighlighterQueryEnricherProcessor.this::enrichHighlight); + } + } + + @Override + public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { + return this; + } + } + + @Override + public String getType() { + return SemanticHighlightingConstants.QUERY_ENRICHER_TYPE; + } + + public static class Factory implements Processor.Factory { + private static final String DEFAULT_MODEL_ID = "default_model_id"; + private static final String SEMANTIC_HIGHLIGHTER_FIELD_DEFAULT_ID = "semantic_highlighter_field_default_id"; + + /** + * Create the processor object. + * + * @return {@link SemanticHighlighterQueryEnricherProcessor} + */ + @Override + public SemanticHighlighterQueryEnricherProcessor create( + Map> processorFactories, + String tag, + String description, + boolean ignoreFailure, + Map config, + PipelineContext pipelineContext + ) throws IllegalArgumentException { + String modelId = readOptionalStringProperty(TYPE, tag, config, DEFAULT_MODEL_ID); + Map fieldMap = ConfigurationUtils.readOptionalMap(TYPE, tag, config, SEMANTIC_HIGHLIGHTER_FIELD_DEFAULT_ID); + + if (modelId == null && fieldMap == null) { + throw new IllegalArgumentException("[default_model_id] or [semantic_highlighter_field_default_id] should be provided"); + } + + if (fieldMap != null) { + List nonStringFields = fieldMap.entrySet() + .stream() + .filter(en -> !(en.getValue() instanceof String)) + .map(Map.Entry::getKey) + .toList(); + if (!nonStringFields.isEmpty()) { + throw new IllegalArgumentException( + "Invalid type in [semantic_highlighter_field_default_id]: value for [" + + String.join(", ", nonStringFields) + + "] must be a model_id of type string" + ); + } + } + + return new SemanticHighlighterQueryEnricherProcessor( + tag, + description, + ignoreFailure, + modelId, + fieldMap != null ? Collections.unmodifiableMap(fieldMap) : null + ); + } + } +} diff --git a/src/main/java/org/opensearch/neuralsearch/stats/events/EventStatName.java b/src/main/java/org/opensearch/neuralsearch/stats/events/EventStatName.java index 1abd36516..daf075a90 100644 --- a/src/main/java/org/opensearch/neuralsearch/stats/events/EventStatName.java +++ b/src/main/java/org/opensearch/neuralsearch/stats/events/EventStatName.java @@ -250,6 +250,13 @@ public enum EventStatName implements StatName { EventStatType.TIMESTAMPED_EVENT_COUNTER, Version.V_3_1_0 ), + /** Tracks executions of the semantic highlighter query enricher processor */ + SEMANTIC_HIGHLIGHTING_QUERY_ENRICHER_EXECUTIONS( + "semantic_highlighting_query_enricher_executions", + "processors.search", + EventStatType.TIMESTAMPED_EVENT_COUNTER, + Version.V_3_8_0 + ), /** Tracks executions of the ML reranking processor */ RERANK_ML_PROCESSOR_EXECUTIONS("rerank_ml_executions", "processors.search", EventStatType.TIMESTAMPED_EVENT_COUNTER, Version.V_3_1_0), diff --git a/src/test/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterQueryEnricherProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterQueryEnricherProcessorIT.java new file mode 100644 index 000000000..3a92a54ff --- /dev/null +++ b/src/test/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterQueryEnricherProcessorIT.java @@ -0,0 +1,203 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.neuralsearch.highlight; + +import java.util.ArrayList; +import java.util.Map; + +import org.apache.hc.core5.http.io.entity.EntityUtils; +import org.junit.Before; +import org.opensearch.client.Request; +import org.opensearch.client.Response; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.common.xcontent.XContentHelper; +import org.opensearch.common.xcontent.XContentType; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.neuralsearch.stats.events.EventStatName; + +import lombok.SneakyThrows; +import lombok.extern.log4j.Log4j2; + +/** + * Integration tests for the semantic highlighter query enricher search request processor. + * + *

The processor lets an operator configure the highlighter {@code model_id} once in a search + * pipeline instead of repeating it in every query body, so each test issues a query that omits + * {@code model_id} and asserts that semantic highlighting still resolves a model. + */ +@Log4j2 +public class SemanticHighlighterQueryEnricherProcessorIT extends BaseSemanticHighlightingIT { + + private static final String TEST_INDEX = "test-semantic-highlight-enricher-index"; + private static final String SEARCH_PIPELINE = "semantic-highlight-enricher-pipeline"; + private static final String QUERY_TEXT = "treatments for neurodegenerative diseases"; + private static final String INVALID_MODEL_ID = "this-model-does-not-exist"; + + private String highlightModelId; + + @Before + @SneakyThrows + public void setUp() { + super.setUp(); + highlightModelId = prepareSentenceHighlightingModel(); + log.info("Prepared local highlighting model, model ID: {}", highlightModelId); + prepareHighlightingIndex(TEST_INDEX); + indexTestDocuments(TEST_INDEX); + } + + /** + * The query omits model_id entirely, the pipeline's default_model_id supplies it. + * Also asserts the processor's own event stat is recorded. + */ + @SneakyThrows + public void testQueryEnricher_whenModelIdOmitted_thenDefaultModelIdApplied() { + enableStats(); + createEnricherPipeline(defaultModelIdConfig(highlightModelId)); + + Map searchResponse = search(semanticHighlightQuery(null)); + + assertSemanticHighlighting(searchResponse, TEST_FIELD, "treatments"); + + Map stats = parseAggregatedNodeStatsResponse(executeNeuralStatRequest(new ArrayList<>(), new ArrayList<>())); + int enricherExecutions = (int) getNestedValue(stats, EventStatName.SEMANTIC_HIGHLIGHTING_QUERY_ENRICHER_EXECUTIONS); + assertEquals("Query enricher should have run exactly once", 1, enricherExecutions); + } + + /** A per field override supplies the model_id for the highlighted field. */ + @SneakyThrows + public void testQueryEnricher_whenFieldDefaultIdConfigured_thenFieldModelIdApplied() { + createEnricherPipeline(fieldDefaultIdConfig(TEST_FIELD, highlightModelId)); + + Map searchResponse = search(semanticHighlightQuery(null)); + + assertSemanticHighlighting(searchResponse, TEST_FIELD, "treatments"); + } + + /** + * A model_id in the query body must win over the pipeline default. The pipeline is configured + * with a model that does not exist, so highlighting can only succeed if the query value was kept. + */ + @SneakyThrows + public void testQueryEnricher_whenQuerySuppliesModelId_thenPipelineDefaultNotApplied() { + createEnricherPipeline(defaultModelIdConfig(INVALID_MODEL_ID)); + + Map searchResponse = search(semanticHighlightQuery(highlightModelId)); + + assertSemanticHighlighting(searchResponse, TEST_FIELD, "treatments"); + } + + /** + * A non semantic highlighter must be left alone. The pipeline default points at a model that + * does not exist, so the request only succeeds if the processor did not enrich it. + */ + @SneakyThrows + public void testQueryEnricher_whenHighlighterIsNotSemantic_thenNotEnriched() { + createEnricherPipeline(defaultModelIdConfig(INVALID_MODEL_ID)); + + XContentBuilder searchBody = XContentFactory.jsonBuilder() + .startObject() + .field("size", 2) + .startObject("query") + .startObject("match") + .field(TEST_FIELD, QUERY_TEXT) + .endObject() + .endObject() + .startObject("highlight") + .startObject("fields") + .startObject(TEST_FIELD) + .field("type", "unified") + .endObject() + .endObject() + .endObject() + .endObject(); + + Map searchResponse = search(searchBody); + + // Unified highlighting still produces fragments, the point is that no model was resolved + assertSemanticHighlighting(searchResponse, TEST_FIELD, "treatments"); + } + + // --------------------------------------------------------------------- + // helpers + // --------------------------------------------------------------------- + + /** + * Builds a match query highlighted with {@code type: semantic}, optionally carrying an + * explicit {@code model_id} in the highlight options. + */ + @SneakyThrows + private XContentBuilder semanticHighlightQuery(String modelId) { + XContentBuilder builder = XContentFactory.jsonBuilder() + .startObject() + .field("size", 2) + .startObject("query") + .startObject("match") + .field(TEST_FIELD, QUERY_TEXT) + .endObject() + .endObject() + .startObject("highlight") + .startObject("fields") + .startObject(TEST_FIELD) + .field("type", SemanticHighlightingConstants.HIGHLIGHTER_TYPE) + .endObject() + .endObject(); + if (modelId != null) { + builder.startObject("options").field(SemanticHighlightingConstants.MODEL_ID, modelId).endObject(); + } + return builder.endObject().endObject(); + } + + @SneakyThrows + private String defaultModelIdConfig(String modelId) { + return XContentFactory.jsonBuilder() + .startObject() + .startArray("request_processors") + .startObject() + .startObject(SemanticHighlightingConstants.QUERY_ENRICHER_TYPE) + .field("default_model_id", modelId) + .endObject() + .endObject() + .endArray() + .endObject() + .toString(); + } + + @SneakyThrows + private String fieldDefaultIdConfig(String fieldName, String modelId) { + return XContentFactory.jsonBuilder() + .startObject() + .startArray("request_processors") + .startObject() + .startObject(SemanticHighlightingConstants.QUERY_ENRICHER_TYPE) + .startObject("semantic_highlighter_field_default_id") + .field(fieldName, modelId) + .endObject() + .endObject() + .endObject() + .endArray() + .endObject() + .toString(); + } + + /** Creates the search pipeline and makes it the index default so plain searches pick it up. */ + @SneakyThrows + private void createEnricherPipeline(String pipelineConfig) { + Request request = new Request("PUT", "/_search/pipeline/" + SEARCH_PIPELINE); + request.setJsonEntity(pipelineConfig); + Response response = client().performRequest(request); + assertEquals(200, response.getStatusLine().getStatusCode()); + updateIndexSettings(TEST_INDEX, Settings.builder().put("index.search.default_pipeline", SEARCH_PIPELINE)); + } + + @SneakyThrows + private Map search(XContentBuilder searchBody) { + Request request = new Request("POST", "/" + TEST_INDEX + "/_search"); + request.setJsonEntity(searchBody.toString()); + Response response = client().performRequest(request); + assertEquals(200, response.getStatusLine().getStatusCode()); + return XContentHelper.convertToMap(XContentType.JSON.xContent(), EntityUtils.toString(response.getEntity()), false); + } +} diff --git a/src/test/java/org/opensearch/neuralsearch/processor/SemanticHighlighterQueryEnricherProcessorTests.java b/src/test/java/org/opensearch/neuralsearch/processor/SemanticHighlighterQueryEnricherProcessorTests.java new file mode 100644 index 000000000..7c7698502 --- /dev/null +++ b/src/test/java/org/opensearch/neuralsearch/processor/SemanticHighlighterQueryEnricherProcessorTests.java @@ -0,0 +1,415 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.neuralsearch.processor; + +import org.apache.lucene.search.join.ScoreMode; +import org.opensearch.OpenSearchParseException; +import org.opensearch.action.search.SearchRequest; +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.InnerHitBuilder; +import org.opensearch.index.query.MatchQueryBuilder; +import org.opensearch.index.query.NestedQueryBuilder; +import org.opensearch.neuralsearch.highlight.SemanticHighlightingConstants; +import org.opensearch.neuralsearch.util.TestUtils; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder; +import org.opensearch.test.OpenSearchTestCase; +import org.junit.Before; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +public class SemanticHighlighterQueryEnricherProcessorTests extends OpenSearchTestCase { + + private static final String DEFAULT_MODEL_ID = "default-model-id"; + private static final String FIELD_MODEL_ID = "field-model-id"; + private static final String CONTENT_FIELD = "content"; + private static final String TITLE_FIELD = "title"; + private static final String SEMANTIC = SemanticHighlightingConstants.HIGHLIGHTER_TYPE; + private static final String MODEL_ID = SemanticHighlightingConstants.MODEL_ID; + + @Before + public void setup() { + TestUtils.initializeEventStatsManager(); + } + + // --------------------------------------------------------------------- + // Factory + // --------------------------------------------------------------------- + + public void testFactory_whenNoModelIdAndNoFieldMap_thenThrowException() { + SemanticHighlighterQueryEnricherProcessor.Factory factory = new SemanticHighlighterQueryEnricherProcessor.Factory(); + IllegalArgumentException e = expectThrows( + IllegalArgumentException.class, + () -> factory.create(Collections.emptyMap(), null, null, false, new HashMap<>(), null) + ); + assertTrue(e.getMessage(), e.getMessage().contains("semantic_highlighter_field_default_id")); + } + + public void testFactory_whenOnlyDefaultModelId_thenSuccess() throws Exception { + SemanticHighlighterQueryEnricherProcessor processor = createProcessor(DEFAULT_MODEL_ID, null); + assertEquals(DEFAULT_MODEL_ID, processor.getModelId()); + assertNull(processor.getFieldDefaultIdMap()); + } + + public void testFactory_whenOnlyFieldMap_thenSuccess() throws Exception { + SemanticHighlighterQueryEnricherProcessor processor = createProcessor(null, Map.of(CONTENT_FIELD, FIELD_MODEL_ID)); + assertNull(processor.getModelId()); + assertEquals(FIELD_MODEL_ID, processor.getFieldDefaultIdMap().get(CONTENT_FIELD)); + } + + public void testFactory_whenModelIdIsNotString_thenThrowException() { + SemanticHighlighterQueryEnricherProcessor.Factory factory = new SemanticHighlighterQueryEnricherProcessor.Factory(); + Map config = new HashMap<>(); + config.put("default_model_id", 12345L); + expectThrows(OpenSearchParseException.class, () -> factory.create(Collections.emptyMap(), null, null, false, config, null)); + } + + public void testFactory_whenFieldMapValueIsNotString_thenThrowException() { + SemanticHighlighterQueryEnricherProcessor.Factory factory = new SemanticHighlighterQueryEnricherProcessor.Factory(); + Map config = new HashMap<>(); + config.put("semantic_highlighter_field_default_id", Map.of(CONTENT_FIELD, 12345L)); + IllegalArgumentException e = expectThrows( + IllegalArgumentException.class, + () -> factory.create(Collections.emptyMap(), null, null, false, config, null) + ); + // The offending field must be named so the operator can find it in the pipeline config + assertTrue(e.getMessage(), e.getMessage().contains(CONTENT_FIELD)); + } + + public void testType() throws Exception { + SemanticHighlighterQueryEnricherProcessor processor = createProcessor(DEFAULT_MODEL_ID, null); + assertEquals(SemanticHighlightingConstants.QUERY_ENRICHER_TYPE, processor.getType()); + } + + // --------------------------------------------------------------------- + // processRequest: degenerate requests + // --------------------------------------------------------------------- + + /** A transport level request built without a body has a null source. */ + public void testProcessRequest_whenSourceIsNull_thenNoOp() throws Exception { + SearchRequest searchRequest = new SearchRequest(); + assertNull(searchRequest.source()); + SemanticHighlighterQueryEnricherProcessor processor = createProcessor(DEFAULT_MODEL_ID, null); + assertSame(searchRequest, processor.processRequest(searchRequest)); + } + + public void testProcessRequest_whenEmptySource_thenNoOp() throws Exception { + SearchRequest searchRequest = requestOf(new SearchSourceBuilder()); + SemanticHighlighterQueryEnricherProcessor processor = createProcessor(DEFAULT_MODEL_ID, null); + processor.processRequest(searchRequest); + assertNull(searchRequest.source().highlighter()); + assertNull(searchRequest.source().query()); + } + + // --------------------------------------------------------------------- + // processRequest: top level highlighter + // --------------------------------------------------------------------- + + /** Global type: semantic with no options at all, the common case the processor exists for. */ + public void testProcessRequest_whenGlobalSemanticAndNoOptions_thenGlobalAndFieldEnriched() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + hlBuilder.highlighterType(SEMANTIC); + hlBuilder.field(new HighlightBuilder.Field(CONTENT_FIELD)); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, null); + + assertEquals(DEFAULT_MODEL_ID, enriched.options().get(MODEL_ID)); + assertEquals(DEFAULT_MODEL_ID, enriched.fields().get(0).options().get(MODEL_ID)); + } + + /** Field level type: semantic with no global type, and no options. */ + public void testProcessRequest_whenFieldSemanticAndNoOptions_thenFieldEnriched() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + HighlightBuilder.Field field = new HighlightBuilder.Field(CONTENT_FIELD); + field.highlighterType(SEMANTIC); + hlBuilder.field(field); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, null); + + assertEquals(DEFAULT_MODEL_ID, enriched.fields().get(0).options().get(MODEL_ID)); + // The global highlighter is not semantic, it must not be touched + assertNull(enriched.options()); + } + + /** A model_id supplied in the query always wins over the pipeline configuration. */ + public void testProcessRequest_whenModelIdAlreadySet_thenPreserved() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + HighlightBuilder.Field field = new HighlightBuilder.Field(CONTENT_FIELD); + field.highlighterType(SEMANTIC); + Map options = new HashMap<>(); + options.put(MODEL_ID, "user-supplied-model-id"); + field.options(options); + hlBuilder.field(field); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, null); + + assertEquals("user-supplied-model-id", enriched.fields().get(0).options().get(MODEL_ID)); + } + + /** + * A model_id declared once in the global highlight options covers every field. Field level + * options win when merged with the global ones, so enriching the field would silently + * override the model the user asked for. + */ + public void testProcessRequest_whenGlobalOptionsCarryModelId_thenFieldNotOverridden() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + Map globalOptions = new HashMap<>(); + globalOptions.put(MODEL_ID, "user-supplied-model-id"); + hlBuilder.options(globalOptions); + HighlightBuilder.Field field = new HighlightBuilder.Field(CONTENT_FIELD); + field.highlighterType(SEMANTIC); + hlBuilder.field(field); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, null); + + assertEquals("user-supplied-model-id", enriched.options().get(MODEL_ID)); + assertNull("The pipeline default must not be pushed onto the field", enriched.fields().get(0).options()); + } + + /** A per field override from the pipeline must not beat a model_id supplied in the query either. */ + public void testProcessRequest_whenGlobalOptionsCarryModelIdAndFieldMapConfigured_thenFieldNotOverridden() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + Map globalOptions = new HashMap<>(); + globalOptions.put(MODEL_ID, "user-supplied-model-id"); + hlBuilder.options(globalOptions); + HighlightBuilder.Field field = new HighlightBuilder.Field(CONTENT_FIELD); + field.highlighterType(SEMANTIC); + hlBuilder.field(field); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, Map.of(CONTENT_FIELD, FIELD_MODEL_ID)); + + assertEquals("user-supplied-model-id", enriched.options().get(MODEL_ID)); + assertNull("The pipeline field override must not beat the query's global model_id", enriched.fields().get(0).options()); + } + + /** Global type semantic with a user supplied global model_id: nothing to enrich anywhere. */ + public void testProcessRequest_whenGlobalSemanticAndGlobalModelIdSet_thenUntouched() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + hlBuilder.highlighterType(SEMANTIC); + Map globalOptions = new HashMap<>(); + globalOptions.put(MODEL_ID, "user-supplied-model-id"); + hlBuilder.options(globalOptions); + hlBuilder.field(new HighlightBuilder.Field(CONTENT_FIELD)); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, null); + + assertEquals("user-supplied-model-id", enriched.options().get(MODEL_ID)); + assertNull(enriched.fields().get(0).options()); + } + + /** + * Global options that carry no model_id must not suppress enrichment. The guard keys on the + * presence of model_id, not on the presence of a global options block. + */ + public void testProcessRequest_whenGlobalOptionsWithoutModelId_thenStillEnriched() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + Map globalOptions = new HashMap<>(); + globalOptions.put("random_option", "some-value"); + hlBuilder.options(globalOptions); + HighlightBuilder.Field field = new HighlightBuilder.Field(CONTENT_FIELD); + field.highlighterType(SEMANTIC); + hlBuilder.field(field); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, null); + + assertEquals(DEFAULT_MODEL_ID, enriched.fields().get(0).options().get(MODEL_ID)); + // The unrelated global option is left alone + assertEquals("some-value", enriched.options().get("random_option")); + } + + /** Enrichment must not drop options the user already set. */ + public void testProcessRequest_whenOtherOptionsSet_thenPreservedAlongsideModelId() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + HighlightBuilder.Field field = new HighlightBuilder.Field(CONTENT_FIELD); + field.highlighterType(SEMANTIC); + Map options = new HashMap<>(); + options.put("random_option", true); + field.options(options); + hlBuilder.field(field); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, null); + + Map enrichedOptions = enriched.fields().getFirst().options(); + assertEquals(DEFAULT_MODEL_ID, enrichedOptions.get(MODEL_ID)); + assertEquals(true, enrichedOptions.get("random_option")); + } + + public void testProcessRequest_whenFieldIsNotSemantic_thenNotEnriched() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + HighlightBuilder.Field field = new HighlightBuilder.Field(CONTENT_FIELD); + field.highlighterType("plain"); + hlBuilder.field(field); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, null); + + assertNull(enriched.fields().get(0).options()); + } + + /** A field that opts out of the semantic global type must not be enriched, the global one still is. */ + public void testProcessRequest_whenGlobalSemanticButFieldOverridesType_thenOnlyGlobalEnriched() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + hlBuilder.highlighterType(SEMANTIC); + HighlightBuilder.Field field = new HighlightBuilder.Field(CONTENT_FIELD); + field.highlighterType("plain"); + hlBuilder.field(field); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, null); + + assertEquals(DEFAULT_MODEL_ID, enriched.options().get(MODEL_ID)); + assertNull(enriched.fields().get(0).options()); + } + + public void testProcessRequest_whenNoFieldsDeclared_thenGlobalStillEnriched() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + hlBuilder.highlighterType(SEMANTIC); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, null); + + assertEquals(DEFAULT_MODEL_ID, enriched.options().get(MODEL_ID)); + assertTrue(enriched.fields().isEmpty()); + } + + // --------------------------------------------------------------------- + // processRequest: per field model id map + // --------------------------------------------------------------------- + + public void testProcessRequest_whenFieldMapConfigured_thenFieldSpecificModelIdWins() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + HighlightBuilder.Field content = new HighlightBuilder.Field(CONTENT_FIELD); + content.highlighterType(SEMANTIC); + hlBuilder.field(content); + HighlightBuilder.Field title = new HighlightBuilder.Field(TITLE_FIELD); + title.highlighterType(SEMANTIC); + hlBuilder.field(title); + + HighlightBuilder enriched = process(hlBuilder, DEFAULT_MODEL_ID, Map.of(CONTENT_FIELD, FIELD_MODEL_ID)); + + assertEquals(FIELD_MODEL_ID, enriched.fields().get(0).options().get(MODEL_ID)); + // No override for title, it falls back to the default model id + assertEquals(DEFAULT_MODEL_ID, enriched.fields().get(1).options().get(MODEL_ID)); + } + + /** Only per field overrides configured, a field absent from the map has nothing to enrich with. */ + public void testProcessRequest_whenOnlyFieldMapAndFieldAbsent_thenNotEnriched() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + HighlightBuilder.Field title = new HighlightBuilder.Field(TITLE_FIELD); + title.highlighterType(SEMANTIC); + hlBuilder.field(title); + + HighlightBuilder enriched = process(hlBuilder, null, Map.of(CONTENT_FIELD, FIELD_MODEL_ID)); + + assertNull(enriched.fields().get(0).options()); + } + + /** Only per field overrides configured, the global highlighter has no default to fall back on. */ + public void testProcessRequest_whenOnlyFieldMapAndGlobalSemantic_thenGlobalNotEnriched() throws Exception { + HighlightBuilder hlBuilder = new HighlightBuilder(); + hlBuilder.highlighterType(SEMANTIC); + hlBuilder.field(new HighlightBuilder.Field(CONTENT_FIELD)); + + HighlightBuilder enriched = process(hlBuilder, null, Map.of(CONTENT_FIELD, FIELD_MODEL_ID)); + + assertNull(enriched.options()); + assertEquals(FIELD_MODEL_ID, enriched.fields().get(0).options().get(MODEL_ID)); + } + + // --------------------------------------------------------------------- + // processRequest: nested inner_hits + // --------------------------------------------------------------------- + + public void testProcessRequest_whenNestedInnerHitsHighlight_thenEnriched() throws Exception { + NestedQueryBuilder nested = nestedWithSemanticInnerHit("chunks", "chunks.text"); + + SearchRequest searchRequest = requestOf(new SearchSourceBuilder().query(nested)); + createProcessor(DEFAULT_MODEL_ID, null).processRequest(searchRequest); + + assertEquals(DEFAULT_MODEL_ID, innerHitOptions(nested).get(MODEL_ID)); + } + + /** The realistic shape: a nested clause wrapped in a bool, reached through getChildVisitor. */ + public void testProcessRequest_whenNestedInsideBool_thenEnriched() throws Exception { + NestedQueryBuilder nested = nestedWithSemanticInnerHit("chunks", "chunks.text"); + + SearchRequest searchRequest = requestOf(new SearchSourceBuilder().query(new BoolQueryBuilder().must(nested))); + createProcessor(DEFAULT_MODEL_ID, null).processRequest(searchRequest); + + assertEquals(DEFAULT_MODEL_ID, innerHitOptions(nested).get(MODEL_ID)); + } + + public void testProcessRequest_whenNestedWithoutInnerHits_thenNoOp() throws Exception { + NestedQueryBuilder nested = new NestedQueryBuilder("chunks", new MatchQueryBuilder("chunks.text", "x"), ScoreMode.Avg); + assertNull(nested.innerHit()); + + SearchRequest searchRequest = requestOf(new SearchSourceBuilder().query(nested)); + SemanticHighlighterQueryEnricherProcessor processor = createProcessor(DEFAULT_MODEL_ID, null); + assertSame(searchRequest, processor.processRequest(searchRequest)); + } + + public void testProcessRequest_whenInnerHitsWithoutHighlight_thenNoOp() throws Exception { + InnerHitBuilder innerHit = new InnerHitBuilder(); + assertNull(innerHit.getHighlightBuilder()); + NestedQueryBuilder nested = new NestedQueryBuilder("chunks", new MatchQueryBuilder("chunks.text", "x"), ScoreMode.Avg).innerHit( + innerHit + ); + + SearchRequest searchRequest = requestOf(new SearchSourceBuilder().query(nested)); + SemanticHighlighterQueryEnricherProcessor processor = createProcessor(DEFAULT_MODEL_ID, null); + assertSame(searchRequest, processor.processRequest(searchRequest)); + } + + public void testProcessRequest_whenNonNestedQuery_thenNoOp() throws Exception { + SearchRequest searchRequest = requestOf(new SearchSourceBuilder().query(new MatchQueryBuilder(CONTENT_FIELD, "x"))); + SemanticHighlighterQueryEnricherProcessor processor = createProcessor(DEFAULT_MODEL_ID, null); + assertSame(searchRequest, processor.processRequest(searchRequest)); + } + + // --------------------------------------------------------------------- + // helpers + // --------------------------------------------------------------------- + + /** Runs the processor over a request carrying only the given highlighter, and returns it for assertions. */ + private HighlightBuilder process(HighlightBuilder hlBuilder, String modelId, Map fieldMap) throws Exception { + SearchRequest searchRequest = requestOf(new SearchSourceBuilder().highlighter(hlBuilder)); + createProcessor(modelId, fieldMap).processRequest(searchRequest); + return searchRequest.source().highlighter(); + } + + /** A nested query whose inner_hits declare a single {@code type: semantic} field. */ + private NestedQueryBuilder nestedWithSemanticInnerHit(String path, String fieldName) { + HighlightBuilder innerHighlight = new HighlightBuilder(); + HighlightBuilder.Field field = new HighlightBuilder.Field(fieldName); + field.highlighterType(SEMANTIC); + innerHighlight.field(field); + + InnerHitBuilder innerHit = new InnerHitBuilder(); + innerHit.setHighlightBuilder(innerHighlight); + return new NestedQueryBuilder(path, new MatchQueryBuilder(fieldName, "x"), ScoreMode.Avg).innerHit(innerHit); + } + + private static Map innerHitOptions(NestedQueryBuilder nested) { + return nested.innerHit().getHighlightBuilder().fields().get(0).options(); + } + + private SearchRequest requestOf(SearchSourceBuilder source) { + SearchRequest searchRequest = new SearchRequest("index"); + searchRequest.source(source); + return searchRequest; + } + + private SemanticHighlighterQueryEnricherProcessor createProcessor(String modelId, Map fieldMap) throws Exception { + SemanticHighlighterQueryEnricherProcessor.Factory factory = new SemanticHighlighterQueryEnricherProcessor.Factory(); + // readOptionalMap/readOptionalStringProperty remove consumed keys, so the config map must be mutable + Map config = new HashMap<>(); + if (modelId != null) { + config.put("default_model_id", modelId); + } + if (fieldMap != null) { + config.put("semantic_highlighter_field_default_id", new HashMap(fieldMap)); + } + return factory.create(Collections.emptyMap(), null, null, false, config, null); + } +}