diff --git a/programming_runs/generators/generator_utils.py b/programming_runs/generators/generator_utils.py index 1d1b381..ea81fe6 100644 --- a/programming_runs/generators/generator_utils.py +++ b/programming_runs/generators/generator_utils.py @@ -205,4 +205,4 @@ def print_messages(system_message_text: str, user_message_text: str) -> None: def print_generated_func_body(func_body_str: str) -> None: print(f"""--------------------- GENERATED FUNC BODY --------------------- {func_body_str} -------------------------------------------""") \ No newline at end of file +------------------------------------------""") diff --git a/programming_runs/generators/parse.py b/programming_runs/generators/parse.py index c9fc230..562c1e4 100644 --- a/programming_runs/generators/parse.py +++ b/programming_runs/generators/parse.py @@ -5,10 +5,17 @@ from typing import Optional def parse_code_block(string: str, lang: str) -> Optional[str]: code_pattern = fr"```{lang}\n(.*?)\n```" match = re.search(code_pattern, string, re.DOTALL) + if match: return match.group(1) - else: - return parse_first_func(string, lang) + + generic_code_pattern = r"```\n(.*?)\n```" + match = re.search(generic_code_pattern, string, re.DOTALL) + + if match: + return match.group(1) + + return parse_first_func(string, lang) def parse_first_func(code: str, lang: str) -> Optional[str]: