def format_solution(code_str: str) -> str:
    """Format the solution to match the format of codex.

    Args:
        code_str (str): The code string to format.

    Returns:
        str: The formatted code string.
    """
    if ('def' in code_str) or code_str.startswith('class'):
        return code_str
    code_arr = code_str.splitlines()
    code = 'def code():\n    ' + '\n    '.join(code_arr)
    return code