import os

from openai import OpenAI
from dotenv import load_dotenv

class OpenAIClient:

	def __init__(self):

		load_dotenv()
		self.client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))

	def text_generation(self, prompt):

		response = self.client.chat.completions.create(
			model = "gpt-3.5-turbo",
			temperature = 1.4,
			messages = [{
				"role": "user",
				"content": prompt
			}]
		)

		return response.choices[0].message.content


	def vision(self, prompt, urls):

		content = []

		content.append({ "type": "text", "text": prompt })

		for url in urls:
			content.append({ "type": "image_url", "image_url": url })

		response = self.client.chat.completions.create(
			model = "gpt-4-turbo",
			messages = [{
				"role": "user",
				"content": content
			}]
		)

		return response.choices[0].message.content


	def image_generation(self, prompt):

		response = self.client.images.generate(
			model = "dall-e-3",
  			prompt = prompt,
  			size = "1024x1024",
  			n = 1
  		)

		return response.data[0].url