Spaces:
Sleeping
Sleeping
File size: 2,411 Bytes
5fdb69e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
from argparse import ArgumentParser
import os
from dotenv import load_dotenv
from openai import OpenAI
import anthropic
from utils import add_doc_string, Model, get_system_message
from pathlib import Path
def main():
# get run time arguments
parser = ArgumentParser(
prog='Generate Doc String for an existing functions',
description='Run Doc String for a given file and model',
)
parser.add_argument(
'-fp',
'--file_path',
help='Enter the file path to the script that will be updated with doc strings',
default=None
)
parser.add_argument(
'-llm',
'--llm_model',
help='Choose the LLM model that will create the doc strings',
default='claude'
)
# get run time arguments
args = parser.parse_args()
file_path = Path(args.file_path)
llm_model = args.llm_model
# check for file path
assert file_path.exists(), f"File Path {str(file_path.as_posix())} doesn't exist. Please try again."
# check for value llm values
assert llm_model in ['gpt', 'claude'], (f"Invalid model chosen '{llm_model}'. "
f"Please choose a valid model ('gpt' or 'claude')")
# load keys and environment variables
load_dotenv()
os.environ['OPENAI_API_KEY'] = os.getenv('OPENAI_API_KEY', 'your-key-if-not-using-env')
os.environ['ANTHROPIC_API_KEY'] = os.getenv('ANTHROPIC_API_KEY', 'your-key-if-not-using-env')
os.environ['HF_TOKEN'] = os.getenv('HF_INF_TOKEN', 'your-key-if-not-using-env')
# get system messages
system_message = get_system_message()
# get model info
model_info = {
'gpt': {
'client': OpenAI(),
'model': Model.OPENAI_MODEL.value,
},
'claude': {
'client': anthropic.Anthropic(),
'model': Model.CLAUDE_MODEL.value
}
}
# add standard argumens
model_info[llm_model].update(
{
'file_path': file_path,
'system_message': system_message
}
)
# convert python code to c++ code using open ai
print(f"\nSTARTED | Doc Strings Using {llm_model.upper()} for file {str(file_path)}\n\n")
add_doc_string(**model_info[llm_model])
print(f"\nFINISHED | Doc Strings Using {llm_model.upper()} for file {str(file_path)}\n\n")
if __name__ == '__main__':
main()
|