src/model_loading.py

from langchain.chat_models import ChatOpenAI
from environs import Env
import typing


def load_open_ai_llm(model_loading_kwargs: typing.Optional[dict] = None) -> ChatOpenAI:
    """Load the openai language model

    Args:
        model_loading_kwargs (typing.Optional[dict], optional): Keyword arguments for loading the model.
            These are passed to the ChatOpenAI class. Defaults to None.

    Returns:
        ChatOpenAI: ChatOpenAI language model
    """
    env = Env()
    env.read_env()
    openai_api_key = env.str("OPENAI_API_KEY")
    openai_organization = env.str("OPENAI_ORGANIZATION")

    # Create the default model kwargs
    default_model_kwargs = dict(
        model_name="gpt-3.5-turbo",
        temperature=0.0,
        openai_api_key=openai_api_key,
        openai_organization=openai_organization,
    )

    # Update the default model kwargs with the provided model loading kwargs

    if model_loading_kwargs:
        default_model_kwargs.update(model_loading_kwargs)

    # Load the model
    openai_llm = ChatOpenAI(**default_model_kwargs)

    return openai_llm