Update app.py

app.py

@@ -76,6 +76,27 @@ def get_response(history, query):
         app.chat_history.append((query, "I have no information about it. Feed me knowledge, please!"))
         return history, f"I have no information about it. Feed me knowledge, please! Error: {str(e)}"
 
+# Function to get response for the current RAG tab
+def get_response_current(history, query):
+    if app.chain is None:
+        raise gr.Error("The chain has not been built yet. Please ensure the vector database is built before querying.")
+    
+    try:
+        result = app.chain.invoke(
+            {"question": query, "chat_history": app.chat_history}
+        )
+        app.chat_history.append((query, result["answer"]))
+        source_docs = result["source_documents"]
+        source_texts = []
+        for doc in source_docs:
+            source_texts.append(f"Page {doc.metadata['page'] + 1}: {doc.page_content}")
+        source_texts_str = "\n\n".join(source_texts)
+        history[-1] = (history[-1][0], result["answer"])
+        return history, source_texts_str
+    except Exception as e:
+        app.chat_history.append((query, "I have no information about it. Feed me knowledge, please!"))
+        return history, f"I have no information about it. Feed me knowledge, please! Error: {str(e)}"
+
 # Function to render file
 def render_file(file) -> Image.Image:
     doc = fitz.open(file.name)
@@ -112,50 +133,50 @@ def set_api_key(api_key):
 
 # Gradio interface
 with gr.Blocks() as demo:
+    api_key_input = gr.Textbox(label="OpenAI API Key", type="password", placeholder="Enter your OpenAI API Key")
+    api_key_btn = gr.Button("Set API Key")
+    api_key_status = gr.Textbox(value="API Key status", interactive=False)
+
+    api_key_btn.click(
+        fn=set_api_key,
+        inputs=[api_key_input],
+        outputs=[api_key_status]
+    )
+
     with gr.Tab("Inst RAG"):
         with gr.Column():
-            api_key_input = gr.Textbox(label="OpenAI API Key", type="password", placeholder="Enter your OpenAI API Key")
-            api_key_btn = gr.Button("Set API Key")
-            api_key_status = gr.Textbox(value="API Key status", interactive=False)
-
-            api_key_btn.click(
-                fn=set_api_key,
-                inputs=[api_key_input],
-                outputs=[api_key_status]
-            )
-
-            with gr.Tab("Upload PDF"):
+            with gr.Row():
                 btn = gr.UploadButton("📁 Upload a PDF", file_types=[".pdf"])
                 show_img = gr.Image(label="Uploaded PDF")
-                
+
                 btn.upload(
                     fn=purge_chat_and_render_first,
                     inputs=[btn],
                     outputs=[show_img],
                 )
 
-            with gr.Tab("Process PDF"):
+            with gr.Row():
                 process_btn = gr.Button("Process PDF")
                 show_img_processed = gr.Image(label="Processed PDF")
                 process_status = gr.Textbox(label="Processing Status", interactive=False)
-                
+
                 process_btn.click(
                     fn=lambda file: (app.process_file(file), "Processing complete!"),
                     inputs=[btn],
                     outputs=[show_img_processed, process_status],
                 )
 
-            with gr.Tab("Build Vector Database"):
+            with gr.Row():
                 build_vector_btn = gr.Button("Build Vector Database")
                 status_text = gr.Textbox(label="Status", value="", interactive=False)
-                
+
                 build_vector_btn.click(
                     fn=app.build_chain,
                     inputs=[btn],
                     outputs=[status_text],
                 )
 
-            with gr.Tab("Chat"):
+            with gr.Row():
                 chatbot = gr.Chatbot(elem_id="chatbot")
                 txt = gr.Textbox(
                     show_label=False,
@@ -193,9 +214,6 @@ with gr.Blocks() as demo:
             refresh_btn_current = gr.Button("Refresh Chat", scale=1)
             source_texts_output_current = gr.Textbox(label="Source Texts", interactive=False)
 
-        def get_response_current(history, query):
-            return get_response(history, query, open("track_training.pdf", 'rb'))
-
         submit_btn_current.click(
             fn=add_text,
             inputs=[chatbot_current, txt_current],