fix state propagation issue between process_pdf() and chat_withdoc()

app.py

@@ -77,43 +77,22 @@ class SessionState:
     def is_db_ready(self):
         return self.db is not None
 
-async def process_pdf(pdf_file, state: gr.State):
-    gr.Info("Processing PDF, please wait...")
+async def process_pdf(pdf_file, state: SessionState):
+    """Processes the PDF and updates the state object."""
     try:
-        # Check if a PDF has already been processed in this session
-        if state and state.is_db_ready():
-            return (
-                gr.update(interactive=False),
-                gr.update(interactive=True),
-                state
-            )
-            
         file_size_mb = os.path.getsize(pdf_file.name) / (1024 * 1024)
         if file_size_mb >= 75:
-            gr.Error("File size exceeds the 75 MB limit. Please upload a smaller PDF.")
-            # Reset components on error
-            return (
-                gr.update(interactive=True),
-                gr.update(interactive=False),
-                gr.State() # Reset state
-            )
+            raise gr.Error("File size exceeds the 75 MB limit. Please upload a smaller PDF.")
 
         print("Opening PDF file...")
         try:
             doc = fitz.open(pdf_file.name)
             text = ""
-            # CRITICAL FIX: Iterate over pages and get text from each page
             for page in doc:
                 text += page.get_text()
             doc.close()
         except Exception as e:
-            print(f"Error processing PDF document: {str(e)}")
-            gr.Error(f"Error processing PDF document: {str(e)}")
-            return (
-                gr.update(interactive=True),
-                gr.update(interactive=False),
-                gr.State()
-            )
+            raise gr.Error(f"Error processing PDF document: {str(e)}")
 
         print("PDF file opened successfully. Splitting text into chunks...")
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
@@ -122,35 +101,23 @@ async def process_pdf(pdf_file, state: gr.State):
 
         embeddings = GoogleGenerativeAIEmbeddings(model=EMBEDDING_MODEL, google_api_key=google_api_key)
         
-        # Initialize a new session state object
-        new_state = SessionState()
-        
-        new_state.db = await Chroma.afrom_documents(
+        state.db = await Chroma.afrom_documents(
             documents=docs,
             embedding=embeddings,
-            persist_directory=new_state.vector_store_path,
-            collection_name=new_state.session_id
+            persist_directory=state.vector_store_path,
+            collection_name=state.session_id
         )
         print("PDF processed successfully! Database is ready.")
-        gr.Info("PDF processed! You can now ask questions about the document.")
-        
-        return (
-            gr.update(interactive=False),
-            gr.update(interactive=True),
-            new_state
-        )
         
     except Exception as e:
-        if state and os.path.exists(state.vector_store_path):
+        if os.path.exists(state.vector_store_path):
             shutil.rmtree(state.vector_store_path)
-        print(f"An error occurred: {str(e)}")
-        gr.Error(f"An error occurred: {str(e)}")
         
-        return (
-            gr.update(interactive=True),
-            gr.update(interactive=False),
-            gr.State()
-        )
+        if isinstance(e, gr.Error):
+            raise  # Re-raise Gradio errors directly
+        else:
+            raise gr.Error(f"An unexpected error occurred: {str(e)}")
+
 
 async def chat_with_pdf(message, history, state: SessionState):
     print("Chat interface called. Checking if database is ready...")
@@ -226,15 +193,23 @@ with gr.Blocks(title="PDF Chatbot") as demo:
 
     async def process_and_show_chat(file, state):
         gr.Info("Processing your PDF, please wait...")
+        new_state = SessionState()
         try:
-            new_state = SessionState()
             await process_pdf(file, new_state)
             gr.Info("PDF processed successfully! You can now chat with it.")
-            return gr.update(visible=True), gr.update(interactive=False), new_state
-        except Exception:
-            # The exception is already a gr.Error, so it will be displayed in the UI.
-            # We just need to return the correct UI updates.
-            return gr.update(visible=False), gr.update(interactive=True), state
+            return {
+                chat_row: gr.update(visible=True),
+                file_upload_input: gr.update(interactive=False),
+                state: new_state,
+            }
+        except gr.Error as e:
+            # Display the Gradio error message to the user
+            gr.Error(str(e))
+            return {
+                chat_row: gr.update(visible=False),
+                file_upload_input: gr.update(interactive=True),
+                state: state, # Return original state on failure
+            }
 
     file_upload_input.upload(
         fn=process_and_show_chat,
@@ -242,4 +217,4 @@ with gr.Blocks(title="PDF Chatbot") as demo:
         outputs=[chat_row, file_upload_input, state]
     )
 
-demo.launch()
+demo.launch()