code_extraction.py

import re


Class Extraction():

    def extract_code(llm_output):
        
        code_match = re.search(r"```python\n(.*?)\n```", llm_output, re.DOTALL)
        if code_match:
            return code_match.group(1)
        return None