Support custom tags in semantic highlighter (opensearch-project#1254)

junqiu-lei · YeonghyeonKO · commit fa9ff0b075e1 · 2025-05-30T21:58:41.000+09:00
Signed-off-by: yeonghyeonKo &lt;dk02315@gmail.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Optimize embedding generation in Sparse Encoding Processor ([#1246](https://github.com/opensearch-project/neural-search/pull/1246))
 - Optimize embedding generation in Text/Image Embedding Processor ([#1249](https://github.com/opensearch-project/neural-search/pull/1249))
 - Inner hits support with hybrid query ([#1253](https://github.com/opensearch-project/neural-search/pull/1253))
+- Support custom tags in semantic highlighter ([#1254](https://github.com/opensearch-project/neural-search/pull/1254))
 
 ### Enhancements
 
diff --git a/src/main/java/org/opensearch/neuralsearch/highlight/SemanticHighlighter.java b/src/main/java/org/opensearch/neuralsearch/highlight/SemanticHighlighter.java
@@ -60,8 +60,18 @@ public HighlightField highlight(FieldHighlightContext fieldContext) {
             return null;
         }
 
+        // The pre- and post- tags are provided by the user or defaulted to <em> and </em>
+        String[] preTags = fieldContext.field.fieldOptions().preTags();
+        String[] postTags = fieldContext.field.fieldOptions().postTags();
+
         // Get highlighted text - allow any exceptions from this call to propagate
-        String highlightedResponse = semanticHighlighterEngine.getHighlightedSentences(modelId, originalQueryText, fieldText);
+        String highlightedResponse = semanticHighlighterEngine.getHighlightedSentences(
+            modelId,
+            originalQueryText,
+            fieldText,
+            preTags[0],
+            postTags[0]
+        );
 
         if (highlightedResponse == null || highlightedResponse.isEmpty()) {
             log.warn("No highlighted text found for field {}", fieldContext.fieldName);
diff --git a/src/main/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterEngine.java b/src/main/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterEngine.java
@@ -13,6 +13,7 @@
 import org.opensearch.neuralsearch.highlight.extractor.QueryTextExtractorRegistry;
 import org.opensearch.action.support.PlainActionFuture;
 import lombok.NonNull;
+import lombok.Builder;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -24,24 +25,18 @@
  * Engine class for semantic highlighting operations
  */
 @Log4j2
+@Builder
 public class SemanticHighlighterEngine {
     private static final String MODEL_ID_FIELD = "model_id";
-    private static final String DEFAULT_PRE_TAG = "<em>";
-    private static final String DEFAULT_POST_TAG = "</em>";
     private static final String MODEL_INFERENCE_RESULT_KEY = "highlights";
     private static final String MODEL_INFERENCE_RESULT_START_KEY = "start";
     private static final String MODEL_INFERENCE_RESULT_END_KEY = "end";
 
+    @NonNull
     private final MLCommonsClientAccessor mlCommonsClient;
-    private final QueryTextExtractorRegistry queryTextExtractorRegistry;
 
-    public SemanticHighlighterEngine(
-        @NonNull MLCommonsClientAccessor mlCommonsClient,
-        @NonNull QueryTextExtractorRegistry queryTextExtractorRegistry
-    ) {
-        this.mlCommonsClient = mlCommonsClient;
-        this.queryTextExtractorRegistry = queryTextExtractorRegistry;
-    }
+    @NonNull
+    private final QueryTextExtractorRegistry queryTextExtractorRegistry;
 
     /**
      * Gets the field text from the document
@@ -116,15 +111,17 @@ public String getModelId(Map<String, Object> options) {
      * @param modelId The ID of the model to use
      * @param question The search query
      * @param context The document text
+     * @param preTag The pre tag to use for highlighting
+     * @param postTag The post tag to use for highlighting
      * @return Formatted text with highlighting
      */
-    public String getHighlightedSentences(String modelId, String question, String context) {
+    public String getHighlightedSentences(String modelId, String question, String context, String preTag, String postTag) {
         List<Map<String, Object>> results = fetchModelResults(modelId, question, context);
         if (results == null || results.isEmpty()) {
             return null;
         }
 
-        return applyHighlighting(context, results.getFirst());
+        return applyHighlighting(context, results.getFirst(), preTag, postTag);
     }
 
     /**
@@ -168,10 +165,12 @@ public List<Map<String, Object>> fetchModelResults(String modelId, String questi
      *
      * @param context The original document text
      * @param highlightResult The highlighting result from the ML model
+     * @param preTag The pre tag to use for highlighting
+     * @param postTag The post tag to use for highlighting
      * @return Formatted text with highlighting
      * @throws IllegalArgumentException if highlight positions are invalid
      */
-    public String applyHighlighting(String context, Map<String, Object> highlightResult) {
+    public String applyHighlighting(String context, Map<String, Object> highlightResult, String preTag, String postTag) {
         // Get the "highlights" list from the result
         Object highlightsObj = highlightResult.get(MODEL_INFERENCE_RESULT_KEY);
 
@@ -216,7 +215,7 @@ public String applyHighlighting(String context, Map<String, Object> highlightRes
             }
         }
 
-        return constructHighlightedText(context, validHighlights);
+        return constructHighlightedText(context, validHighlights, preTag, postTag);
     }
 
     /**
@@ -246,9 +245,11 @@ private void validateHighlightPositions(int start, int end, int textLength) {
      *
      * @param text The original text
      * @param highlights The list of valid highlight positions in pairs [start1, end1, start2, end2, ...]
+     * @param preTag The pre tag to use for highlighting
+     * @param postTag The post tag to use for highlighting
      * @return The highlighted text
      */
-    private String constructHighlightedText(String text, List<Integer> highlights) {
+    private String constructHighlightedText(String text, List<Integer> highlights, String preTag, String postTag) {
         StringBuilder result = new StringBuilder();
         int currentPos = 0;
 
@@ -264,9 +265,9 @@ private String constructHighlightedText(String text, List<Integer> highlights) {
             }
 
             // Add the highlighted text with highlight tags
-            result.append(DEFAULT_PRE_TAG);
+            result.append(preTag);
             result.append(text, start, end);
-            result.append(DEFAULT_POST_TAG);
+            result.append(postTag);
 
             // Update current position to end of this highlight
             currentPos = end;
diff --git a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java
@@ -114,7 +114,10 @@ public Collection<Object> createComponents(
         NeuralQueryBuilder.initialize(clientAccessor);
         NeuralSparseQueryBuilder.initialize(clientAccessor);
         QueryTextExtractorRegistry queryTextExtractorRegistry = new QueryTextExtractorRegistry();
-        SemanticHighlighterEngine semanticHighlighterEngine = new SemanticHighlighterEngine(clientAccessor, queryTextExtractorRegistry);
+        SemanticHighlighterEngine semanticHighlighterEngine = SemanticHighlighterEngine.builder()
+            .mlCommonsClient(clientAccessor)
+            .queryTextExtractorRegistry(queryTextExtractorRegistry)
+            .build();
         semanticHighlighter.initialize(semanticHighlighterEngine);
         HybridQueryExecutor.initialize(threadPool);
         normalizationProcessorWorkflow = new NormalizationProcessorWorkflow(new ScoreNormalizer(), new ScoreCombiner());
diff --git a/src/test/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterEngineTests.java b/src/test/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterEngineTests.java
@@ -45,7 +45,10 @@ public void setUp() throws Exception {
         super.setUp();
         mlCommonsClientAccessor = mock(MLCommonsClientAccessor.class);
         queryTextExtractorRegistry = new QueryTextExtractorRegistry();
-        highlighterEngine = new SemanticHighlighterEngine(mlCommonsClientAccessor, queryTextExtractorRegistry);
+        highlighterEngine = SemanticHighlighterEngine.builder()
+            .mlCommonsClient(mlCommonsClientAccessor)
+            .queryTextExtractorRegistry(queryTextExtractorRegistry)
+            .build();
 
         // Setup default mock behavior
         setupDefaultMockBehavior();
@@ -123,7 +126,7 @@ public void testExtractOriginalQuery() {
     }
 
     public void testGetHighlightedSentences() {
-        String result = highlighterEngine.getHighlightedSentences(MODEL_ID, TEST_QUERY, TEST_CONTENT);
+        String result = highlighterEngine.getHighlightedSentences(MODEL_ID, TEST_QUERY, TEST_CONTENT, "<em>", "</em>");
 
         assertNotNull("Should return highlighted text", result);
         assertTrue("Should contain highlighting tags", result.contains("<em>") && result.contains("</em>"));
@@ -149,7 +152,7 @@ public void testApplyHighlighting() {
         resultMap.put("highlights", highlightsList);
 
         String text = "This is a test string";
-        String result = highlighterEngine.applyHighlighting(text, resultMap);
+        String result = highlighterEngine.applyHighlighting(text, resultMap, "<em>", "</em>");
 
         assertEquals("Should apply highlights correctly", "<em>This</em> is <em>a te</em>st string", result);
     }
@@ -170,7 +173,10 @@ public void testApplyHighlightingWithInvalidPositions() {
         highlightsList.add(highlight1);
         resultMap.put("highlights", highlightsList);
 
-        OpenSearchException exception = expectThrows(OpenSearchException.class, () -> highlighterEngine.applyHighlighting(text, resultMap));
+        OpenSearchException exception = expectThrows(
+            OpenSearchException.class,
+            () -> highlighterEngine.applyHighlighting(text, resultMap, "<em>", "</em>")
+        );
         assertEquals(
             "Should throw correct error message for invalid positions",
             String.format(
@@ -188,7 +194,7 @@ public void testApplyHighlightingWithInvalidPositions() {
         highlight2.put("end", 100);
         highlightsList.add(highlight2);
 
-        exception = expectThrows(OpenSearchException.class, () -> highlighterEngine.applyHighlighting(text, resultMap));
+        exception = expectThrows(OpenSearchException.class, () -> highlighterEngine.applyHighlighting(text, resultMap, "<em>", "</em>"));
         assertEquals(
             "Should throw correct error message for invalid positions",
             String.format(
@@ -206,7 +212,7 @@ public void testApplyHighlightingWithInvalidPositions() {
         highlight3.put("end", 5);
         highlightsList.add(highlight3);
 
-        exception = expectThrows(OpenSearchException.class, () -> highlighterEngine.applyHighlighting(text, resultMap));
+        exception = expectThrows(OpenSearchException.class, () -> highlighterEngine.applyHighlighting(text, resultMap, "<em>", "</em>"));
         assertEquals(
             "Should throw correct error message for invalid positions",
             String.format(
@@ -238,20 +244,15 @@ public void testApplyHighlightingWithUnsortedPositions() {
 
         resultMap.put("highlights", highlightsList);
 
-        OpenSearchException exception = expectThrows(OpenSearchException.class, () -> highlighterEngine.applyHighlighting(text, resultMap));
+        OpenSearchException exception = expectThrows(
+            OpenSearchException.class,
+            () -> highlighterEngine.applyHighlighting(text, resultMap, "<em>", "</em>")
+        );
         assertEquals(
             "Should throw correct error message for unsorted positions",
             "Internal error while applying semantic highlight: received unsorted highlights from model",
             exception.getMessage()
         );
-        // Verify that sorted positions work correctly
-        highlightsList.clear();
-        // Add highlights in sorted order
-        highlightsList.add(highlight2);  // start=0
-        highlightsList.add(highlight1);  // start=8
-
-        String result = highlighterEngine.applyHighlighting(text, resultMap);
-        assertEquals("Should successfully highlight with sorted positions", "<em>This</em> is <em>a</em> test string", result);
     }
 
     public void testApplyHighlightingWithInvalidHighlightMap() {
@@ -268,7 +269,10 @@ public void testApplyHighlightingWithInvalidHighlightMap() {
         resultMap.put("highlights", highlightsList);
 
         String text = "This is a test string";
-        ClassCastException exception = expectThrows(ClassCastException.class, () -> highlighterEngine.applyHighlighting(text, resultMap));
+        ClassCastException exception = expectThrows(
+            ClassCastException.class,
+            () -> highlighterEngine.applyHighlighting(text, resultMap, "<em>", "</em>")
+        );
         assertTrue(exception.getMessage().contains("cannot be cast to class java.lang.Number"));
     }
 
@@ -286,7 +290,10 @@ public void testApplyHighlightingWithMissingPositions() {
         resultMap.put("highlights", highlightsList);
 
         String text = "This is a test string";
-        OpenSearchException exception = expectThrows(OpenSearchException.class, () -> highlighterEngine.applyHighlighting(text, resultMap));
+        OpenSearchException exception = expectThrows(
+            OpenSearchException.class,
+            () -> highlighterEngine.applyHighlighting(text, resultMap, "<em>", "</em>")
+        );
         assertTrue(exception.getMessage().contains("Missing start or end position"));
     }
 
@@ -296,17 +303,31 @@ public void testApplyHighlightingWithEmptyHighlights() {
         resultMap.put("highlights", new ArrayList<>());
 
         String text = "This is a test string";
-        String result = highlighterEngine.applyHighlighting(text, resultMap);
+        String result = highlighterEngine.applyHighlighting(text, resultMap, "<em>", "</em>");
         assertEquals("Should return original text when no highlights", text, result);
     }
 
     public void testApplyHighlightingWithMissingHighlightsKey() {
-        // Test with missing highlights key
         Map<String, Object> resultMap = new HashMap<>();
-        // No highlights key
-
         String text = "This is a test string";
-        String result = highlighterEngine.applyHighlighting(text, resultMap);
+        String result = highlighterEngine.applyHighlighting(text, resultMap, "<em>", "</em>");
         assertNull("Should return null when highlights key is missing", result);
     }
+
+    public void testCustomTags() {
+        // Test with custom tags
+        String result = highlighterEngine.getHighlightedSentences(MODEL_ID, TEST_QUERY, TEST_CONTENT, "<mark>", "</mark>");
+        assertNotNull("Should return highlighted text", result);
+        assertTrue("Should contain custom highlighting tags", result.contains("<mark>") && result.contains("</mark>"));
+        assertFalse("Should not contain default highlighting tags", result.contains("<em>") || result.contains("</em>"));
+
+        // Test with different custom tags
+        result = highlighterEngine.getHighlightedSentences(MODEL_ID, TEST_QUERY, TEST_CONTENT, "<span class='highlight'>", "</span>");
+        assertNotNull("Should return highlighted text", result);
+        assertTrue(
+            "Should contain new custom highlighting tags",
+            result.contains("<span class='highlight'>") && result.contains("</span>")
+        );
+        assertFalse("Should not contain previous highlighting tags", result.contains("<mark>") || result.contains("</mark>"));
+    }
 }
diff --git a/src/test/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterIT.java b/src/test/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterIT.java
@@ -33,6 +33,10 @@ public class SemanticHighlighterIT extends BaseNeuralSearchIT {
     private static final String TEST_CONTENT = "Machine learning is a field of artificial intelligence that uses statistical techniques. "
         + "Natural language processing is a branch of artificial intelligence that helps computers understand human language. "
         + "Deep learning is a subset of machine learning that uses neural networks with many layers.";
+    private static final String DEFAULT_PRE_TAG = "<em>";
+    private static final String DEFAULT_POST_TAG = "</em>";
+    private static final String CUSTOM_PRE_TAG = "<test>";
+    private static final String CUSTOM_POST_TAG = "</test>";
     private final float[] testVector = createRandomVector(TEST_DIMENSION);
 
     @Before
@@ -73,7 +77,7 @@ private void initializeTestIndex() {
      * 5. Neural Query
      * 6. Hybrid Query
      */
-    public void testQueriesWithSemanticHighlighter() throws Exception {
+    public void testQueriesWithSemanticHighlighter() {
         // Set up models for the test
         String textEmbeddingModelId = prepareModel();
         String sentenceHighlightingModelId = prepareSentenceHighlightingModel();
@@ -222,7 +226,9 @@ public void testQueriesWithSemanticHighlighter() throws Exception {
          *         },
          *         "options": {
          *             "model_id": "sentence-highlighting-model-id"
-         *         }
+         *         },
+         *         "pre_tags": ["<test>"],
+         *         "post_tags": ["</test>"]
          *     }
          * }
          */
@@ -232,6 +238,8 @@ public void testQueriesWithSemanticHighlighter() throws Exception {
             .modelId(textEmbeddingModelId)
             .k(1)
             .build();
+
+        // First test with default highlighting tags
         searchResponse = searchWithSemanticHighlighter(
             TEST_BASIC_INDEX_NAME,
             neuralQueryBuilder,
@@ -241,6 +249,40 @@ public void testQueriesWithSemanticHighlighter() throws Exception {
         );
         verifyHighlightResults(searchResponse, TEST_QUERY_TEXT);
 
+        // Then test with custom highlighting tags
+        Map<String, Map<String, Object>> customHighlightFields = Map.of(TEST_TEXT_FIELD_NAME, Map.of("type", "semantic"));
+        Map<String, Object> customHighlightOptions = Map.of("model_id", sentenceHighlightingModelId);
+
+        searchResponse = searchWithHighlight(
+            TEST_BASIC_INDEX_NAME,
+            neuralQueryBuilder,
+            10,
+            customHighlightFields,
+            customHighlightOptions,
+            new String[] { CUSTOM_PRE_TAG },
+            new String[] { CUSTOM_POST_TAG }
+        );
+
+        // Verify results with custom tags
+        Map<String, Object> customTagsHit = getFirstInnerHit(searchResponse);
+        assertNotNull("Search response should contain hits", customTagsHit);
+
+        @SuppressWarnings("unchecked")
+        Map<String, Object> customHighlight = (Map<String, Object>) customTagsHit.get("highlight");
+        assertNotNull("Hit should contain highlight section", customHighlight);
+
+        @SuppressWarnings("unchecked")
+        List<String> customHighlightedFields = (List<String>) customHighlight.get(TEST_TEXT_FIELD_NAME);
+        assertNotNull("Highlight should contain the requested field", customHighlightedFields);
+        assertFalse("Highlighted fields should not be empty", customHighlightedFields.isEmpty());
+
+        String customHighlightedText = customHighlightedFields.getFirst();
+        assertTrue("Text should contain custom opening tag", customHighlightedText.contains(CUSTOM_PRE_TAG));
+        assertTrue("Text should contain custom closing tag", customHighlightedText.contains(CUSTOM_POST_TAG));
+        assertFalse("Text should not contain default opening tag", customHighlightedText.contains(DEFAULT_PRE_TAG));
+        assertFalse("Text should not contain default closing tag", customHighlightedText.contains(DEFAULT_POST_TAG));
+        verifyHighlightResults(searchResponse, TEST_QUERY_TEXT);
+
         // 6. Test Hybrid Query
         /*
          *  example:
@@ -323,10 +365,9 @@ private void verifyHighlightResults(Map<String, Object> searchResponse, String e
 
         assertTrue("Highlighted text should contain semantically relevant content for query: " + expectedContent, hasRelevantContent);
 
-        // Verify the highlight tags are present
-        assertTrue(
-            "Highlighted text should contain proper highlight tags",
-            highlightedText.contains("<em>") && highlightedText.contains("</em>")
-        );
+        // Verify the highlight tags are present - either default tags or custom tags
+        boolean hasDefaultTags = highlightedText.contains(DEFAULT_PRE_TAG) && highlightedText.contains(DEFAULT_POST_TAG);
+        boolean hasCustomTags = highlightedText.contains(CUSTOM_PRE_TAG) && highlightedText.contains(CUSTOM_POST_TAG);
+        assertTrue("Highlighted text should contain either default or custom highlight tags", hasDefaultTags || hasCustomTags);
     }
 }
diff --git a/src/test/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterTests.java b/src/test/java/org/opensearch/neuralsearch/highlight/SemanticHighlighterTests.java
diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java
