Update to use PlainActionFuture to fetch model result and use NeuralHighlighterEngine class name

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

Author: junqiu-lei

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

@@ -19,13 +19,15 @@
 public class NeuralHighlighter implements Highlighter {
     public static final String NAME = "neural";
 
-    private NeuralHighlighterManager neuralHighlighterManager;
+    private NeuralHighlighterEngine neuralHighlighterEngine;
 
     public void initialize(MLCommonsClientAccessor mlClient) {
-        if (neuralHighlighterManager != null) {
-            throw new IllegalStateException("NeuralHighlighter has already been initialized. Multiple initializations are not permitted.");
+        if (neuralHighlighterEngine != null) {
+            throw new IllegalStateException(
+                "NeuralHighlighterEngine has already been initialized. Multiple initializations are not permitted."
+            );
         }
-        this.neuralHighlighterManager = new NeuralHighlighterManager(mlClient);
+        this.neuralHighlighterEngine = new NeuralHighlighterEngine(mlClient);
     }
 
     @Override
@@ -41,26 +43,26 @@ public boolean canHighlight(MappedFieldType fieldType) {
      */
     @Override
     public HighlightField highlight(FieldHighlightContext fieldContext) {
-        if (neuralHighlighterManager == null) {
+        if (neuralHighlighterEngine == null) {
             throw new IllegalStateException("NeuralHighlighter has not been initialized");
         }
 
         // Extract field text
-        String fieldText = neuralHighlighterManager.getFieldText(fieldContext);
+        String fieldText = neuralHighlighterEngine.getFieldText(fieldContext);
 
         // Get model ID
-        String modelId = neuralHighlighterManager.getModelId(fieldContext.field.fieldOptions().options());
+        String modelId = neuralHighlighterEngine.getModelId(fieldContext.field.fieldOptions().options());
 
         // Try to extract query text
-        String originalQueryText = neuralHighlighterManager.extractOriginalQuery(fieldContext.query, fieldContext.fieldName);
+        String originalQueryText = neuralHighlighterEngine.extractOriginalQuery(fieldContext.query, fieldContext.fieldName);
 
         if (originalQueryText == null || originalQueryText.isEmpty()) {
             log.warn("No query text found for field {}", fieldContext.fieldName);
             return null;
         }
 
         // Get highlighted text - allow any exceptions from this call to propagate
-        String highlightedText = neuralHighlighterManager.getHighlightedSentences(modelId, originalQueryText, fieldText);
+        String highlightedText = neuralHighlighterEngine.getHighlightedSentences(modelId, originalQueryText, fieldText);
 
         // Create highlight field
         Text[] fragments = new Text[] { new Text(highlightedText) };

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

@@ -8,26 +8,25 @@
 import org.apache.commons.lang.StringUtils;
 import org.apache.lucene.search.Query;
 import org.opensearch.OpenSearchException;
-import org.opensearch.core.action.ActionListener;
 import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor;
 import org.opensearch.neuralsearch.processor.SentenceHighlightingRequest;
 import org.opensearch.search.fetch.subphase.highlight.FieldHighlightContext;
 import org.opensearch.neuralsearch.highlight.extractor.QueryTextExtractorRegistry;
+import org.opensearch.action.support.PlainActionFuture;
+import lombok.NonNull;
 
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Manager class for neural highlighting operations that handles the core highlighting logic
  */
 @Log4j2
-public class NeuralHighlighterManager {
+public class NeuralHighlighterEngine {
     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>";
@@ -38,8 +37,8 @@ public class NeuralHighlighterManager {
     private final MLCommonsClientAccessor mlCommonsClient;
     private final QueryTextExtractorRegistry queryTextExtractorRegistry;
 
-    public NeuralHighlighterManager(MLCommonsClientAccessor mlCommonsClient) {
-        this.mlCommonsClient = Objects.requireNonNull(mlCommonsClient, "ML Commons client cannot be null");
+    public NeuralHighlighterEngine(@NonNull MLCommonsClientAccessor mlCommonsClient) {
+        this.mlCommonsClient = mlCommonsClient;
         this.queryTextExtractorRegistry = new QueryTextExtractorRegistry();
     }
 
@@ -121,40 +120,24 @@ public String getHighlightedSentences(String modelId, String question, String co
      * @return The highlighting results
      */
     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<>();
-        AtomicReference<Exception> exceptionRef = new AtomicReference<>();
+        PlainActionFuture<List<Map<String, Object>>> future = PlainActionFuture.newFuture();
 
         SentenceHighlightingRequest request = SentenceHighlightingRequest.builder()
             .modelId(modelId)
             .question(question)
             .context(context)
             .build();
 
-        mlCommonsClient.inferenceSentenceHighlighting(request, ActionListener.wrap(result -> {
-            resultRef.set(result);
-            latch.countDown();
-        }, exception -> {
-            exceptionRef.set(exception);
-            latch.countDown();
-        }));
+        mlCommonsClient.inferenceSentenceHighlighting(request, future);
 
         try {
-            latch.await();
-        } catch (InterruptedException e) {
-            Thread.currentThread().interrupt();
+            return future.actionGet();
+        } catch (Exception e) {
             throw new OpenSearchException(
-                String.format(Locale.ROOT, "Interrupted while waiting for sentence highlighting inference from model [%s]", modelId),
+                String.format(Locale.ROOT, "Error during sentence highlighting inference from model [%s]", modelId),
                 e
             );
         }
-
-        if (exceptionRef.get() != null) {
-            throw new OpenSearchException("Error during sentence highlighting inference", exceptionRef.get());
-        }
-
-        return resultRef.get();
     }
 
     /**

src/main/java/org/opensearch/neuralsearch/highlight/extractor/QueryTextExtractorRegistry.java

@@ -32,7 +32,7 @@ public QueryTextExtractorRegistry() {
     /**
      * Initializes the registry with default extractors
      */
-    public void initialize() {
+    private void initialize() {
         register(NeuralKNNQuery.class, new NeuralQueryTextExtractor());
         register(TermQuery.class, new TermQueryTextExtractor());
 

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

@@ -8,7 +8,6 @@
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.TermQuery;
-import org.opensearch.OpenSearchException;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor;
 import org.opensearch.neuralsearch.processor.SentenceHighlightingRequest;
@@ -19,16 +18,14 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 /**
- * Tests for the NeuralHighlighterManager class
+ * Tests for the NeuralHighlighterEngine class
  */
 public class NeuralHighlighterManagerTests extends OpenSearchTestCase {
 
@@ -37,14 +34,14 @@ public class NeuralHighlighterManagerTests extends OpenSearchTestCase {
     private static final String TEST_CONTENT = "This is a test content. For highlighting purposes. With multiple sentences.";
     private static final String TEST_QUERY = "test content";
 
-    private NeuralHighlighterManager manager;
+    private NeuralHighlighterEngine manager;
     private MLCommonsClientAccessor mlCommonsClientAccessor;
 
     @Override
     public void setUp() throws Exception {
         super.setUp();
         mlCommonsClientAccessor = mock(MLCommonsClientAccessor.class);
-        manager = new NeuralHighlighterManager(mlCommonsClientAccessor);
+        manager = new NeuralHighlighterEngine(mlCommonsClientAccessor);
 
         // Setup default mock behavior
         setupDefaultMockBehavior();
@@ -213,67 +210,4 @@ public void testApplyHighlightingWithInvalidPositions() {
         // Should only apply the valid highlight
         assertEquals("Should only apply valid highlights", "<em>This</em> is a test string", result);
     }
-
-    public void testFetchModelResultsWithTimeout() throws Exception {
-        // Create a custom mock that delays the response
-        MLCommonsClientAccessor delayedMlClient = mock(MLCommonsClientAccessor.class);
-        NeuralHighlighterManager customManager = new NeuralHighlighterManager(delayedMlClient);
-
-        // Use a CountDownLatch to control the test timing
-        CountDownLatch latch = new CountDownLatch(1);
-
-        // Mock response with delay
-        doAnswer(invocation -> {
-            ActionListener<List<Map<String, Object>>> listener = invocation.getArgument(1);
-
-            // Start a new thread to delay the response
-            new Thread(() -> {
-                try {
-                    // Simulate a delay longer than any reasonable timeout
-                    Thread.sleep(500);
-
-                    // Create mock response
-                    List<Map<String, Object>> mockResponse = new ArrayList<>();
-                    Map<String, Object> resultMap = new HashMap<>();
-                    List<Map<String, Object>> highlights = new ArrayList<>();
-
-                    Map<String, Object> highlight = new HashMap<>();
-                    highlight.put("start", 0);
-                    highlight.put("end", 4);
-                    highlights.add(highlight);
-
-                    resultMap.put("highlights", highlights);
-                    mockResponse.add(resultMap);
-
-                    listener.onResponse(mockResponse);
-                } catch (InterruptedException e) {
-                    listener.onFailure(e);
-                } finally {
-                    latch.countDown();
-                }
-            }).start();
-
-            return null;
-        }).when(delayedMlClient).inferenceSentenceHighlighting(any(SentenceHighlightingRequest.class), any());
-
-        // Call the method in a separate thread so we can interrupt it
-        Thread testThread = new Thread(() -> {
-            try {
-                customManager.fetchModelResults(MODEL_ID, TEST_QUERY, TEST_CONTENT);
-                fail("Should have been interrupted");
-            } catch (OpenSearchException e) {
-                // Expected exception
-                assertTrue(e.getMessage().contains("Interrupted while waiting"));
-            }
-        });
-
-        testThread.start();
-
-        // Wait a bit and then interrupt the thread
-        Thread.sleep(100);
-        testThread.interrupt();
-
-        // Wait for the mock to finish
-        latch.await(1, TimeUnit.SECONDS);
-    }
 }

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

@@ -29,8 +29,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
@@ -50,7 +48,7 @@ public class NeuralHighlighterTests extends OpenSearchTestCase {
     private NeuralHighlighter highlighter;
     private MLCommonsClientAccessor mlCommonsClientAccessor;
     private MappedFieldType fieldType;
-    private NeuralHighlighterManager manager;
+    private NeuralHighlighterEngine manager;
 
     @Override
     public void setUp() throws Exception {
@@ -59,7 +57,7 @@ public void setUp() throws Exception {
         highlighter = new NeuralHighlighter();
         highlighter.initialize(mlCommonsClientAccessor);
         fieldType = new TextFieldMapper.TextFieldType(TEST_FIELD);
-        manager = new NeuralHighlighterManager(mlCommonsClientAccessor);
+        manager = new NeuralHighlighterEngine(mlCommonsClientAccessor);
 
         // Setup default mock behavior
         setupDefaultMockBehavior();
@@ -199,7 +197,10 @@ public void testMultipleInitialization() {
         IllegalStateException exception = expectThrows(IllegalStateException.class, () -> testHighlighter.initialize(mlClient2));
 
         assertNotNull(exception);
-        assertEquals("NeuralHighlighter has already been initialized. Multiple initializations are not permitted.", exception.getMessage());
+        assertEquals(
+            "NeuralHighlighterEngine has already been initialized. Multiple initializations are not permitted.",
+            exception.getMessage()
+        );
     }
 
     public void testUninitializedHighlighter() {
@@ -211,7 +212,7 @@ public void testUninitializedHighlighter() {
         assertTrue(exception.getMessage().contains("has not been initialized"));
     }
 
-    // Tests for the NeuralHighlighterManager class
+    // Tests for the NeuralHighlighterEngine class
 
     public void testGetModelId() {
         Map<String, Object> options = new HashMap<>();
@@ -353,71 +354,6 @@ public void testApplyHighlightingWithInvalidPositions() {
         assertEquals("Should only apply valid highlights", "<em>This</em> is a test string", result);
     }
 
-    public void testFetchHighlightingResultsWithTimeout() throws Exception {
-        // Create a custom mock that delays the response
-        MLCommonsClientAccessor delayedMlClient = mock(MLCommonsClientAccessor.class);
-        NeuralHighlighterManager customManager = new NeuralHighlighterManager(delayedMlClient);
-
-        // Use a CountDownLatch to control the test timing
-        CountDownLatch latch = new CountDownLatch(1);
-
-        // Mock response with delay
-        doAnswer(invocation -> {
-            ActionListener<List<Map<String, Object>>> listener = invocation.getArgument(1);
-
-            // Start a new thread to delay the response
-            new Thread(() -> {
-                try {
-                    // Simulate a delay longer than any reasonable timeout
-                    Thread.sleep(500);
-
-                    // Create mock response
-                    List<Map<String, Object>> mockResponse = new ArrayList<>();
-                    Map<String, Object> resultMap = new HashMap<>();
-                    List<Map<String, Object>> highlights = new ArrayList<>();
-
-                    Map<String, Object> highlight = new HashMap<>();
-                    highlight.put("start", 0);
-                    highlight.put("end", 4);
-                    highlights.add(highlight);
-
-                    resultMap.put("highlights", highlights);
-                    mockResponse.add(resultMap);
-
-                    listener.onResponse(mockResponse);
-                } catch (InterruptedException e) {
-                    listener.onFailure(e);
-                } finally {
-                    latch.countDown();
-                }
-            }).start();
-
-            return null;
-        }).when(delayedMlClient).inferenceSentenceHighlighting(any(SentenceHighlightingRequest.class), any());
-
-        // Call the method in a separate thread so we can interrupt it
-        Thread testThread = new Thread(() -> {
-            try {
-                customManager.fetchModelResults(MODEL_ID, TEST_QUERY, TEST_CONTENT);
-                fail("Should have been interrupted");
-            } catch (OpenSearchException e) {
-                // Expected exception
-                assertTrue(e.getMessage().contains("Interrupted while waiting"));
-            }
-        });
-
-        testThread.start();
-
-        // Wait a bit and then interrupt the thread
-        Thread.sleep(100);
-        testThread.interrupt();
-
-        // Wait for the mock to finish
-        latch.await(1, TimeUnit.SECONDS);
-    }
-
-    // Integration and error handling tests
-
     public void testIntegrationWithTermQuery() {
         // Integration test logic for term queries
         TermQuery query = new TermQuery(new Term(TEST_FIELD, "test"));

src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java

@@ -251,7 +251,7 @@ public void testNeuralQueryWithNeuralHighlighting() {
         List<String> highlightedFields = (List<String>) highlight.get(TEST_TEXT_FIELD_NAME_1);
         assertNotNull("Highlight should contain the requested field", highlightedFields);
 
-        String highlightedText = highlightedFields.get(0);
+        String highlightedText = highlightedFields.getFirst();
         assertTrue(
             "Highlighted text should contain content from the original text",
             highlightedText.contains("neural highlighting")