changes app

app.py

@@ -14,22 +14,8 @@ from utils.vectordb.vector_db import VectorDb
 CONFIG_PATH = os.path.join(current_dir,'config.yaml')
 PERSIST_DIRECTORY = os.path.join(current_dir,f"data/my-vector-db") # changed to current_dir
 
-#class ChatState:
-#    def __init__(self):
-#        self.conversation = None
-#        self.chat_history = []
-#        self.show_sources = True
-#        self.sources_history = []
-#        self.vectorstore = None
-#        self.input_disabled = True
-#        self.document_retrieval = None
-
 chat_history = gr.State()
 chat_history = []
-vectorstore = gr.State()
-document_retrieval = gr.State()
-
-document_retrieval = DocumentRetrieval()
 
 def handle_userinput(user_question, conversation):
     if user_question:
@@ -48,31 +34,25 @@ def handle_userinput(user_question, conversation):
         return "An error occurred", ""
     #return chat_history, "" #, state.sources_history
 
-def process_documents(files, conversation, save_location=None):
+def process_documents(files, document_retrieval, vectorstore, conversation, save_location=None):
     try:
-        #for doc in files:
+        document_retrieval = DocumentRetrieval()
         _, _, text_chunks = parse_doc_universal(doc=files)
         print(text_chunks)
-        #text_chunks = chat_state.document_retrieval.parse_doc(files)
         embeddings = document_retrieval.load_embedding_model()
         collection_name = 'ekr_default_collection' if not config['prod_mode'] else None
         vectorstore = document_retrieval.create_vector_store(text_chunks, embeddings, output_db=save_location, collection_name=collection_name)
-        #vectorstore = vectorstore
         document_retrieval.init_retriever(vectorstore)
         conversation = document_retrieval.get_qa_retrieval_chain()
         #input_disabled = False
-        return conversation, "Complete! You can now ask questions."
+        return conversation, vectorstore, document_retrieval, "Complete! You can now ask questions."
     except Exception as e:
-        return conversation, f"An error occurred while processing: {str(e)}"
+        return conversation, vectorstore, document_retrieval, f"An error occurred while processing: {str(e)}"
 
-def reset_conversation():
+def reset_conversation(chat_history):
     chat_history = []
-    #chat_state.sources_history = []
     return chat_history, ""
 
-def show_selection(model):
-    return f"You selected: {model}"
-
 # Read config file
 with open(CONFIG_PATH, 'r') as yaml_file:
     config = yaml.safe_load(yaml_file)
@@ -87,7 +67,9 @@ caution_text = """⚠️ Note: depending on the size of your document, this coul
 """
 
 with gr.Blocks() as demo:
+    vectorstore = gr.State()
     conversation = gr.State()
+    document_retrieval = gr.State()
 
     gr.Markdown("# Enterprise Knowledge Retriever", 
             elem_id="title")
@@ -108,7 +90,7 @@ with gr.Blocks() as demo:
     gr.Markdown(caution_text)
       
     # Preprocessing events
-    process_btn.click(process_documents, inputs=[docs, conversation], outputs=[conversation, setup_output])
+    process_btn.click(process_documents, inputs=[docs, document_retrieval, vectorstore, conversation], outputs=[conversation, vectorstore, document_retrieval, setup_output], concurrency_limit=10)
         #process_save_btn.click(process_documents, inputs=[file_upload, save_location], outputs=setup_output)
         #load_db_btn.click(load_existing_db, inputs=[db_path], outputs=setup_output)
     
@@ -122,8 +104,9 @@ with gr.Blocks() as demo:
 
     # Chatbot events
     #msg.submit(handle_userinput, inputs=[msg], outputs=[chatbot, sources_output])
-    msg.submit(handle_userinput, inputs=[msg, conversation], outputs=[chatbot, msg])
-    clear_btn.click(reset_conversation, outputs=[chatbot,msg])
+    msg.submit(handle_userinput, inputs=[msg, conversation], outputs=[chatbot, msg], queue=False)
+    clear_btn.click(lambda: [None, ""], inputs=None, outputs=[chatbot, msg], queue=False)
+    #clear_btn.click(reset_conversation, inputs=[], outputs=[chatbot,msg])
     #show_sources.change(lambda x: gr.update(visible=x), show_sources, sources_output)
 
 if __name__ == "__main__":