Integrate Code Formatting and Linting Tools to Improve Function Workflow

gpt_engineer/applications/cli/main.py

@@ -283,6 +283,7 @@ def main(
     if improve_mode:
         fileselector = FileSelector(project_path)
         original_files_dict = fileselector.ask_for_files()
+        original_files_dict = store.lint_files_dict(original_files_dict)
         files_dict = handle_improve_mode(prompt, agent, memory, original_files_dict)
         if not prompt_yesno("\nDo you want to apply these changes?"):
             return

gpt_engineer/core/default/file_store.py

@@ -3,6 +3,8 @@
 from pathlib import Path
 from typing import Union
 
+import autopep8
+
 from gpt_engineer.core.files_dict import FilesDict
 
 
@@ -54,3 +56,47 @@ def download(self) -> FilesDict:
                         content = "binary file"
                     files[str(path.relative_to(self.working_dir))] = content
         return FilesDict(files)
+
+    def lint_files_dict(self, files_dict: FilesDict) -> FilesDict:
+        """
+        Lints the given files dictionary. The dictionary keys are file names and the values are the file contents.
+        The function supports linting for Python and JavaScript files.
+
+        Parameters
+        ----------
+        files_dict : dict
+            The dictionary containing file names as keys and file contents as values.
+
+        Returns
+        -------
+        dict
+            The dictionary containing linted file contents.
+        """
+        lint_types = {
+            ".py": lint_python_code,
+            # ".js": lint_javascript_code,
+        }
+        for file_name, file_content in files_dict.items():
+            for file_type, lint_func in lint_types.items():
+                if file_name.endswith(file_type):
+                    files_dict[file_name] = lint_func(file_content)
+                    print(f"Formatted {file_name} with {lint_func.__name__}")
+        return files_dict
+
+
+# Static method to lint Python code
+def lint_python_code(code: str) -> str:
+    """
+    Lints the given Python code using autopep8.
+
+    Parameters
+    ----------
+    code : str
+        The Python code to be linted.
+
+    Returns
+    -------
+    str
+        The linted Python code.
+    """
+    return autopep8.fix_code(code)

gpt_engineer/core/default/steps.py

@@ -282,6 +282,7 @@ def improve(
     messages = [
         SystemMessage(content=setup_sys_prompt_existing_code(preprompts)),
     ]
+
     # Add files as input
     messages.append(HumanMessage(content=f"{files_dict.to_chat()}"))
     messages.append(HumanMessage(content=f"Request: {prompt}"))

pyproject.toml

@@ -34,6 +34,8 @@ tomlkit = "^0.12.4"
 pyperclip = "^1.8.2"
 langchain-anthropic = "^0.1.1"
 regex = "^2023.12.25"
+autopep8 = "1.6.0"
+jsbeautifier = "1.14.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=7.3.1"