import os
from langchain import OpenAI
from langchain.chat_models import ChatOpenAI


def get_llm(llm_name, model_temperature, api_key, max_tokens=256):
    os.environ["OPENAI_API_KEY"] = api_key
    if llm_name == "text-davinci-003":
        return OpenAI(
            temperature=model_temperature, model_name=llm_name, max_tokens=max_tokens
        )
    else:
        return ChatOpenAI(
            temperature=model_temperature, model_name=llm_name, max_tokens=max_tokens
        )