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)