import openai
from dotenv import load_dotenv
import os
import chainlit as cl

load_dotenv()

class ChatOpenAI:
    def __init__(self, model_name: str = "gpt-3.5-turbo"):
        self.model_name = model_name
        self.openai_api_key = os.getenv("OPENAI_API_KEY")
        if self.openai_api_key is None:
            raise ValueError("OPENAI_API_KEY is not set")

    def run(self, messages, text_only: bool = True):
        if not isinstance(messages, list):
            raise ValueError("messages must be a list")

        openai.api_key = self.openai_api_key
        response = openai.ChatCompletion.create(
            model=self.model_name, messages=messages
        )

        if text_only:
            return response.choices[0].message.content

        return response

    async def stream_with_cl_message(self, message_history, chainlit_msg: cl.Message, text_only: bool = True, settings: dict = {}):
        print("streaming with cl message", message_history);
        async for stream_resp in await openai.ChatCompletion.acreate(
            model=self.model_name, messages=message_history, stream=True, **settings
        ):
            token = stream_resp.choices[0]["delta"].get("content", "")
            await chainlit_msg.stream_token(token)