Merge pull request #537 from koid/feature/ignore-header-description-in-ai-response

Enhancement of AI Response Parsing Mechanism

pr_agent/algo/utils.py

@@ -343,18 +343,30 @@ def try_fix_yaml(response_text: str) -> dict:
     except:
         get_logger().info(f"Failed to parse AI prediction after adding |-\n")
 
-    # second fallback - try to remove last lines
+    # second fallback - try to extract only range from first ```yaml to ````
+    snippet_pattern = r'```(yaml)?[\s\S]*?```'
+    snippet = re.search(snippet_pattern, '\n'.join(response_text_lines_copy))
+    if snippet:
+        snippet_text = snippet.group()
+        try:
+            data = yaml.safe_load(snippet_text.removeprefix('```yaml').rstrip('`'))
+            get_logger().info(f"Successfully parsed AI prediction after extracting yaml snippet")
+            return data
+        except:
+            pass
+
+    # thrid fallback - try to remove last lines
     data = {}
     for i in range(1, len(response_text_lines)):
         response_text_lines_tmp = '\n'.join(response_text_lines[:-i])
         try:
             data = yaml.safe_load(response_text_lines_tmp,)
             get_logger().info(f"Successfully parsed AI prediction after removing {i} lines")
-            break
+            return data
         except:
             pass
     
-    # thrid fallback - try to remove leading and trailing curly brackets
+    # fourth fallback - try to remove leading and trailing curly brackets
     response_text_copy = response_text.strip().rstrip().removeprefix('{').removesuffix('}')
     try:
         data = yaml.safe_load(response_text_copy,)

tests/unittest/test_try_fix_yaml.py

@@ -19,6 +19,19 @@ class TestTryFixYaml:
         expected_output = {"relevant line": "value: 3"}
         assert try_fix_yaml(review_text) == expected_output
 
+    # The function extracts YAML snippet
+    def test_extract_snippet(self):
+        review_text = '''\
+Here is the answer in YAML format:
+
+```yaml
+name: John Smith
+age: 35
+```
+'''
+        expected_output = {'name': 'John Smith', 'age': 35}
+        assert try_fix_yaml(review_text) == expected_output
+
     # The function removes the last line(s) of the YAML string and successfully parses the YAML string.
     def test_remove_last_line(self):
         review_text = "key: value\nextra invalid line\n"
@@ -28,4 +41,4 @@ class TestTryFixYaml:
     # The YAML string is empty.
     def test_empty_yaml_fixed(self):
         review_text = ""
-        assert try_fix_yaml(review_text) is None
+        assert try_fix_yaml(review_text) is None