import black
from smolagents import Tool

class CodeFormatterTool(Tool):
    name = "CodeFormatterTool"
    description = """
        This tool formats python code using Black.
        You can avoid formatting errors like "IndentationError" with this tool.
    """

    inputs = {
        "code_string": {
            "type": "string",
            "description": "python code string that can be unformatted",
        }
    }

    output_type = "string"

    def forward(self, code_string: str) -> str:
        """
        Format Python code using Black.
        """
        try:
            formatted_code = black.format_str(code_string, mode=black.Mode())
            return formatted_code
        except black.NothingChanged:
            return code_string
        except Exception as e:
            print(f"An error occurred during formatting: {e}")
            return code_string