Update app.py

app.py

@@ -40,6 +40,54 @@ YOUTUBE_URL_3 = "https://www.youtube.com/watch?v=vw-KWfKwvTQ"
 
 MODEL_NAME  = "gpt-4"
 
+def document_storage_chroma():
+    # 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
+    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
+    vector_db = Chroma.from_documents(documents = splits, 
+                                      embedding = OpenAIEmbeddings(disallowed_special = ()), 
+                                      persist_directory = CHROMA_DIR)
+
 def document_retrieval_chroma(llm, prompt):
     vector_db = Chroma(embedding_function = OpenAIEmbeddings(),
                        persist_directory = CHROMA_DIR)
@@ -60,10 +108,10 @@ def document_retrieval_mongodb(llm, prompt):
     result = rag_chain({"query": prompt})
     return result["result"]
 
-def invoke(openai_api_key, rag, prompt):
+def invoke(openai_api_key, rag_option, prompt):
     if (openai_api_key == ""):
         raise gr.Error("OpenAI API Key is required.")
-    if (rag is None):
+    if (rag_option is None):
         raise gr.Error("Retrieval Augmented Generation is required.")
     if (prompt == ""):
         raise gr.Error("Prompt is required.")
@@ -73,33 +121,11 @@ def invoke(openai_api_key, rag, prompt):
                          openai_api_key = openai_api_key, 
                          temperature = 0)
 
-        if (rag == "Chroma"):
-            # 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
-            #vector_db = Chroma.from_documents(documents = splits, 
-            #                                  embedding = OpenAIEmbeddings(disallowed_special = ()), 
-            #                                  persist_directory = CHROMA_DIR)
-            # Document retrieval
+        if (rag_option == "Chroma"):
+            #document_storage_chroma()
             result = document_retrieval_chroma(llm, prompt)
-        elif (rag == "MongoDB"):
-            # Document retrieval
+        elif (rag_option == "MongoDB"):
+            #document_storage_mongodb()
             result = document_retrieval_mongodb(llm, prompt)
         else:
             chain = LLMChain(llm = llm, prompt = LLM_CHAIN_PROMPT)