Function extract_code_blocks

ms_agent/utils/utils.py:701–742

Extract code blocks from the given text. py:a.py Args: text: The text to extract code blocks from. target_filename: The filename target to extract. file_wrapper: The file wrapper Returns: Tuple: 0: The extracted code bl

(
    text: str,
    target_filename: Optional[str] = None,
    file_wrapper: Optional[List[str]] = None,
)

def extract_code_blocks(
    text: str,
    target_filename: Optional[str] = None,
    file_wrapper: Optional[List[str]] = None,
) -> Tuple[List, str]:
    """Extract code blocks from the given text.

    <result>py:a.py
    </result>

    Args:
        text: The text to extract code blocks from.
        target_filename: The filename target to extract.
        file_wrapper: The file wrapper

    Returns:
        Tuple:
            0: The extracted code blocks.
            1: The left content of the input text.
    """
    if not file_wrapper:
        file_wrapper = ['<result>', '</result>']
    assert len(file_wrapper) == 2
    pattern = rf'{file_wrapper[0]}[a-zA-Z]*:([^\n\r`]+)\n(.*?){file_wrapper[1]}'
    matches = re.findall(pattern, text, re.DOTALL)
    result = []

    for filename, code in matches:
        filename = filename.strip()
        if target_filename is None or filename == target_filename:
            result.append({'filename': filename, 'code': code.strip()})

    if target_filename is not None:
        remove_pattern = rf'{file_wrapper[0]}[a-zA-Z]*:{re.escape(target_filename)}\n.*?{file_wrapper[1]}'
    else:
        remove_pattern = pattern

    remaining_text = re.sub(remove_pattern, '', text, flags=re.DOTALL)
    remaining_text = re.sub(r'\n\s*\n\s*\n', '\n\n', remaining_text)
    remaining_text = remaining_text.strip()

    return result, remaining_text


@contextmanager