Skip to content

Commit

Permalink
set scripts for gpt
Browse files Browse the repository at this point in the history
  • Loading branch information
mzouink committed Feb 15, 2024
1 parent aeecb66 commit a317282
Show file tree
Hide file tree
Showing 3 changed files with 121 additions and 0 deletions.
111 changes: 111 additions & 0 deletions .github/workflows/add_docstring.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
# Import necessary libraries
import os
import sys
import time
import subprocess
import openai
from redbaron import RedBaron

# Set OpenAI API key
openai.api_key = os.getenv("OPENAI_API_KEY")

# Set starting prompt and history for OpenAI chatbot
# Modify it according to your use case (this is just an example)
starting_prompt = dict(
{
"role": "system",
"content": "I will send you a code of Python function. You need to analyse the code and return to me a string that I can use as the docstring for that function, so as to improve my documentation. The functions can also be routes of a Web App, handle those cases too. Donot write any explanations, just send me a string that I can use as the docstring. The language style of the docstring should be simple and easy to understand and it should be in Google Style Multi-Line format",
}
)
history = [
starting_prompt,
]


# Define function to add docstring to Python functions
def addDocstring(filePath):
"""
Adds docstring to Python functions using OpenAI API
Args:
filePath (str): Path to the Python file
Returns:
None
"""
currentTime = time.time()

# Open the Python file using RedBaron library
with open(filePath, "r", encoding="utf-8") as file:
code = RedBaron(file.read())

# Loop through all functions in the Python file
for node in code.find_all("def"):
# Check if function already has a docstring
if not node.value[0].type == "string":
# To avoid OpenAI rate limit (only free trial accounts have rate limit, comment the code below if you have a paid account)
# Free trial accounts have a hard cap of 1 request every 20 seconds
if time.time() - currentTime < 20:
# Sleep for remaining time
time.sleep(20 - (time.time() - currentTime) + 1)

# Extract the function code
function_code = node.dumps()

# Send the function code to ChatGPT API for generating docstring (offcourse use GPT4 API if you hace access to it)
response = openai.ChatCompletion.create(
model="gpt-3.5-turbo",
temperature=0.2,
messages=[
*history,
{"role": "user", "content": function_code},
],
)

currentTime = time.time()

# Extract the generated docstring from the OpenAI response
docstring = response.choices[0].message.content

# Remove the quotes from the generated docstring if present
if docstring.startswith('"""') or docstring.startswith("'''"):
docstring = docstring[3:-3]
if docstring.startswith('"'):
docstring = docstring[1:-1]

# Add the function code and generated docstring to history
history.append({"role": "user", "content": function_code})
history.append(
{
"role": "assistant",
"content": docstring,
}
)

# Insert the generated docstring to the Function node
if node.next and node.next.type == "comment":
node.next.insert_after(f'"""\n{docstring}\n"""')
else:
node.value.insert(0, f'"""\n{docstring}\n"""')

# Write the modified Python file back to disk
with open(filePath, "w", encoding="utf-8") as file:
file.write(code.dumps())

# Format the new file with autoflake and black
subprocess.run(
[
"autoflake",
"--in-place",
"--remove-unused-variables",
"--remove-all-unused-imports",
filePath,
]
)
subprocess.run(["black", filePath])


# Run the function if this script is called directly
if __name__ == "__main__":
filePath = sys.argv[1]
addDocstring(filePath)
4 changes: 4 additions & 0 deletions .github/workflows/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
openai
redbaron
autoflake
black
6 changes: 6 additions & 0 deletions .github/workflows/run_add_docstring.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
#!/bin/bash
add_docstring_script=$1
for file in $(find . -name "add_docstring.py" -prune -o -name "*.py" -print)
do
python $add_docstring_script $file
done

0 comments on commit a317282

Please sign in to comment.