Backport/backport 208 to 2.x

src/main/java/org/opensearch/neuralsearch/ml/MLCommonsClientAccessor.java

@@ -36,6 +36,8 @@ public class MLCommonsClientAccessor {
     private static final List<String> TARGET_RESPONSE_FILTERS = List.of("sentence_embedding");
     private final MachineLearningNodeClient mlClient;
 
+    private static final int MAX_RETRY = 3;
+
     /**
      * Wrapper around {@link #inferenceSentences} that expected a single input text and produces a single floating
      * point vector as a response.

src/main/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessor.java

@@ -105,7 +105,7 @@ public void execute(IngestDocument ingestDocument, BiConsumer<IngestDocument, Ex
                 handler.accept(ingestDocument, null);
             } else {
                 mlCommonsClientAccessor.inferenceSentences(this.modelId, inferenceList, ActionListener.wrap(vectors -> {
-                    appendVectorFieldsToDocument(ingestDocument, knnMap, vectors);
+                    setVectorFieldsToDocument(ingestDocument, knnMap, vectors);
                     handler.accept(ingestDocument, null);
                 }, e -> { handler.accept(null, e); }));
             }
@@ -115,11 +115,11 @@ public void execute(IngestDocument ingestDocument, BiConsumer<IngestDocument, Ex
 
     }
 
-    void appendVectorFieldsToDocument(IngestDocument ingestDocument, Map<String, Object> knnMap, List<List<Float>> vectors) {
+    void setVectorFieldsToDocument(IngestDocument ingestDocument, Map<String, Object> knnMap, List<List<Float>> vectors) {
         Objects.requireNonNull(vectors, "embedding failed, inference returns null result!");
         log.debug("Text embedding result fetched, starting build vector output!");
         Map<String, Object> textEmbeddingResult = buildTextEmbeddingResult(knnMap, vectors, ingestDocument.getSourceAndMetadata());
-        textEmbeddingResult.forEach(ingestDocument::appendFieldValue);
+        textEmbeddingResult.forEach(ingestDocument::setFieldValue);
     }
 
     @SuppressWarnings({ "unchecked" })

src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java

@@ -17,6 +17,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
+import java.util.UUID;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
@@ -501,6 +502,7 @@ private String registerModelGroup() throws IOException, URISyntaxException {
         String modelGroupRegisterRequestBody = Files.readString(
             Path.of(classLoader.getResource("processor/CreateModelGroupRequestBody.json").toURI())
         );
+        modelGroupRegisterRequestBody = modelGroupRegisterRequestBody.replace("<MODEL_GROUP_NAME>", UUID.randomUUID().toString());
         Response modelGroupResponse = makeRequest(
             client(),
             "POST",

src/test/java/org/opensearch/neuralsearch/ml/MLCommonsClientAccessorTests.java

@@ -32,6 +32,8 @@
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.transport.NodeNotConnectedException;
 
+import com.google.common.collect.ImmutableMap;
+
 public class MLCommonsClientAccessorTests extends OpenSearchTestCase {
 
     @Mock
@@ -168,7 +170,9 @@ private ModelTensorOutput createModelTensorOutput(final Float[] output) {
             output,
             new long[] { 1, 2 },
             MLResultDataType.FLOAT64,
-            ByteBuffer.wrap(new byte[12])
+            ByteBuffer.wrap(new byte[12]),
+            "mockResult",
+            ImmutableMap.of("mockKey", "mockValue")
         );
         mlModelTensorList.add(tensor);
         final ModelTensors modelTensors = new ModelTensors(mlModelTensorList);

src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorTests.java

@@ -7,7 +7,14 @@
 
 import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.isA;
+import static org.mockito.Mockito.isNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -350,7 +357,7 @@ public void testProcessResponse_successful() throws Exception {
         Map<String, Object> knnMap = processor.buildMapWithKnnKeyAndOriginalValue(ingestDocument);
 
         List<List<Float>> modelTensorList = createMockVectorResult();
-        processor.appendVectorFieldsToDocument(ingestDocument, knnMap, modelTensorList);
+        processor.setVectorFieldsToDocument(ingestDocument, knnMap, modelTensorList);
         assertEquals(12, ingestDocument.getSourceAndMetadata().size());
     }
 
@@ -398,6 +405,20 @@ public void testBuildVectorOutput_withNestedMap_successful() {
         assertNotNull(actionGamesKnn);
     }
 
+    public void test_updateDocument_appendVectorFieldsToDocument_successful() {
+        Map<String, Object> config = createPlainStringConfiguration();
+        IngestDocument ingestDocument = createPlainIngestDocument();
+        TextEmbeddingProcessor processor = createInstanceWithNestedMapConfiguration(config);
+        Map<String, Object> knnMap = processor.buildMapWithKnnKeyAndOriginalValue(ingestDocument);
+        List<List<Float>> modelTensorList = createMockVectorResult();
+        processor.setVectorFieldsToDocument(ingestDocument, knnMap, modelTensorList);
+
+        List<List<Float>> modelTensorList1 = createMockVectorResult();
+        processor.setVectorFieldsToDocument(ingestDocument, knnMap, modelTensorList1);
+        assertEquals(12, ingestDocument.getSourceAndMetadata().size());
+        assertEquals(2, ((List<?>) ingestDocument.getSourceAndMetadata().get("oriKey6_knn")).size());
+    }
+
     private List<List<Float>> createMockVectorResult() {
         List<List<Float>> modelTensorList = new ArrayList<>();
         List<Float> number1 = ImmutableList.of(1.234f, 2.354f);

src/test/resources/processor/CreateModelGroupRequestBody.json

@@ -1,4 +1,4 @@
 {
-  "name": "test_model_group_public",
+  "name": "<MODEL_GROUP_NAME>",
   "description": "This is a public model group"
 }