from openai import OpenAI

class ChatGPT:

    def __init__(self, api_key):
        self.api_key = api_key
        self.client = OpenAI(
            api_key=api_key
        )
        self.role = {'role': 'system', 'content': 'You are a helpful personal assistent called Steve.'}

    def completion(self, message: str) -> str:
        messages = [
            self.role,
            {'role': 'user', 'content': message}
        ]
        response = self.client.chat.completions.create(
            model='gpt-3.5-turbo',
            messages=messages
        )

        return response.choices[0].message.content
    
    def image(self, prompt: str) -> str:
        response = self.client.images.generate(
            model="dall-e-3",
            prompt=prompt,
            size="1024x1024",
            quality="standard",
            n=1
        )

        return response.data[0].url