Update app.py

app.py

@@ -13,6 +13,8 @@ from langchain.prompts import PromptTemplate
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
 
+from pymongo import MongoClient
+
 from dotenv import load_dotenv, find_dotenv
 _ = load_dotenv(find_dotenv())
 
@@ -40,7 +42,7 @@ YOUTUBE_URL_3 = "https://www.youtube.com/watch?v=vw-KWfKwvTQ"
 
 MODEL_NAME  = "gpt-4"
 
-def document_storage_chroma():
+def document_loading_splitting():
     # Document loading
     docs = []
     # Load PDF
@@ -59,31 +61,14 @@ def document_storage_chroma():
     text_splitter = RecursiveCharacterTextSplitter(chunk_overlap = 150,
                                                    chunk_size = 1500)
     splits = text_splitter.split_documents(docs)
-    # Document storage
+    return splits
+
+def document_storage_chroma(splits):
     vector_db = Chroma.from_documents(documents = splits, 
                                       embedding = OpenAIEmbeddings(disallowed_special = ()), 
                                       persist_directory = CHROMA_DIR)
 
-def document_storage_mongodb():
-    # Document loading
-    docs = []
-    # Load PDF
-    loader = PyPDFLoader(PDF_URL)
-    docs.extend(loader.load())
-    # Load Web
-    loader = WebBaseLoader(WEB_URL_1)
-    docs.extend(loader.load())
-    # Load YouTube
-    loader = GenericLoader(YoutubeAudioLoader([YOUTUBE_URL_1,
-                                               YOUTUBE_URL_2,
-                                               YOUTUBE_URL_3], YOUTUBE_DIR), 
-                           OpenAIWhisperParser())
-    docs.extend(loader.load())
-    # Document splitting
-    text_splitter = RecursiveCharacterTextSplitter(chunk_overlap = 150,
-                                                   chunk_size = 1500)
-    splits = text_splitter.split_documents(docs)
-    # Document storage
+def document_storage_mongodb(splits):
     vector_db = Chroma.from_documents(documents = splits, 
                                       embedding = OpenAIEmbeddings(disallowed_special = ()), 
                                       persist_directory = CHROMA_DIR)
@@ -115,17 +100,16 @@ def invoke(openai_api_key, rag_option, prompt):
         raise gr.Error("Retrieval Augmented Generation is required.")
     if (prompt == ""):
         raise gr.Error("Prompt is required.")
-    
     try:
         llm = ChatOpenAI(model_name = MODEL_NAME, 
                          openai_api_key = openai_api_key, 
                          temperature = 0)
-
+        #splits = document_loading_splitting()
         if (rag_option == "Chroma"):
-            #document_storage_chroma()
+            #document_storage_chroma(splits)
             result = document_retrieval_chroma(llm, prompt)
         elif (rag_option == "MongoDB"):
-            #document_storage_mongodb()
+            #document_storage_mongodb(splits)
             result = document_retrieval_mongodb(llm, prompt)
         else:
             chain = LLMChain(llm = llm, prompt = LLM_CHAIN_PROMPT)