diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py index 0ec1397b..8d008e92 100644 --- a/pr_agent/algo/utils.py +++ b/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,) diff --git a/tests/unittest/try_fix_yaml.py b/tests/unittest/test_try_fix_yaml.py similarity index 75% rename from tests/unittest/try_fix_yaml.py rename to tests/unittest/test_try_fix_yaml.py index fe37bb1c..21ba9211 100644 --- a/tests/unittest/try_fix_yaml.py +++ b/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 \ No newline at end of file + assert try_fix_yaml(review_text) is None