Close #2, break whole api, put unsafe stuff in unsafe route

morganpartee · morganpartee · commit b4ac4686c103 · 2022-12-11T09:40:43.000-06:00
diff --git a/README.md b/README.md
@@ -2,6 +2,10 @@
 
 A tool for using GPT just a little quicker. A nearly truly automated footgun. Learn how to revert with git before trying please.
 
+Posting about progress here:
+
+[![Twitter Follow](https://img.shields.io/twitter/follow/_JohnPartee?style=social)](https://twitter.com/_JohnPartee)
+
 # Getting Started
 
 `pip install codegpt`
@@ -15,9 +19,13 @@ Windows users can also use `setx` like:
 
 from an admin console.
 
-Then find a file you hate (Back it up! Don't do it live!) and give it a shot.
+## Be careful! But let's go.
+
+The fun stuff is in the `unsafe` command.
+
+Find a file you hate (Back it up! Don't do it live!) and give it a shot.
 
-`codegpt refactor .\helper.py "Break this up into smaller functions where you can. Add google style docstrings. Feel free to rewrite any code doesn't make sense."`
+`codegpt unsafe edit .\helper.py "Break this up into smaller functions where you can. Add google style docstrings. Feel free to rewrite any code doesn't make sense."`
 
 You'll see something like:
 
@@ -32,9 +40,9 @@ Explanation: The code has been refactored into smaller functions to improve read
 
 Other things to try:
 
-- `codegpt edit` - For editing markdown files, including code blocks. Hello, blog editor!
-- `codegpt varnames` - Changes variable names (and supposed to only be variable names...) to be readable
-- `codegpt comment` - Automatically add comments to a file.
+- `codegpt unsafe edit` - Try it with anything. Markdown blog posts, js, yaml, python, whatever.
+- `codegpt unsafe varnames` - Changes variable names (and supposed to only be variable names...) to be readable
+- `codegpt unsafe comment` - Automatically add comments to a file.
 
 Propose endpoints as issues, I've got a few ideas:
 
diff --git a/codegpt/codegpt.py b/codegpt/codegpt.py
@@ -1,189 +1,51 @@
-import openai
 import os
 import typer
-import nltk
 
-app = typer.Typer()
-
-
-def generate_prompt(refactor_or_edit_instructions, code, language):
-    """Generate a prompt from the given instructions and code.
-
-    Args:
-        refactor_or_edit_instructions (str): Instructions for refactoring or editing the code.
-        code (str): The code to be refactored or edited.
-        language (str): The language of the code.
-
-    Returns:
-        str: The generated prompt.
-    """
-    return f"""
-    {'Refactor' if 'refactor' in refactor_or_edit_instructions.lower() else 'Edit'} the following {language} code: {refactor_or_edit_instructions}
-
-    Please provide an extremely succinct human explanation of the changes made to the code
-    and return the edited code in a new section, delimited by '==='. Don't use '===' other than
-    between the sections (don't remove it if it's present though!), and don't add space between sections.
-
-    Ensure that code is well documented and formatted.
-    {" Use google docstrings and black formatting."if language == "python" else ""}
-
-    You must explain what you did, even if you don't make a change.
-
-    Code:
-    {code}""".strip()
-
-
-def refactor_or_edit(
-    file_path: str,
-    refactor_or_edit_instructions: str,
-    explanation_file: str = None,
-    model: str = "text-davinci-003",
-    language: str = "python",
-    debug: bool = False,
-):
-    """Refactor or edit the given file.
-
-    Args:
-        file_path (str): The path to the file to be refactored or edited.
-        refactor_or_edit_instructions (str): Instructions for refactoring or editing the code.
-        explanation_file (str, optional): The path to the file to save the explanation. Defaults to None.
-        model (str, optional): GPT-3 model to use. Defaults to "text-davinci-003".
-        language (str, optional): The language of the code. Defaults to "python".
-        debug (bool, optional): If True, save the response from the model in a JSON file. Defaults to False.
-    """
-    openai.api_key = os.getenv("OPENAI_API_KEY")
-
-    # open the file and escape the code as a code block
-    with open(file_path, "r") as file:
-        code = f"```{language}\n" + file.read() + "\n```"
-
-    # specify the prompt
-    prompt = generate_prompt(refactor_or_edit_instructions, code, language)
-    tokens = nltk.word_tokenize(prompt)
-
-    #! Yeah this math is BS, closeish though...
-    max_tokens = round(4097 - (7 / 4) * len(tokens))
-
-    typer.confirm(
-        f"This prompt is {len(tokens)} tokens, are you sure you want to continue?\nThe most GPT-3 can return in response is {max_tokens}.",
-        default=True,
-        abort=True,
-    )
-
-    # send the prompt to the model
-    response = openai.Completion.create(
-        max_tokens=max_tokens,
-        engine=model,
-        prompt=prompt,
-        n=1,
-        stop=None,
-        temperature=0.6,
-    )
-
-    if debug:
-        import json
-
-        # write the response
-        with open(file_path + ".resp.json", "w") as file:
-            file.write(json.dumps(response))
-
-    # print the response from the model
-    refactored_code = response["choices"][0]["text"]
-
-    explanation = refactored_code.split("===")[0]
-    refactored_code = "".join(refactored_code.split("===")[1:])
-    print(explanation)
+from refactor import edit_file
 
-    old_file_path = file_path + ".old"
-    os.rename(file_path, old_file_path)
-
-    # write the refactored code to the original file
-    with open(file_path, "w") as file:
-        file.write(refactored_code)
-
-    # write the refactored code to the original file
-    with open(
-        explanation_file if explanation_file else file_path + ".explained.txt", "w"
-    ) as file:
-        file.write(explanation)
+app = typer.Typer()
 
 
 @app.command()
-def refactor(
+def unsafe(
+    task: str,
     file_path: str,
-    refactor_instructions: str,
     explanation_file: str = None,
     model: str = "text-davinci-003",
-    language: str = "python",
     debug: bool = False,
 ):
-    """Refactor the given file according to the given instructions.
-
-    Args:
-        file_path (str): The path to the file to be refactored.
-        refactor_instructions (str): Instructions for refactoring the code.
-        explanation_file (str, optional): The path to the file to save the explanation. Defaults to None.
-        model (str, optional): GPT-3 model to use. Defaults to "text-davinci-003".
-        language (str, optional): The language of the code. Defaults to "python".
-        debug (bool, optional): If True, save the response from the model in a JSON file. Defaults to False.
     """
-    refactor_or_edit(
-        file_path, refactor_instructions, explanation_file, model, language, debug
-    )
+    Refactor the given file according to the specified task.
 
+    :param task: is the most important - The task to perform.
 
-@app.command()
-def varnames(
-    file_path: str,
-    refactor_instructions: str = "In the following code, rename variables as you see appropriate for it to be easier to read. Don't touch any of the code otherwise, other than to update comments.",
-    explanation_file: str = None,
-    model: str = "text-davinci-003",
-    language: str = "python",
-    debug: bool = False,
-):
-    """Refactor the given file to rename variables as appropriate.
+    "varnames" : Changes variable names to be more human readable.
+    "comment" : Adds comments to the code, no other changes
+    "edit" : A raw interface to send a prompt to gpt3. "Break this up into more functions" or something is fun to try!
 
     Args:
+        task (str): The type of refactoring to perform on the file. Can be "varnames", "comment", or "edit".
         file_path (str): The path to the file to be refactored.
-        refactor_instructions (str): Instructions for refactoring the code.
         explanation_file (str, optional): The path to the file to save the explanation. Defaults to None.
         model (str, optional): GPT-3 model to use. Defaults to "text-davinci-003".
-        language (str, optional): The language of the code. Defaults to "python".
         debug (bool, optional): If True, save the response from the model in a JSON file. Defaults to False.
     """
-    refactor_or_edit(
-        file_path, refactor_instructions, explanation_file, model, language, debug
-    )
-
-
-@app.command()
-def comment(
-    file_path: str,
-    refactor_instructions: str = "In the following code, make no code changes but add comments. Keep them succinct, but explain everything you can if it's helpful. Add function or class strings where you can.",
-    explanation_file: str = None,
-    model: str = "text-davinci-003",
-    language: str = "python",
-    debug: bool = False,
-):
-    """Edit the given file to add comments."""
-    refactor_or_edit(
-        file_path, refactor_instructions, explanation_file, model, language, debug
-    )
-
-
-@app.command()
-def edit(
-    file_path: str,
-    edit_instructions: str,
-    explanation_file: str = None,
-    model: str = "text-davinci-003",
-    language: str = "md",
-    debug: bool = False,
-):
-    """A Generic edit option, meant for editing markdown blog posts. Basically refactor with some extra instructions."""
-    refactor_or_edit(
-        file_path, edit_instructions, explanation_file, model, language, debug
-    )
+    # Dictionary of instructions for each refactoring task
+    refactor_tasks = {
+        "varnames": "In the following code, rename variables as you see appropriate for it to be easier to read. Don't touch any of the code otherwise, other than to update comments.",
+        "comment": "In the following code, make no code changes but add comments. Keep them succinct, but explain everything you can if it's helpful. Add function or class strings where you can.",
+        "edit": "",
+    }
+    # If the task is "edit", prompt the user for instructions
+    if task == "edit":
+        refactor_instructions = typer.prompt(
+            "Please provide the instructions for editing the file: "
+        )
+    # Otherwise, look up the instructions in the dictionary
+    else:
+        refactor_instructions = refactor_tasks[task]
+    # Call the edit_file function with the specified instructions
+    edit_file(file_path, refactor_instructions, explanation_file, model, debug)
 
 
 @app.command()
diff --git a/codegpt/refactor.py b/codegpt/refactor.py
@@ -0,0 +1,107 @@
+import nltk
+import openai
+import os
+import typer
+
+
+def edit_file(
+    file_path: str,
+    refactor_or_edit_instructions: str,
+    explanation_file: str = None,
+    model: str = "text-davinci-003",
+    debug: bool = False,
+):
+    """Edit the given file.
+
+    Args:
+        file_path (str): The path to the file to be refactored or edited.
+        refactor_or_edit_instructions (str): Instructions for refactoring or editing the code.
+        explanation_file (str, optional): The path to the file to save the explanation. Defaults to None.
+        model (str, optional): GPT-3 model to use. Defaults to "text-davinci-003".
+        language (str, optional): The language of the code. Defaults to "python".
+        debug (bool, optional): If True, save the response from the model in a JSON file. Defaults to False.
+    """
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+
+    language = file_path.split(".")[-1]
+
+    # open the file and escape the code as a code block
+    with open(file_path, "r") as file:
+        code = f"###{file_path}\n```{language}\n" + file.read() + "\n```"
+
+    # specify the prompt
+    prompt = generate_prompt(refactor_or_edit_instructions, code, language)
+    tokens = nltk.word_tokenize(prompt)
+
+    #! Yeah this math is BS, closeish though...
+    max_tokens = round(4097 - (7 / 4) * len(tokens))
+
+    typer.confirm(
+        f"This prompt is {len(tokens)} tokens, are you sure you want to continue?\nThe most GPT-3 can return in response is {max_tokens}.",
+        default=True,
+        abort=True,
+    )
+
+    # send the prompt to the model
+    response = openai.Completion.create(
+        max_tokens=max_tokens,
+        engine=model,
+        prompt=prompt,
+        n=1,
+        stop=None,
+        temperature=0.6,
+    )
+
+    if debug:
+        import json
+
+        # write the response
+        with open(file_path + ".resp.json", "w") as file:
+            file.write(json.dumps(response))
+
+    # print the response from the model
+    refactored_code = response["choices"][0]["text"]
+
+    explanation = refactored_code.split("===")[0]
+    refactored_code = "".join(refactored_code.split("===")[1:])
+    print(explanation)
+
+    old_file_path = file_path + ".old"
+    os.rename(file_path, old_file_path)
+
+    # write the refactored code to the original file
+    with open(file_path, "w") as file:
+        file.write(refactored_code)
+
+    # write the refactored code to the original file
+    with open(
+        explanation_file if explanation_file else file_path + ".explained.txt", "w"
+    ) as file:
+        file.write(explanation)
+
+
+def generate_prompt(refactor_or_edit_instructions, code, language):
+    """Generate a prompt from the given instructions and code.
+
+    Args:
+        refactor_or_edit_instructions (str): Instructions for refactoring or editing the code.
+        code (str): The code to be refactored or edited.
+        language (str): The language of the code.
+
+    Returns:
+        str: The generated prompt.
+    """
+    return f"""
+    {'Refactor' if 'refactor' in refactor_or_edit_instructions.lower() else 'Edit'} the code below, to: {refactor_or_edit_instructions}
+
+    Write a short summary of the changes made to the code and return the edited code in a new section,
+    delimited by '==='. Don't use '===' other than between the sections, and don't add space between sections.
+
+    You are an expert, ensure that code is technically correct, well documented and formatted.
+    {"Use google docstrings and black formatting."if language == "py" else ""}
+    {"This is likely an article or blog post, maintain the authors tone and voice while you edit." if language == "md" else ""}
+
+    You must explain what you did, even if you don't make a change.
+
+    Code:
+    {code}""".strip()
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "codegpt"
-version = "0.0.7"
+version = "0.0.7-8"
 license = "GPL-3.0-or-later"
 description = "A CLI tool for refactoring Python code using OpenAI's text-davinci-003 model"
 authors = ["John Partee"]
