add json return, improve prompts

Author: morganpartee

.gitignore

@@ -4,4 +4,4 @@
 *.resp.json
 /dist
 **/__pycache__/**
-.oai.txt
+.oai

codegpt/files.py

@@ -12,8 +12,9 @@ def load_text(filenames):
 def write_text(files, backup=False):
     # If the backup option is specified and the file exists,
     # write the existing file to <filename>.bak
-    for out in files:
-        filename = out['filename']
+    for i, out in enumerate(files):
+        filename = out.get('filename', f"{i}.txt")
+        typer.secho(f"Hmm, didn't find a filename, writing to {filename}", color=typer.colors.MAGENTA)
         if backup and os.path.exists(filename):
             with open(filename, "r") as f_in:
                 with open(f"{filename}.bak", "w") as f_out:

codegpt/gpt_interface.py

@@ -33,36 +33,37 @@ def send_iffy_edit(prompt: str, code: Dict[str, str], clipboard: bool = False, y
 
     for filename, code_string in code.items():
         if not clipboard:
-            full_prompt += f"\n\nFile: {filename}\n"
-        full_prompt += f"```\n{code_string}\n```"
+            full_prompt += f"\n\nfilename:\n> {filename}\n"
+        nl = '\n'
+        full_prompt += f"code:\n{nl.join(f'> {x}' for x in code_string.splitlines())}\n\n"
 
     full_prompt += f"\n\nDo the following: {prompt}"
 
-    if not clipboard:
+    if clipboard:
         full_prompt += dedent("""
-        You MUST output complete files.
         
-        For each file you wish to output (only if you modified it or made a new one),
-        answer in this exact format:
+        Answer in the following format:
         
-        filename:
-        > <the filename to be output>
         explanation:
         > <The changes that you made>
         code:
-        > <code line 1>
-        > <code line 2...>""")
+        > <the code to be output line 1>
+        > <the code to be output, line n...>""")
+
     else:
         full_prompt += dedent("""
-        
-        Answer in the following format:
-        
+
+        You may only output complete files.
+
+        If you add or modify a file, return it in this exact format:
+
+        filename:
+        > <the filename to be output>
         explanation:
         > <The changes that you made>
         code:
-        > <the code to be output line 1>
-        > <the code to be output line 2>
-        > <the code to be output, line n...>""")
+        > <code line 1>
+        > <code line n...>""")
 
     max_tokens = confirm_send(full_prompt, yes=yes)
 

codegpt/main.py

@@ -1,5 +1,6 @@
 import os
 import typer
+import json
 
 from codegpt import prompts
 
@@ -21,23 +22,27 @@ def edit_file(
         ..., help="Instruction to edit the file(s). Keep it short! Wrap with quotes.",
     ),
     filenames: List[Path] = typer.Argument(
-        ..., help="List of filenames to edit. If not provided, will prompt for input.",
+        [], help="List of filenames to edit. If not provided, will prompt for input.",
     ),
     backup: bool = typer.Option(
         False, "--backup", "-b", help="Whether to create a backup of the original file(s).",
     ),
     yes: bool = typer.Option(False, "--yes", "-y", help="Don't ask for confirmation.",),
+    raw_code: str = typer.Option(None, "--raw-code", "-c", help="Raw code to edit. Overrides filenames")
     ):
     """
-    Edit one or more files using codegpt.
+    Do something given some code for context. Asking for documents, queries, etc. should work okay. Edits are iffy, but work a lot of the time.
+    
+    Your code better be in git before you use this.
 
     FILENAMES: list of filenames to edit. If not provided, will prompt for input.
     INSTRUCTION: the instruction to edit the file(s). Keep it short!
     """
-    if not filenames:
-        filenames = typer.prompt("Enter the filenames to edit, separated by spaces").split()
-    code = files.load_text(filenames)
-    result = gpt.send_iffy_edit(instruction, code, yes=yes)
+
+    if not filenames and not raw_code:
+        raise typer.BadParameter("Either FILENAMES or --raw-code (-c) must be provided.")
+    code = {"code": raw_code} if raw_code else files.load_text(filenames)
+    result = gpt.send_iffy_edit(instruction, code, yes=yes, clipboard=bool(raw_code))
     files.write_text(result, backup)
     typer.secho("Done!", color=typer.colors.BRIGHT_BLUE)
 
@@ -46,28 +51,48 @@ def edit_file(
 def quick_edit_file(
     option: str = typer.Argument(..., help=f"{{{'|'.join(prompts.prompts.keys())}}}"),
     filenames: List[str] = typer.Argument(..., help="Enter the filenames to edit, separated by spaces"),
-    backup: bool = typer.Option(False, '--backup', '-b', help="Whether to create a backup of the original file(s)."),
+        backup: bool = typer.Option(
+        False, "--backup", "-b", help="Whether to create a backup of the original file(s).",
+    ),
+    yes: bool = typer.Option(False, "--yes", "-y", help="Don't ask for confirmation.",),
+    raw_code: str = typer.Option(None, "--raw-code", "-c", help="Raw code to edit. Overrides filenames"),
+    json: bool = typer.Option(False, "--json", "-j", help="Output in JSON format"),
     ):
     """
-    Edit a file using codegpt's built in prompts
+    Edit a file using codegpt's built in prompts.
+
+    Arguments for `option`:
+    - comment - Adds or updates comments
+    - varnames - Makes variable names reasonable
+    - ugh - Do anything GPT can to make the code suck less (might break stuff...)
+    - docs - Generate (or update) docs, including README.md
+    - bugs - Comment in code where the bugs are if GPT sees them (iffy)
+    - vulns - Comment in code where the vulns are if GPT sees them (iffy)
     """
     if option not in prompts.prompts:
         raise typer.BadParameter(f"{option} is not a valid option. Must be one of {list(prompts.prompts.keys())}")
-    code = files.load_text(filenames)
-    result = gpt.send_iffy_edit(prompts.prompts[option], code)
+
+    if not filenames and not raw_code:
+        raise typer.BadParameter("Either FILENAMES or --raw-code (-c) must be provided.")
+
+    code = {"code": raw_code} if raw_code else files.load_text(filenames)
+    result = gpt.send_iffy_edit(prompts.prompts[option], code, yes=yes, clipboard=bool(raw_code))
+
+    if json:
+        return json.dumps(result, sort_keys=True, indent=4)
+
     files.write_text(result, backup)
-    typer.secho("done", color=typer.colors.BRIGHT_BLUE)
+    typer.secho("Done!", color=typer.colors.BRIGHT_BLUE)
 
 
-@app.command()
+@app.command("config")
 def config():
     """
     Configuration instructions for the OpenAI secret key for the codegpt CLI.
     """
     # check if the secret key is already set in the environment variables
     if "OPENAI_SECRET_KEY" in os.environ:
-        print("The OPENAI_SECRET_KEY is already set in the environment variables.")
-        return
+        typer.secho("OPENAI_SECRET_KEY is already set in the environment! You probably don't need this.", typer.colors.BRIGHT_BLUE)
     else:
         typer.confirm(
             """

codegpt/prompts.py

@@ -2,6 +2,7 @@
     "comment": "Add or update comments according to the given language's standards. Add or update function, module, or class level comments if they're appropriate.",
     "varnames": "Change variable names, but nothing else, to make the code more readable. For example, instead of using 'x' and 'y', use 'width' and 'height'.",
     "ugh": "Do anything you can to make this code more readable. Add comments, change variable and function names, add whitespace, whatever. Add a readme to explain what the code does and where it could be improved.",
+    "docs": "Generate documentation in markdown format for the given files. Make sure to include a README.md for github. If provided markdown files, update them and use their current structure.",
     "bugs": """Find any bugs you can, note them in comments prefixed with BUG:
 
 Before: