Skip to content

Commit f1569c1

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 f1569c1

7 files changed

Lines changed: 815 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/issues/1916))
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: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
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+
if (hlBuilder.fields() == null) {
86+
return;
87+
}
88+
for (HighlightBuilder.Field field : hlBuilder.fields()) {
89+
// Enrich if either:
90+
// - the global type is semantic and the field specific type is unset
91+
// - the field specific type is set to semantic
92+
if ((globalIsSemantic && field.highlighterType() == null)
93+
|| SemanticHighlightingConstants.HIGHLIGHTER_TYPE.equals(field.highlighterType())) {
94+
String fieldModelId = (String) Optional.ofNullable(this.fieldDefaultIdMap)
95+
.orElseGet(Collections::emptyMap)
96+
.getOrDefault(field.name(), modelId);
97+
if (fieldModelId != null) {
98+
field.options(enrichWithModelId(field.options(), fieldModelId));
99+
}
100+
// else: no default model_id and no per-field override for this field, nothing to enrich
101+
}
102+
}
103+
}
104+
105+
private Map<String, Object> enrichWithModelId(@Nullable Map<String, Object> options, String modelId) {
106+
if (options != null && options.containsKey(SemanticHighlightingConstants.MODEL_ID)) {
107+
return options;
108+
}
109+
Map<String, Object> enrichedOptions = options != null ? new HashMap<>(options) : new HashMap<>();
110+
enrichedOptions.put(SemanticHighlightingConstants.MODEL_ID, modelId);
111+
return enrichedOptions;
112+
}
113+
114+
private class NestedQueryHighlightVisitor implements QueryBuilderVisitor {
115+
@Override
116+
public void accept(QueryBuilder qb) {
117+
if (qb instanceof NestedQueryBuilder nested) {
118+
Optional.ofNullable(nested.innerHit())
119+
.map(InnerHitBuilder::getHighlightBuilder)
120+
.ifPresent(SemanticHighlighterQueryEnricherProcessor.this::enrichHighlight);
121+
}
122+
}
123+
124+
@Override
125+
public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
126+
return this;
127+
}
128+
}
129+
130+
@Override
131+
public String getType() {
132+
return SemanticHighlightingConstants.QUERY_ENRICHER_TYPE;
133+
}
134+
135+
public static class Factory implements Processor.Factory<SearchRequestProcessor> {
136+
private static final String DEFAULT_MODEL_ID = "default_model_id";
137+
private static final String SEMANTIC_HIGHLIGHTER_FIELD_DEFAULT_ID = "semantic_highlighter_field_default_id";
138+
139+
/**
140+
* Create the processor object.
141+
*
142+
* @return {@link SemanticHighlighterQueryEnricherProcessor}
143+
*/
144+
@Override
145+
public SemanticHighlighterQueryEnricherProcessor create(
146+
Map<String, Processor.Factory<SearchRequestProcessor>> processorFactories,
147+
String tag,
148+
String description,
149+
boolean ignoreFailure,
150+
Map<String, Object> config,
151+
PipelineContext pipelineContext
152+
) throws IllegalArgumentException {
153+
String modelId = readOptionalStringProperty(TYPE, tag, config, DEFAULT_MODEL_ID);
154+
Map<String, Object> fieldMap = ConfigurationUtils.readOptionalMap(TYPE, tag, config, SEMANTIC_HIGHLIGHTER_FIELD_DEFAULT_ID);
155+
156+
if (modelId == null && fieldMap == null) {
157+
throw new IllegalArgumentException("[default_model_id] or [semantic_highlighter_field_default_id] should be provided");
158+
}
159+
160+
if (fieldMap != null) {
161+
List<String> nonStringFields = fieldMap.entrySet()
162+
.stream()
163+
.filter(en -> !(en.getValue() instanceof String))
164+
.map(Map.Entry::getKey)
165+
.toList();
166+
if (!nonStringFields.isEmpty()) {
167+
throw new IllegalArgumentException(
168+
"Invalid type in [semantic_highlighter_field_default_id]: value for ["
169+
+ String.join(", ", nonStringFields)
170+
+ "] must be a model_id of type string"
171+
);
172+
}
173+
}
174+
175+
return new SemanticHighlighterQueryEnricherProcessor(
176+
tag,
177+
description,
178+
ignoreFailure,
179+
modelId,
180+
fieldMap != null ? Collections.unmodifiableMap(fieldMap) : null
181+
);
182+
}
183+
}
184+
}

‎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)