utils/chatgpt.py

import random

import openai


class ChatGPTAPI:
    def __init__(self, api_key: str = '', max_input_length: int = 1024):
        openai.api_key = self.load_api_key(api_key)
        self.max_input_length = max_input_length

    @staticmethod
    def load_api_key(api_key: str):
        if not api_key:
            try:
                api_key = open('data/api_key.txt', 'r').read()
            except Exception as e:
                raise Exception(f'ChatGPT Error: No API key provided {e}')

        if '\n' in api_key:
            api_key_list = api_key.strip().split('\n')
            api_key = random.choice(api_key_list)
        return api_key

    def __call__(self, content: str):
        assert isinstance(content, str), 'ChatGPT Error: content must be a string'
        content = content.strip()
        messages = [{'role': 'user', 'content': content}]
        try:
            resp = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=messages
            )
            output: str = resp['choices'][0]['message']['content']
            output = output.strip()
        except Exception as e:
            raise Exception(f'ChatGPT Error: {e}')
        return output


if __name__ == '__main__':
    chatgpt = ChatGPTAPI()
    response = chatgpt('Hello, how are you?')
    print(response)