Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import openai, os, shutil
+import openai, os
 
 from langchain.chains import RetrievalQA
 from langchain.chat_models import ChatOpenAI
@@ -30,10 +30,7 @@ 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)
+    llm = ChatOpenAI(model_name = MODEL_NAME, openai_api_key = openai_api_key, temperature = 0)
     if (use_rag):
         if (os.path.isdir(CHROMA_DIR)):
             vector_db = Chroma(persist_directory = CHROMA_DIR, embedding_function = OpenAIEmbeddings())
@@ -43,13 +40,13 @@ def invoke(openai_api_key, use_rag, prompt):
             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})
+        rag_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 = rag_chain({"query": prompt})
     else:
-        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})
+        #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})
+        chain = LLMChain(llm = llm)
+        result = chain({"query": prompt})
     #print(result)
     return result["result"]