From a31728253090fa669f89d80e50e47a1b82018b23 Mon Sep 17 00:00:00 2001
From: Marwan Zouinkhi <zouinkhi.marwan@gmail.com>
Date: Thu, 15 Feb 2024 18:48:49 -0500
Subject: [PATCH] set scripts for gpt

---
 .github/workflows/add_docstring.py     | 111 +++++++++++++++++++++++++
 .github/workflows/requirements.txt     |   4 +
 .github/workflows/run_add_docstring.sh |   6 ++
 3 files changed, 121 insertions(+)
 create mode 100644 .github/workflows/add_docstring.py
 create mode 100644 .github/workflows/requirements.txt
 create mode 100644 .github/workflows/run_add_docstring.sh

diff --git a/.github/workflows/add_docstring.py b/.github/workflows/add_docstring.py
new file mode 100644
index 000000000..bdcf3bbb4
--- /dev/null
+++ b/.github/workflows/add_docstring.py
@@ -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)
\ No newline at end of file
diff --git a/.github/workflows/requirements.txt b/.github/workflows/requirements.txt
new file mode 100644
index 000000000..8596d447f
--- /dev/null
+++ b/.github/workflows/requirements.txt
@@ -0,0 +1,4 @@
+openai
+redbaron
+autoflake
+black
\ No newline at end of file
diff --git a/.github/workflows/run_add_docstring.sh b/.github/workflows/run_add_docstring.sh
new file mode 100644
index 000000000..40ad509a4
--- /dev/null
+++ b/.github/workflows/run_add_docstring.sh
@@ -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
\ No newline at end of file