Update rag_langchain.py

rag_langchain.py

@@ -15,15 +15,13 @@ from langchain.vectorstores import MongoDBAtlasVectorSearch
 
 from pymongo import MongoClient
 
-RAG_CHROMA  = "Chroma"
-
 PDF_URL       = "https://arxiv.org/pdf/2303.08774.pdf"
 WEB_URL       = "https://openai.com/research/gpt-4"
 YOUTUBE_URL_1 = "https://www.youtube.com/watch?v=--khbXchTeE"
 YOUTUBE_URL_2 = "https://www.youtube.com/watch?v=hdhZwyf24mE"
 
-YOUTUBE_DIR = "/data/yt"
 CHROMA_DIR  = "/data/db"
+YOUTUBE_DIR = "/data/yt"
 
 MONGODB_ATLAS_CLUSTER_URI = os.environ["MONGODB_ATLAS_CLUSTER_URI"]
 MONGODB_DB_NAME           = "langchain_db"
@@ -120,16 +118,12 @@ def llm_chain(config, prompt):
     
     return completion, llm_chain, cb
 
-def rag_chain(config, rag_option, prompt):
-    llm = get_llm(config)
-
-    if (rag_option == RAG_CHROMA):
-        db = retrieve_chroma()
-    else:
-        db = retrieve_mongodb()
+def rag_chain(config, prompt):
+    #db = retrieve_chroma()
+    db = retrieve_mongodb()
 
     rag_chain = RetrievalQA.from_chain_type(
-        llm, 
+        get_llm(config), 
         chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT,
                              "verbose": True}, 
         retriever = db.as_retriever(search_kwargs = {"k": config["k"]}),