Skip to content

Commit 584cbea

Browse files
nomoanomoa
committed
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 <dcausse@wikimedia.org>
1 parent 5f9fba8 commit 584cbea

7 files changed

Lines changed: 812 additions & 1 deletion

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
66
## [Unreleased 3.x](https://github.com/opensearch-project/neural-search/compare/main...HEAD)
77

88
### Features
9+
* [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))
910

1011
### Enhancements
1112

‎src/main/java/org/opensearch/neuralsearch/highlight/SemanticHighlightingConstants.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ public final class SemanticHighlightingConstants {
1111
// System-generated factory and processor types
1212
public static final String SYSTEM_FACTORY_TYPE = "semantic-highlighter";
1313
public static final String PROCESSOR_TYPE = "semantic_highlighting";
14+
public static final String QUERY_ENRICHER_TYPE = "semantic_highlighter_query_enricher";
1415

1516
// Default processor tags and descriptions
1617
public static final String DEFAULT_PROCESSOR_TAG = "semantic-highlighter";

‎src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import com.google.common.collect.ImmutableList;
3333
import lombok.extern.log4j.Log4j2;
3434
import org.opensearch.action.ActionRequest;
35+
import org.opensearch.neuralsearch.processor.SemanticHighlighterQueryEnricherProcessor;
3536
import org.opensearch.neuralsearch.query.NeuralQueryBuilder;
3637
import org.opensearch.neuralsearch.query.HybridQueryBuilder;
3738
import org.opensearch.neuralsearch.query.NeuralSparseQueryBuilder;
@@ -389,7 +390,9 @@ public Map<String, org.opensearch.search.pipeline.Processor.Factory<SearchReques
389390
NeuralSparseTwoPhaseProcessor.TYPE,
390391
new NeuralSparseTwoPhaseProcessor.Factory(),
391392
AgenticQueryTranslatorProcessor.TYPE,
392-
new AgenticQueryTranslatorProcessor.Factory(clientAccessor, xContentRegistry, settingsAccessor)
393+
new AgenticQueryTranslatorProcessor.Factory(clientAccessor, xContentRegistry, settingsAccessor),
394+
SemanticHighlighterQueryEnricherProcessor.TYPE,
395+
new SemanticHighlighterQueryEnricherProcessor.Factory()
393396
);
394397
}
395398

‎src/main/java/org/opensearch/neuralsearch/processor/SemanticHighlighterQueryEnricherProcessor.java‎

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package org.opensearch.neuralsearch.processor;
6+
7+
import lombok.Getter;
8+
9+
import org.apache.lucene.search.BooleanClause;
10+
import org.opensearch.action.search.SearchRequest;
11+
import org.opensearch.common.Nullable;
12+
import org.opensearch.index.query.InnerHitBuilder;
13+
import org.opensearch.index.query.NestedQueryBuilder;
14+
import org.opensearch.index.query.QueryBuilder;
15+
import org.opensearch.index.query.QueryBuilderVisitor;
16+
import org.opensearch.ingest.ConfigurationUtils;
17+
import org.opensearch.neuralsearch.highlight.SemanticHighlightingConstants;
18+
import org.opensearch.neuralsearch.stats.events.EventStatName;
19+
import org.opensearch.neuralsearch.stats.events.EventStatsManager;
20+
import org.opensearch.search.builder.SearchSourceBuilder;
21+
import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder;
22+
import org.opensearch.search.pipeline.AbstractProcessor;
23+
import org.opensearch.search.pipeline.Processor;
24+
import org.opensearch.search.pipeline.SearchRequestProcessor;
25+
26+
import java.util.Collections;
27+
import java.util.HashMap;
28+
import java.util.List;
29+
import java.util.Map;
30+
import java.util.Optional;
31+
32+
import static org.opensearch.ingest.ConfigurationUtils.readOptionalStringProperty;
33+
34+
/**
35+
* Query enricher that will populate the model_id option when the semantic highlighter is used and no model_id is
36+
* specified in the search query body.
37+
*/
38+
@Getter
39+
public class SemanticHighlighterQueryEnricherProcessor extends AbstractProcessor implements SearchRequestProcessor {
40+
public static final String TYPE = SemanticHighlightingConstants.QUERY_ENRICHER_TYPE;
41+
42+
private final String modelId;
43+
private final Map<String, Object> fieldDefaultIdMap;
44+
45+
private SemanticHighlighterQueryEnricherProcessor(
46+
String tag,
47+
String description,
48+
boolean ignoreFailure,
49+
@Nullable String modelId,
50+
@Nullable Map<String, Object> fieldDefaultIdMap
51+
) {
52+
super(tag, description, ignoreFailure);
53+
this.modelId = modelId;
54+
this.fieldDefaultIdMap = fieldDefaultIdMap;
55+
}
56+
57+
@Override
58+
public SearchRequest processRequest(SearchRequest searchRequest) {
59+
EventStatsManager.increment(EventStatName.SEMANTIC_HIGHLIGHTING_QUERY_ENRICHER_EXECUTIONS);
60+
Optional<SearchSourceBuilder> source = Optional.ofNullable(searchRequest.source());
61+
source.map(SearchSourceBuilder::highlighter).ifPresent(this::enrichHighlight);
62+
source.map(SearchSourceBuilder::query).ifPresent(qb -> qb.visit(new NestedQueryHighlightVisitor()));
63+
// NOTE: we explicitly do not enrich TopHitsAggregationBuilder highlighters because it is not useful yet — the batch path ignores
64+
// aggregations entirely (HighlightConfigResolver only walks source.highlighter() and inner_hits, HighlightContextBuilder only reads
65+
// response.getHits()), so under ext.semantic_highlighting_batch the model_id would be set but never used, and highlights would go
66+
// missing silently. Needs the highlighting feature to support aggregations first.
67+
return searchRequest;
68+
}
69+
70+
private void enrichHighlight(HighlightBuilder hlBuilder) {
71+
Map<String, Object> globalOptions = hlBuilder.options();
72+
boolean userSuppliedGlobalModelId = globalOptions != null && globalOptions.containsKey(SemanticHighlightingConstants.MODEL_ID);
73+
if (userSuppliedGlobalModelId) {
74+
// if the user provided a global model_id there's no need to enrich anything.
75+
return;
76+
}
77+
78+
boolean globalIsSemantic = false;
79+
if (SemanticHighlightingConstants.HIGHLIGHTER_TYPE.equals(hlBuilder.highlighterType())) {
80+
globalIsSemantic = true;
81+
if (modelId != null) {
82+
hlBuilder.options(enrichWithModelId(globalOptions, modelId));
83+
}
84+
}
85+
for (HighlightBuilder.Field field : Optional.ofNullable(hlBuilder.fields()).orElseGet(Collections::emptyList)) {
86+
// Enrich if either:
87+
// - the global type is semantic and the field specific type is unset
88+
// - the field specific type is set to semantic
89+
if ((globalIsSemantic && field.highlighterType() == null)
90+
|| SemanticHighlightingConstants.HIGHLIGHTER_TYPE.equals(field.highlighterType())) {
91+
String fieldModelId = (String) Optional.ofNullable(this.fieldDefaultIdMap)
92+
.orElseGet(Collections::emptyMap)
93+
.getOrDefault(field.name(), modelId);
94+
if (fieldModelId != null) {
95+
field.options(enrichWithModelId(field.options(), fieldModelId));
96+
}
97+
// else: no default model_id and no per-field override for this field, nothing to enrich
98+
}
99+
}
100+
}
101+
102+
private Map<String, Object> enrichWithModelId(@Nullable Map<String, Object> options, String modelId) {
103+
if (options != null && options.containsKey(SemanticHighlightingConstants.MODEL_ID)) {
104+
return options;
105+
}
106+
Map<String, Object> enrichedOptions = options != null ? new HashMap<>(options) : new HashMap<>();
107+
enrichedOptions.put(SemanticHighlightingConstants.MODEL_ID, modelId);
108+
return enrichedOptions;
109+
}
110+
111+
private class NestedQueryHighlightVisitor implements QueryBuilderVisitor {
112+
@Override
113+
public void accept(QueryBuilder qb) {
114+
if (qb instanceof NestedQueryBuilder nested) {
115+
Optional.ofNullable(nested.innerHit())
116+
.map(InnerHitBuilder::getHighlightBuilder)
117+
.ifPresent(SemanticHighlighterQueryEnricherProcessor.this::enrichHighlight);
118+
}
119+
}
120+
121+
@Override
122+
public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
123+
return this;
124+
}
125+
}
126+
127+
@Override
128+
public String getType() {
129+
return SemanticHighlightingConstants.QUERY_ENRICHER_TYPE;
130+
}
131+
132+
public static class Factory implements Processor.Factory<SearchRequestProcessor> {
133+
private static final String DEFAULT_MODEL_ID = "default_model_id";
134+
private static final String SEMANTIC_HIGHLIGHTER_FIELD_DEFAULT_ID = "semantic_highlighter_field_default_id";
135+
136+
/**
137+
* Create the processor object.
138+
*
139+
* @return {@link SemanticHighlighterQueryEnricherProcessor}
140+
*/
141+
@Override
142+
public SemanticHighlighterQueryEnricherProcessor create(
143+
Map<String, Processor.Factory<SearchRequestProcessor>> processorFactories,
144+
String tag,
145+
String description,
146+
boolean ignoreFailure,
147+
Map<String, Object> config,
148+
PipelineContext pipelineContext
149+
) throws IllegalArgumentException {
150+
String modelId = readOptionalStringProperty(TYPE, tag, config, DEFAULT_MODEL_ID);
151+
Map<String, Object> fieldMap = ConfigurationUtils.readOptionalMap(TYPE, tag, config, SEMANTIC_HIGHLIGHTER_FIELD_DEFAULT_ID);
152+
153+
if (modelId == null && fieldMap == null) {
154+
throw new IllegalArgumentException("[default_model_id] or [semantic_highlighter_field_default_id] should be provided");
155+
}
156+
157+
if (fieldMap != null) {
158+
List<String> nonStringFields = fieldMap.entrySet()
159+
.stream()
160+
.filter(en -> !(en.getValue() instanceof String))
161+
.map(Map.Entry::getKey)
162+
.toList();
163+
if (!nonStringFields.isEmpty()) {
164+
throw new IllegalArgumentException(
165+
"Invalid type in [semantic_highlighter_field_default_id]: value for ["
166+
+ String.join(", ", nonStringFields)
167+
+ "] must be a model_id of type string"
168+
);
169+
}
170+
}
171+
172+
return new SemanticHighlighterQueryEnricherProcessor(
173+
tag,
174+
description,
175+
ignoreFailure,
176+
modelId,
177+
fieldMap != null ? Collections.unmodifiableMap(fieldMap) : null
178+
);
179+
}
180+
}
181+
}

‎src/main/java/org/opensearch/neuralsearch/stats/events/EventStatName.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -250,6 +250,13 @@ public enum EventStatName implements StatName {
250250
EventStatType.TIMESTAMPED_EVENT_COUNTER,
251251
Version.V_3_1_0
252252
),
253+
/** Tracks executions of the semantic highlighter query enricher processor */
254+
SEMANTIC_HIGHLIGHTING_QUERY_ENRICHER_EXECUTIONS(
255+
"semantic_highlighting_query_enricher_executions",
256+
"processors.search",
257+
EventStatType.TIMESTAMPED_EVENT_COUNTER,
258+
Version.V_3_8_0
259+
),
253260
/** Tracks executions of the ML reranking processor */
254261
RERANK_ML_PROCESSOR_EXECUTIONS("rerank_ml_executions", "processors.search", EventStatType.TIMESTAMPED_EVENT_COUNTER, Version.V_3_1_0),
255262

0 commit comments

Comments
 (0)