gpt.py

import os
import time
from openai import OpenAI


class ChatGPT:
    def __init__(self, model_name="gpt-3.5-turbo", configs=None) -> None:
        '''
        model_name: api version of a openai model (https://platform.openai.com/docs/models/overview)
        '''
        self.model_name = model_name
        self.client = OpenAI(
            base_url=os.getenv('BASE_URL') + '/v1', 
            api_key=os.getenv('API_KEY'))

        self.config = {
                'n': 1,
                'temperature': 1,
                'top_p': 1
            }
        if configs is not None:
            self.config = configs
    
    def fit_message(self, msg, system_prompt=None):
        if system_prompt is not None:
            conversation = [
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": msg}
                    ]
        else:
            conversation = [
                        {"role": "user", "content": msg}
                    ]
        return conversation

    def query(self, msg, system_prompt=None, **kwargs):
        start = time.time()

        while True:
            try:
                raw_response = self.client.chat.completions.create(
                            model=self.model_name,
                            messages=self.fit_message(msg, system_prompt),
                            **kwargs)
                self.raw_response = raw_response

                return str(raw_response.choices[0].message.content)
            except Exception as e:
                print(e)

            current = time.time()
            if current - start > 60: return None
            time.sleep(10)
    
    def __call__(self, msg, system_prompt=None):
        return self.query(msg, system_prompt, **self.config)


# Example usage
# chatgpt = ChatGPT()
# print(chatgpt("Who are you?"))