changes to app.py and document_retrieval.py

app.py

@@ -1,9 +1,7 @@
 import os
 import sys
-import logging
 import yaml
 import gradio as gr
-import time
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
 print(current_dir)
@@ -16,61 +14,61 @@ 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
 
-logging.basicConfig(level=logging.INFO)
-logging.info("Gradio app is running")
+#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
 
-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()
 
-chat_state = ChatState()
+document_retrieval = DocumentRetrieval()
 
-chat_state.document_retrieval = DocumentRetrieval()
-
-def handle_userinput(user_question):
+def handle_userinput(user_question, conversation):
     if user_question:
         try:
-            response_time = time.time()
-            response = chat_state.conversation.invoke({"question": user_question})
-            response_time = time.time() - response_time
-            chat_state.chat_history.append((user_question, response["answer"]))
+            response = conversation.invoke({"question": user_question})
+            chat_history.append((user_question, response["answer"]))
 
             #sources = set([f'{sd.metadata["filename"]}' for sd in response["source_documents"]])
             #sources_text = "\n".join([f"{i+1}. {source}" for i, source in enumerate(sources)])
             #state.sources_history.append(sources_text)
 
-            return chat_state.chat_history, "" #, state.sources_history
+            return chat_history, "" #, state.sources_history
         except Exception as e:
             return f"An error occurred: {str(e)}", "" #, state.sources_history
-    return chat_state.chat_history, "" #, state.sources_history
+    else:
+        return "An error occurred", ""
+    #return chat_history, "" #, state.sources_history
 
-def process_documents(files, save_location=None):
+def process_documents(files, conversation, save_location=None):
     try:
         #for doc in files:
         _, _, text_chunks = parse_doc_universal(doc=files)
         print(text_chunks)
         #text_chunks = chat_state.document_retrieval.parse_doc(files)
-        embeddings = chat_state.document_retrieval.load_embedding_model()
+        embeddings = document_retrieval.load_embedding_model()
         collection_name = 'ekr_default_collection' if not config['prod_mode'] else None
-        vectorstore = chat_state.document_retrieval.create_vector_store(text_chunks, embeddings, output_db=save_location, collection_name=collection_name)
-        chat_state.vectorstore = vectorstore
-        chat_state.document_retrieval.init_retriever(vectorstore)
-        chat_state.conversation = chat_state.document_retrieval.get_qa_retrieval_chain()
-        chat_state.input_disabled = False
-        return "Complete! You can now ask questions."
+        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."
     except Exception as e:
-        return f"An error occurred while processing: {str(e)}"
+        return conversation, f"An error occurred while processing: {str(e)}"
 
 def reset_conversation():
-    chat_state.chat_history = []
+    chat_history = []
     #chat_state.sources_history = []
-    return chat_state.chat_history, ""
+    return chat_history, ""
 
 def show_selection(model):
     return f"You selected: {model}"
@@ -89,7 +87,8 @@ caution_text = """⚠️ Note: depending on the size of your document, this coul
 """
 
 with gr.Blocks() as demo:
-    #gr.Markdown("# SambaNova Analyst Assistant") # title
+    conversation = gr.State()
+
     gr.Markdown("# Enterprise Knowledge Retriever", 
             elem_id="title")
     
@@ -108,8 +107,8 @@ with gr.Blocks() as demo:
     process_btn = gr.Button("🔄 Process")
     gr.Markdown(caution_text)
       
-
-    process_btn.click(process_documents, inputs=[docs], outputs=setup_output, concurrency_limit=10)
+    # Preprocessing events
+    process_btn.click(process_documents, inputs=[docs, conversation], outputs=[conversation, setup_output])
         #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)
     
@@ -117,13 +116,14 @@ with gr.Blocks() as demo:
     gr.Markdown("## 3️⃣ Chat with your document")
     chatbot = gr.Chatbot(label="Chatbot", show_label=True, show_share_button=False, show_copy_button=True, likeable=True)
     msg = gr.Textbox(label="Ask questions about your data", show_label=True, placeholder="Enter your message...")
-    clear = gr.Button("Clear chat")
+    clear_btn = gr.Button("Clear chat")
     #show_sources = gr.Checkbox(label="Show sources", value=True)
     sources_output = gr.Textbox(label="Sources", visible=False)
 
+    # Chatbot events
     #msg.submit(handle_userinput, inputs=[msg], outputs=[chatbot, sources_output])
-    msg.submit(handle_userinput, inputs=[msg], outputs=[chatbot, msg])
-    clear.click(reset_conversation, outputs=[chatbot,msg])
+    msg.submit(handle_userinput, inputs=[msg, conversation], outputs=[chatbot, msg])
+    clear_btn.click(reset_conversation, outputs=[chatbot,msg])
     #show_sources.change(lambda x: gr.update(visible=x), show_sources, sources_output)
 
 if __name__ == "__main__":

src/document_retrieval.py

@@ -21,7 +21,7 @@ repo_dir = os.path.abspath(os.path.join(kit_dir, '..'))
 sys.path.append(kit_dir)
 sys.path.append(repo_dir)
 
-import streamlit as st
+#import streamlit as st
 
 from utils.model_wrappers.api_gateway import APIGateway
 from utils.vectordb.vector_db import VectorDb
@@ -30,7 +30,7 @@ from utils.visual.env_utils import get_wandb_key
 CONFIG_PATH = os.path.join(kit_dir, 'config.yaml')
 PERSIST_DIRECTORY = os.path.join(kit_dir, 'data/my-vector-db')
 
-load_dotenv(os.path.join(kit_dir, '.env'))
+#load_dotenv(os.path.join(kit_dir, '.env'))
 
 
 from utils.parsing.sambaparse import parse_doc_universal
@@ -153,13 +153,15 @@ class DocumentRetrieval:
         return api_info, llm_info, embedding_model_info, retrieval_info, prompts, prod_mode
 
     def set_llm(self):
-        if self.prod_mode:
-            sambanova_api_key = st.session_state.SAMBANOVA_API_KEY
-        else:
-            if 'SAMBANOVA_API_KEY' in st.session_state:
-                sambanova_api_key = os.environ.get('SAMBANOVA_API_KEY') or st.session_state.SAMBANOVA_API_KEY
-            else:
-                sambanova_api_key = os.environ.get('SAMBANOVA_API_KEY')
+        #if self.prod_mode:
+        #    sambanova_api_key = st.session_state.SAMBANOVA_API_KEY
+        #else:
+        #    if 'SAMBANOVA_API_KEY' in st.session_state:
+        #        sambanova_api_key = os.environ.get('SAMBANOVA_API_KEY') or st.session_state.SAMBANOVA_API_KEY
+        #    else:
+        #        sambanova_api_key = os.environ.get('SAMBANOVA_API_KEY')
+        
+        sambanova_api_key = os.environ.get('SAMBANOVA_API_KEY')
 
         llm = APIGateway.load_llm(
             type=self.api_info,