update api response

Signed-off-by: Hailong Cui <[email protected]>
Hailong-am · Nov 28, 2023 · c8e2dc6 · c8e2dc6
1 parent f88db95
commit c8e2dc6
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java
@@ -432,14 +432,21 @@ private void runReAct(
                             );
 
                         List<ModelTensors> finalModelTensors = new ArrayList<>();
-                        Map<String, Object> additionalInfoMap = new HashMap<>(additionalInfo);
-                        additionalInfoMap.put("response", finalAnswer);
                         finalModelTensors
                             .add(
                                 ModelTensors
                                     .builder()
                                     .mlModelTensors(
-                                        Arrays.asList(ModelTensor.builder().name("response").dataAsMap(additionalInfoMap).build())
+                                        Arrays
+                                            .asList(
+                                                ModelTensor
+                                                    .builder()
+                                                    .name("response")
+                                                    .dataAsMap(
+                                                        ImmutableMap.of("response", finalAnswer, ADDITIONAL_INFO_FIELD, additionalInfo)
+                                                    )
+                                                    .build()
+                                            )
                                     )
                                     .build()
                             );