PEASS-ng/AICoder.py

import argparse
import os
import sys
import string
import random
from typing import List
import openai
import json
import subprocess
import tiktoken
import requests
from github import Github

#########################
#### OPENAI FUNCTIONS ###
#########################

def reportTokens(prompt, model="gpt-4"):
    encoding = tiktoken.encoding_for_model(model)
    print("\033[37m" + str(len(encoding.encode(prompt))) + " tokens\033[0m" + " in prompt: " + "\033[92m" + prompt[:50] + "\033[0m" + ("..." if len(prompt) > 50 else ""))

def write_file(file_path: str, content: str):
    """Write content to a file creating the needed directories first"""
    os.makedirs(os.path.dirname(file_path), exist_ok=True)

    with open(file_path, "w") as file:
        file.write(content)

def delete_file(file_path: str):
    """Delete a file if it exists"""

    if os.path.isfile(file_path):
        os.remove(file_path)

openai_available_functions = {
    "write_file": write_file, "delete_file": delete_file
}

openai_functions = [
    {
        "name": "write_file",
        "description": "Write a file giving the path and the content",
        "parameters": {
            "type": "object",
            "properties": {
                "file_path": {
                    "type": "string",
                    "description": "Path to the file to write",
                },
                "content": {
                    "type": "string",
                    "description": "Content to write in the file",
                },
            },
            "required": ["file_path", "content"],
        },
    },
    {
        "name": "delete_file",
        "description": "Delete a file",
        "parameters": {
            "type": "object",
            "properties": {
                "file_path": {
                    "type": "string",
                    "description": "Path to the file to write",
                }
            },
            "required": ["file_path"],
        },
    }
]


#########################
#### GIT FUNCTIONS ######
#########################


def create_pull_request(branch_name, commit_message, github_token):
    github = Github(github_token)
    repo = github.get_repo(os.environ["GITHUB_REPOSITORY"])

    # Create a new branch
    base_branch = repo.get_branch(repo.default_branch)
    repo.create_git_ref(ref=f"refs/heads/{branch_name}", sha=base_branch.commit.sha)

    # Commit changes to the new branch
    subprocess.run(["git", "checkout", branch_name])
    subprocess.run(["git", "add", "."])
    subprocess.run(["git", "commit", "-m", commit_message])
    subprocess.run(["git", "push", "origin", branch_name])

    # Create a pull request
    pr = repo.create_pull(
        title=commit_message,
        body="Generated by OpenAI Github Action",
        head=branch_name,
        base=repo.default_branch
    )

    return pr.html_url


#########################
#### FILE PROCESSING ####
#########################


def process_file(prompt: str, api_key: str, file_path: str, model: str="gpt-4") -> str:
    with open(file_path, "r") as file:
        file_content = file.read()

    messages = [
        {"role": "system", "content": f"You are a developer and your goal is to generate code. The user will ask you to improve and modify some code. Your response must be a valid JSON with the path of each file to write as keys and the content of the files as values. Several files can be written at the same time."},
        {"role": "user", "content": prompt},
        {"role": "user", "content": f"This is the code from the file '{file_path}':\n\n{file_content}"}
    ]
    openai.api_key = api_key

    reportTokens(f"This is the code from the file '{file_path}':\n\n{file_content}")

    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=0
    )
    response_message = response["choices"][0]["message"]

     # Step 2: check if GPT wanted to call a function
    if response_message.get("function_call"):
        
        function_name = response_message["function_call"]["name"]
        fuction_to_call = openai_available_functions[function_name]
        function_args = json.loads(response_message["function_call"]["arguments"])
        fuction_to_call(**function_args)


def process_folder(prompt: str, api_key: str, folder_path: str, model: str="gpt-4") -> List[str]:
    responses = []
    for root, _, files in os.walk(folder_path):
        for file in files:
            file_path = os.path.join(root, file)
            response = process_file(prompt, api_key, file_path, model)
            responses.append(response)


#########################
#### MAIN FUNCTION ######
#########################


def get_random_string(length):
    # With combination of lower and upper case
    letters = string.ascii_letters
    result_str = ''.join(random.choice(letters) for i in range(length))
    return result_str

def main(prompt: str, api_key: str, file_path: str, github_token: str, model: str="gpt-4"):
    if os.path.isfile(file_path):
        process_file(prompt, api_key, file_path, model)
    elif os.path.isdir(file_path):
        process_folder(prompt, api_key, file_path, model)
    else:
        print("Error: Invalid file path.")
        sys.exit(1)

    try:
        create_pull_request(get_random_string(5), f"Modified {file_path}", github_token)
    except Exception as e:
        print(f"Error: Failed to create pull request. {e}")
        sys.exit(1)

if __name__ == "__main__":
    # Setup the argument parser
    parser = argparse.ArgumentParser()

    # Add arguments for prompt, api_key, file_path and github_token
    parser.add_argument('--prompt', default=None, type=str, help='Input prompt')
    parser.add_argument('--api-key', default=None, type=str, help='Input API key')
    parser.add_argument('--path', default=None, type=str, help='Input file/folder path')
    parser.add_argument('--github-token', default=None, type=str, help='Github token')
    parser.add_argument('--model', default="gpt-4", type=str, help='Model to use')

    # Parse the arguments
    args = parser.parse_args()
    prompt = os.environ.get("INPUT_PROMPT", args.prompt)
    api_key = os.environ.get("INPUT_API_KEY", args.api_key)
    file_path = os.environ.get("INPUT_FILE_PATH", args.path)
    github_token = os.environ.get("GITHUB_TOKEN", args.github_token)
    model = os.environ.get("INPUT_MODEL", args.model)

    if not prompt or not api_key or not file_path:
        print("Error: Missing required inputs.")
        sys.exit(1)
    
    #if not github_token:
    #    print("Error: Missing github token.")
    #    sys.exit(1)

    if os.path.exists(prompt):
        with open(prompt, "r") as file:
            prompt = file.read()
    
    if prompt.startswith("http"):
        prompt = requests.get(prompt).text
    
    main(prompt, api_key, file_path, github_token, model)