From 3e9b15de15e222d2b2ecf3242b62b0c4b055a6c8 Mon Sep 17 00:00:00 2001
From: Xinyu Wu <57612792+XinyuWuu@users.noreply.github.com>
Date: Fri, 7 Nov 2025 15:30:16 +0800
Subject: [PATCH 1/2] Update utils.py

more robust try_fix_yaml
---
 pr_agent/algo/utils.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py
index 72a0624c..a6f49d28 100644
--- a/pr_agent/algo/utils.py
+++ b/pr_agent/algo/utils.py
@@ -812,14 +812,19 @@ def try_fix_yaml(response_text: str,
             pass
 
     # second fallback - try to extract only range from first ```yaml to the last ```
-    snippet_pattern = r'```yaml([\s\S]*?)```(?=\s*$|")'
+    snippet_pattern = r'```(yaml|yml)?([\s\S]*?)```(?=\s*$|")'
     snippet = re.search(snippet_pattern, '\n'.join(response_text_lines_copy))
     if not snippet:
         snippet = re.search(snippet_pattern, response_text_original) # before we removed the "```"
     if snippet:
         snippet_text = snippet.group()
+        prefix = (
+            '```yaml'
+            if snippet_text.startswith('```yaml')
+            else ('```yml' if snippet_text.startswith('```yml') else '```')
+        )
         try:
-            data = yaml.safe_load(snippet_text.removeprefix('```yaml').rstrip('`'))
+            data = yaml.safe_load(snippet_text.removeprefix(prefix).rstrip('`'))
             get_logger().info(f"Successfully parsed AI prediction after extracting yaml snippet")
             return data
         except:

From 46f889b3146190e9d09e4fb6c9b8e4beb51aabff Mon Sep 17 00:00:00 2001
From: Xinyu Wu <57612792+XinyuWuu@users.noreply.github.com>
Date: Fri, 7 Nov 2025 15:43:28 +0800
Subject: [PATCH 2/2] Update test_try_fix_yaml.py

test extract snippet with prefix "```" or "```yml"
---
 tests/unittest/test_try_fix_yaml.py | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/tests/unittest/test_try_fix_yaml.py b/tests/unittest/test_try_fix_yaml.py
index 03823d7b..591fc247 100644
--- a/tests/unittest/test_try_fix_yaml.py
+++ b/tests/unittest/test_try_fix_yaml.py
@@ -21,16 +21,34 @@ class TestTryFixYaml:
 
     # The function extracts YAML snippet
     def test_extract_snippet(self):
-        review_text = '''\
+        review_text1 = '''\
 Here is the answer in YAML format:
 
 ```yaml
 name: John Smith
 age: 35
 ```
+'''
+        review_text2 = '''\
+Here is the answer in YAML format:
+
+```yml
+name: John Smith
+age: 35
+```
+'''
+        review_text3 = '''\
+Here is the answer in YAML format:
+
+```
+name: John Smith
+age: 35
+```
 '''
         expected_output = {'name': 'John Smith', 'age': 35}
-        assert try_fix_yaml(review_text) == expected_output
+        assert try_fix_yaml(review_text1) == expected_output
+        assert try_fix_yaml(review_text2) == expected_output
+        assert try_fix_yaml(review_text3) == expected_output
 
 
     # The YAML string is empty.