Fix final answer with extra meaningless symbol

Signed-off-by: Heng Qian <[email protected]>
qianheng-aws · Jul 18, 2024 · 51322d3 · 51322d3
1 parent 007b914
commit 51322d3
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java
@@ -295,7 +295,7 @@ public static void parseThoughtResponse(Map<String, String> modelOutput, String
     public static String extractFinalAnswer(String text) {
         String result = null;
         if (text.contains("\"final_answer\"")) {
-            String pattern = "\"final_answer\"\\s*:\\s*\"(.*?)$";
+            String pattern = "\"final_answer\"\\s*:\\s*\"(.*)\"";
             Pattern jsonBlockPattern = Pattern.compile(pattern, Pattern.DOTALL);
             Matcher jsonBlockMatcher = jsonBlockPattern.matcher(text);
             if (jsonBlockMatcher.find()) {

diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/AgentUtilsTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/AgentUtilsTest.java
@@ -75,6 +75,10 @@ public class AgentUtilsTest {
         + "\"thought\": \"Now I know the final answer\",\n  "
         + "\"final_answer\": \"PPLTool generates such query ```json source=iris_data | fields petal_length_in_cm,petal_width_in_cm | kmeans centroids=3 ```.\"\n}\n```";
 
+    private String responseForFinalAnswerWithMultilines = "---------------------```json\n{\n  "
+        + "\"thought\": \"Now I know the final answer\",\n  "
+        + "\"final_answer\": \"PPLTool generates such query \n```json source=iris_data | fields petal_length_in_cm,petal_width_in_cm | kmeans centroids=3 ```.\"\n}\n```";
+
     private String wrongResponseForAction = "---------------------```json\n{\n  "
         + "\"thought\": \"Let's try VectorDBTool\",\n  "
         + "\"action\": \"After checking online weather forecasts, it looks like tomorrow will be sunny with a high of 25 degrees Celsius.\"\n}\n```";
@@ -120,7 +124,7 @@ public void setup() {
                 THOUGHT,
                 "Unfortunately the tools did not provide the weather forecast directly. Let me check online sources:",
                 FINAL_ANSWER,
-                "After checking online weather forecasts, it looks like tomorrow will be sunny with a high of 25 degrees Celsius.\"\n}\n```"
+                "After checking online weather forecasts, it looks like tomorrow will be sunny with a high of 25 degrees Celsius."
             );
         llmResponseExpectedParseResults.put(responseForFinalAnswerInvalidJson, responseForFinalAnswerExpectedResultExpectedResult);
         Map responseForFinalAnswerWithJsonExpectedResultExpectedResult = Map
@@ -144,6 +148,15 @@ public void setup() {
             );
         llmResponseExpectedParseResults.put(wrongResponseForAction, wrongResponseForActionExpectedResultExpectedResult);
 
+        Map responseForFinalAnswerWithMultilinesExpectedResult = Map
+            .of(
+                THOUGHT,
+                "Now I know the final answer",
+                FINAL_ANSWER,
+                "PPLTool generates such query \n```json source=iris_data | fields petal_length_in_cm,petal_width_in_cm | kmeans centroids=3 ```."
+            );
+        llmResponseExpectedParseResults.put(responseForFinalAnswerWithMultilines, responseForFinalAnswerWithMultilinesExpectedResult);
+
     }
 
     @Test
@@ -442,7 +455,7 @@ public void testExtractMethods_FinalAnswer() {
             Assert.assertNull(actionInput);
             Assert
                 .assertEquals(
-                    "After checking online weather forecasts, it looks like tomorrow will be sunny with a high of 25 degrees Celsius.\"\n}\n```",
+                    "After checking online weather forecasts, it looks like tomorrow will be sunny with a high of 25 degrees Celsius.",
                     finalAnswer
                 );
         }