File size: 499 Bytes
901c6a4
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
from langchain.chat_models import ChatOpenAI
from langchain.llms import OpenAI

def get_llm(model_name, model_temperature, api_key, max_tokens=None):
    if model_name == "text-davinci-003":
        return OpenAI(temperature=model_temperature, model_name=model_name, max_tokens=max_tokens,
                      openai_api_key=api_key)
    else:
        return ChatOpenAI(temperature=model_temperature, model_name=model_name, max_tokens=max_tokens,
                          openai_api_key=api_key)