Update app.py

app.py

@@ -30,23 +30,26 @@ YOUTUBE_URL = "https://www.youtube.com/watch?v=--khbXchTeE"
 MODEL_NAME  = "gpt-4"
 
 def invoke(openai_api_key, use_rag, prompt):
-    if (os.path.isdir(CHROMA_DIR)):
-        shutil.rmtree(CHROMA_DIR)
-    if (os.path.isdir(YOUTUBE_DIR)):
-        shutil.rmtree(YOUTUBE_DIR)
+#    if (os.path.isdir(CHROMA_DIR)):
+#        shutil.rmtree(CHROMA_DIR)
+#    if (os.path.isdir(YOUTUBE_DIR)):
+#        shutil.rmtree(YOUTUBE_DIR)
     if (use_rag):
-        loader = GenericLoader(YoutubeAudioLoader([YOUTUBE_URL], YOUTUBE_DIR), OpenAIWhisperParser())
-        docs = loader.load()
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size = 1500, chunk_overlap = 150)
-        splits = text_splitter.split_documents(docs)
-        vector_db = Chroma.from_documents(documents = splits, embedding = OpenAIEmbeddings(), persist_directory = CHROMA_DIR)
+        if (os.path.isdir(CHROMA_DIR)):
+            vector_db = Chroma(persist_directory = CHROMA_DIR, embedding_function = OpenAIEmbeddings())
+        else:
+            loader = GenericLoader(YoutubeAudioLoader([YOUTUBE_URL], YOUTUBE_DIR), OpenAIWhisperParser())
+            docs = loader.load()
+            text_splitter = RecursiveCharacterTextSplitter(chunk_size = 1500, chunk_overlap = 150)
+            splits = text_splitter.split_documents(docs)
+            vector_db = Chroma.from_documents(documents = splits, embedding = OpenAIEmbeddings(), persist_directory = CHROMA_DIR)
         llm = ChatOpenAI(model_name = MODEL_NAME, openai_api_key = openai_api_key, temperature = 0)
         qa_chain = RetrievalQA.from_chain_type(llm, retriever = vector_db.as_retriever(search_kwargs = {"k": 3}), return_source_documents = True, chain_type_kwargs = {"prompt": QA_CHAIN_PROMPT})
+        result = qa_chain({"query": prompt})
     else:
-        #vector_db = Chroma(persist_directory = CHROMA_DIR, embedding_function = OpenAIEmbeddings())
         llm = ChatOpenAI(model_name = MODEL_NAME, openai_api_key = openai_api_key, temperature = 0)
         qa_chain = RetrievalQA.from_chain_type(llm, retriever = None, return_source_documents = True, cchain_type_kwargs = {"prompt": QA_CHAIN_PROMPT})
-    result = qa_chain({"query": prompt})
+        result = qa_chain({"query": prompt})
     #print(result)
     return result["result"]