from langchain.chat_models import ChatOpenAI
from langchain.llms import OpenAI

from .secret import Secret


class OpneAIWrapper(OpenAI):
    def __init__(
        self,
        openai_api_key: Secret,
        temperature: float = 0,
        max_tokens: int = 2048,
        model_name: str = "text-davinci-003",
    ):
        if max_tokens == 0:
            super(OpneAIWrapper, self).__init__(
                model_name=model_name,
                openai_api_key=openai_api_key,
                temperature=temperature,
            )

        else:
            super(OpneAIWrapper, self).__init__(
                model_name=model_name,
                openai_api_key=openai_api_key,
                temperature=temperature,
                max_tokens=max_tokens,
            )


class ChatOpenAIWrapper(ChatOpenAI):
    def __init__(
        self,
        openai_api_key: Secret,
        temperature: float = 0,
        max_tokens: int = 2048,
        model_name: str = "gpt-3.5-turbo",
    ):
        if max_tokens == 0:
            super(ChatOpenAIWrapper, self).__init__(
                model_name=model_name,
                openai_api_key=openai_api_key,
                temperature=temperature,
            )
        else:
            super(ChatOpenAIWrapper, self).__init__(
                model_name=model_name,
                openai_api_key=openai_api_key,
                temperature=temperature,
                max_tokens=max_tokens,
            )