Update fetchModelResults name and optimize highlight result handle

Signed-off-by: Junqiu Lei <junqiu@amazon.com>

Author: junqiu-lei

src/main/java/org/opensearch/neuralsearch/highlight/NeuralHighlighterManager.java

@@ -104,7 +104,7 @@ public String getModelId(Map<String, Object> options) {
      * @return Formatted text with highlighting
      */
     public String getHighlightedSentences(String modelId, String question, String context) {
-        List<Map<String, Object>> result = fetchHighlightingResults(modelId, question, context);
+        List<Map<String, Object>> result = fetchModelResults(modelId, question, context);
         if (result == null) {
             return StringUtils.EMPTY;
         }
@@ -120,7 +120,7 @@ public String getHighlightedSentences(String modelId, String question, String co
      * @param context The document text
      * @return The highlighting results
      */
-    public List<Map<String, Object>> fetchHighlightingResults(String modelId, String question, String context) {
+    public List<Map<String, Object>> fetchModelResults(String modelId, String question, String context) {
 
         CountDownLatch latch = new CountDownLatch(1);
         AtomicReference<List<Map<String, Object>>> resultRef = new AtomicReference<>();
@@ -165,44 +165,31 @@ public List<Map<String, Object>> fetchHighlightingResults(String modelId, String
      * @return Formatted text with highlighting
      */
     public String applyHighlighting(String context, List<Map<String, Object>> highlightResults) {
-        // Collect all valid highlight positions
         List<int[]> validHighlights = new ArrayList<>();
 
-        // Process each highlight result
-        for (Map<String, Object> result : highlightResults) {
-            if (result == null) {
-                continue;
-            }
-
-            // Get the "highlights" list from the result
-            Object highlightsObj = result.get(MODEL_INFERENCE_RESULT_KEY);
-
-            // Safely check if the object is a List
-            if (!(highlightsObj instanceof List<?> highlightsList)) {
-                continue;
-            }
-
-            // Process the list safely
-            if (highlightsList.isEmpty()) {
-                continue;
-            }
-
-            // Process each item in the list, checking if it's a Map
-            for (Object item : highlightsList) {
-                if (item instanceof Map<?, ?> map) {
-                    // Create a type-safe map
-                    Map<String, Object> safeMap = new java.util.HashMap<>();
-                    for (Map.Entry<?, ?> entry : map.entrySet()) {
-                        safeMap.put(entry.getKey().toString(), entry.getValue());
-                    }
-
-                    // Extract and validate positions
-                    Object startObj = safeMap.get(MODEL_INFERENCE_RESULT_START_KEY);
-                    Object endObj = safeMap.get(MODEL_INFERENCE_RESULT_END_KEY);
-
-                    int[] positions = validateHighlightPositions(startObj, endObj, context.length());
-                    if (positions != null) {
-                        validHighlights.add(positions);
+        if (highlightResults != null && !highlightResults.isEmpty()) {
+            Map<String, Object> result = highlightResults.getFirst();
+            if (result != null) {
+                // Get the "highlights" list from the result
+                Object highlightsObj = result.get(MODEL_INFERENCE_RESULT_KEY);
+
+                if (highlightsObj instanceof List<?> highlightsList) {
+                    for (Object item : highlightsList) {
+                        if (item instanceof Map<?, ?> map) {
+                            Map<String, Object> safeMap = new java.util.HashMap<>();
+                            for (Map.Entry<?, ?> entry : map.entrySet()) {
+                                safeMap.put(entry.getKey().toString(), entry.getValue());
+                            }
+
+                            // Extract and validate positions
+                            Object startObj = safeMap.get(MODEL_INFERENCE_RESULT_START_KEY);
+                            Object endObj = safeMap.get(MODEL_INFERENCE_RESULT_END_KEY);
+
+                            int[] positions = validateHighlightPositions(startObj, endObj, context.length());
+                            if (positions != null) {
+                                validHighlights.add(positions);
+                            }
+                        }
                     }
                 }
             }
@@ -216,7 +203,6 @@ public String applyHighlighting(String context, List<Map<String, Object>> highli
         // Sort highlights by start position (ascending)
         validHighlights.sort(Comparator.comparingInt(pos -> pos[0]));
 
-        // Construct the highlighted text in O(n) time
         return constructHighlightedText(context, validHighlights);
     }
 
@@ -273,7 +259,7 @@ private List<int[]> mergeOverlappingHighlights(List<int[]> highlights) {
         }
 
         List<int[]> merged = new ArrayList<>();
-        int[] current = highlights.get(0);
+        int[] current = highlights.getFirst();
         merged.add(current);
 
         for (int i = 1; i < highlights.size(); i++) {

src/test/java/org/opensearch/neuralsearch/highlight/NeuralHighlighterManagerTests.java

@@ -214,7 +214,7 @@ public void testApplyHighlightingWithInvalidPositions() {
         assertEquals("Should only apply valid highlights", "<em>This</em> is a test string", result);
     }
 
-    public void testFetchHighlightingResultsWithTimeout() throws Exception {
+    public void testFetchModelResultsWithTimeout() throws Exception {
         // Create a custom mock that delays the response
         MLCommonsClientAccessor delayedMlClient = mock(MLCommonsClientAccessor.class);
         NeuralHighlighterManager customManager = new NeuralHighlighterManager(delayedMlClient);
@@ -259,7 +259,7 @@ public void testFetchHighlightingResultsWithTimeout() throws Exception {
         // Call the method in a separate thread so we can interrupt it
         Thread testThread = new Thread(() -> {
             try {
-                customManager.fetchHighlightingResults(MODEL_ID, TEST_QUERY, TEST_CONTENT);
+                customManager.fetchModelResults(MODEL_ID, TEST_QUERY, TEST_CONTENT);
                 fail("Should have been interrupted");
             } catch (OpenSearchException e) {
                 // Expected exception

src/test/java/org/opensearch/neuralsearch/highlight/NeuralHighlighterTests.java

@@ -398,7 +398,7 @@ public void testFetchHighlightingResultsWithTimeout() throws Exception {
         // Call the method in a separate thread so we can interrupt it
         Thread testThread = new Thread(() -> {
             try {
-                customManager.fetchHighlightingResults(MODEL_ID, TEST_QUERY, TEST_CONTENT);
+                customManager.fetchModelResults(MODEL_ID, TEST_QUERY, TEST_CONTENT);
                 fail("Should have been interrupted");
             } catch (OpenSearchException e) {
                 // Expected exception