reggpt/llms/llm.py

"""
 /*************************************************************************
 * 
 * CONFIDENTIAL
 * __________________
 * 
 *  Copyright (2023-2024) AI Labs, IronOne Technologies, LLC
 *  All Rights Reserved
 *
 *  Author  : Theekshana Samaradiwakara
 *  Description :Python Backend API to chat with private data  
 *  CreatedDate : 14/11/2023
 *  LastModifiedDate : 18/03/2024
 *************************************************************************/
 """

import os
# import time
import logging
logger = logging.getLogger(__name__)
from dotenv import load_dotenv

# from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain_openai import ChatOpenAI

load_dotenv()

openai_api_key = os.environ.get('OPENAI_API_KEY')
# openai_api_key = "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"
print(f"--- ---- ---- openai_api_key: {openai_api_key}")

verbose = os.environ.get('VERBOSE')

def get_model(model_type):
   
    match model_type:
        case "openai":
            llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0, openai_api_key=openai_api_key)
        case _default:
            # raise exception if model_type is not supported
            msg=f"Model type '{model_type}' is not supported. Please choose a valid one"
            logger.error(msg)
            return Exception(msg)


    logger.info(f"model_type: {model_type} loaded:")
    return llm