Update rag.py

rag.py

@@ -52,7 +52,8 @@ def document_loading():
     #loader = GenericLoader(YoutubeAudioLoader([YOUTUBE_URL], YOUTUBE_DIR), 
     #                       OpenAIWhisperParser())
     #docs.extend(loader.load())
-
+    print("### Load YouTube")
+    
     return docs
 
 def document_splitting(config, docs):
@@ -63,23 +64,25 @@ def document_splitting(config, docs):
     return text_splitter.split_documents(docs)
     
 def document_storage_chroma(chunks):
-    print("### Store")
+    print("### Store Chroma")
     Chroma.from_documents(documents = chunks, 
                           embedding = OpenAIEmbeddings(disallowed_special = ()), 
                           persist_directory = CHROMA_DIR)
 
 def document_storage_mongodb(chunks):
-    print("### Store")    
+    print("### Store MongoDB")
     MongoDBAtlasVectorSearch.from_documents(documents = chunks,
                                             embedding = OpenAIEmbeddings(disallowed_special = ()),
                                             collection = collection,
                                             index_name = MONGODB_INDEX_NAME)
 
 def document_retrieval_chroma():
+    print("### Retrieve Chroma")
     return Chroma(embedding_function = OpenAIEmbeddings(disallowed_special = ()),
                   persist_directory = CHROMA_DIR)
 
 def document_retrieval_mongodb():
+    print("### Retrieve MongoDB")
     return MongoDBAtlasVectorSearch.from_connection_string(MONGODB_ATLAS_CLUSTER_URI,
                                                            MONGODB_DB_NAME + "." + MONGODB_COLLECTION_NAME,
                                                            OpenAIEmbeddings(disallowed_special = ()),
@@ -118,7 +121,6 @@ def rag_chain(config, openai_api_key, rag_option, prompt):
     rag_chain = RetrievalQA.from_chain_type(llm, 
                                             chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT}, 
                                             retriever = db.as_retriever(search_kwargs = {"k": config["k"]}), 
-                                            include_run_info = True,
                                             return_source_documents = True,
                                             verbose = True)