week4/community-contributions/doc_string_exercise/generate_doc_string.py

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()