Update app.py

app.py

@@ -30,8 +30,8 @@ RAG_CHAIN_PROMPT = PromptTemplate(input_variables = ["context", "question"],
 CHROMA_DIR  = "/data/chroma"
 YOUTUBE_DIR = "/data/youtube"
 
-#YOUTUBE_URL = "https://www.youtube.com/watch?v=--khbXchTeE"
-YOUTUBE_URL = "https://www.youtube.com/watch?v=Iy1IpvcJH7I&list=PL2yQDdvlhXf9XsB2W76_seM6dJxcE2Pdc&index=2"
+YOUTUBE_URL_01 = "https://www.youtube.com/watch?v=--khbXchTeE"
+YOUTUBE_URL_02 = "https://www.youtube.com/watch?v=Iy1IpvcJH7I&list=PL2yQDdvlhXf9XsB2W76_seM6dJxcE2Pdc&index=2"
 
 MODEL_NAME  = "gpt-4"
 
@@ -40,26 +40,24 @@ def invoke(openai_api_key, use_rag, prompt):
                      openai_api_key = openai_api_key, 
                      temperature = 0)
     if (use_rag):
-#        if (os.path.isdir(CHROMA_DIR)):
-#            vector_db = Chroma(embedding_function = OpenAIEmbeddings(),
-#                               persist_directory = CHROMA_DIR)
-#            print("Load DB")
-#        else:
-            loader = GenericLoader(YoutubeAudioLoader([YOUTUBE_URL], YOUTUBE_DIR), 
-                                   OpenAIWhisperParser())
-            docs = loader.load()
-            text_splitter = RecursiveCharacterTextSplitter(chunk_overlap = 150,
-                                                           chunk_size = 1500)
-            splits = text_splitter.split_documents(docs)
-            vector_db = Chroma.from_documents(documents = splits, 
-                                              embedding = OpenAIEmbeddings(), 
-                                              persist_directory = CHROMA_DIR)
-            print("Make DB")
+        # Document loading, splitting, and storage
+        loader = GenericLoader(YoutubeAudioLoader([YOUTUBE_URL_01,
+                                                   YOUTUBE_URL_02], YOUTUBE_DIR), 
+                               OpenAIWhisperParser())
+        docs = loader.load()
+        text_splitter = RecursiveCharacterTextSplitter(chunk_overlap = 150,
+                                                       chunk_size = 1500)
+        splits = text_splitter.split_documents(docs)
+        vector_db = Chroma.from_documents(documents = splits, 
+                                          embedding = OpenAIEmbeddings(), 
+                                          persist_directory = CHROMA_DIR)
+        # Document retrieval
+        #vector_db = Chroma(embedding_function = OpenAIEmbeddings(),
+        #                   persist_directory = CHROMA_DIR)
         rag_chain = RetrievalQA.from_chain_type(llm, 
                                                 chain_type_kwargs = {"prompt": RAG_CHAIN_PROMPT}, 
                                                 retriever = vector_db.as_retriever(search_kwargs = {"k": 3}), 
                                                 return_source_documents = True)
-        print(os.listdir("/data/chroma/"))
         result = rag_chain({"query": prompt})
         result = result["result"]
     else: