fix issue of seeing context from previously uploaded doc due to gradio state being shared incorrectly

app.py

@@ -22,101 +22,86 @@ LLM_MODEL = "gemini-1.5-flash"
 EMBEDDING_MODEL = "models/embedding-001"
 CHROMA_DB_PATH = tempfile.gettempdir() + "/chroma_db"
 
-class PDFChatbot:
+class SessionState:
     def __init__(self):
-        self.state = SessionState()
-
-    async def process_pdf(self, pdf_file):
-        try:
-            if self.state.is_db_ready():
-                print("Database is already ready.")
-                return
-
-            file_size_mb = os.path.getsize(pdf_file.name) / (1024 * 1024)
-            if file_size_mb >= 75:
-                print("File size exceeds the 75 MB limit.")
-                gr.Error("File size exceeds the 75 MB limit. Please upload a smaller PDF.")
-                return
-
-            self.state = SessionState()
-            print("Opening PDF file...")
-            try:
-                doc = fitz.open(pdf_file.name)
-                text = ""
-                for page in doc:
-                    text += page.get_text()
-                doc.close()
-            except Exception as e:
-                print(f"Error processing PDF document: {str(e)}")
-                return
-
-            print("PDF file opened successfully. Splitting text into chunks...")
-            text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-            docs = text_splitter.create_documents([text])
-            print("Text split into chunks successfully.")
-
-            embeddings = GoogleGenerativeAIEmbeddings(model=EMBEDDING_MODEL, google_api_key=google_api_key)
-            self.state.db = await Chroma.afrom_documents(
-                documents=docs,
-                embedding=embeddings,
-                persist_directory=self.state.vector_store_path,
-                collection_name=self.state.session_id
-            )
-            print("PDF processed successfully! Database is ready.")
-        except Exception as e:
-            if os.path.exists(self.state.vector_store_path):
-                shutil.rmtree(self.state.vector_store_path)
-            print(f"An error occurred: {str(e)}")
+        self.session_id = str(uuid.uuid4())
+        self.db = None
+        self.vector_store_path = os.path.join(CHROMA_DB_PATH, self.session_id)
 
     def is_db_ready(self):
-        return self.state.db is not None
+        return self.db is not None
 
-    async def chat_with_pdf(self, message, history):
-        print("Chat interface called. Checking if database is ready...")
-        if not self.is_db_ready():
-            print("Database is not ready.")
-            yield "Error: Database not ready."
+async def process_pdf(pdf_file, state: SessionState):
+    try:
+        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.")
             return
 
-        print("Database is ready. Retrieving relevant documents...")
-        retriever = self.state.db.as_retriever()
-        llm = ChatGoogleGenerativeAI(model=LLM_MODEL, temperature=0.7, google_api_key=google_api_key)
-
-        prompt_template = PromptTemplate(
-            template="""
-            You are a helpful assistant for a PDF document.
-            Answer the user's question based on the following context.
-            If you don't know the answer, just say that you don't know, don't try to make up an answer.
-            ----------------
-            Context: {context}
-            Question: {question}
-            """,
-            input_variables=["context", "question"],
-        )
-
-        rag_chain = (
-            {"context": retriever, "question": RunnablePassthrough()}
-            | prompt_template
-            | llm
-            | StrOutputParser()
-        )
+        print("Opening PDF file...")
+        try:
+            doc = fitz.open(pdf_file.name)
+            text = ""
+            for page in doc:
+                text += page.get_text()
+            doc.close()
+        except Exception as e:
+            print(f"Error processing PDF document: {str(e)}")
+            return
 
-        response = await rag_chain.ainvoke(
-            message
+        print("PDF file opened successfully. Splitting text into chunks...")
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        docs = text_splitter.create_documents([text])
+        print("Text split into chunks successfully.")
+
+        embeddings = GoogleGenerativeAIEmbeddings(model=EMBEDDING_MODEL, google_api_key=google_api_key)
+        state.db = await Chroma.afrom_documents(
+            documents=docs,
+            embedding=embeddings,
+            persist_directory=state.vector_store_path,
+            collection_name=state.session_id
         )
-        yield response
+        print("PDF processed successfully! Database is ready.")
+    except Exception as e:
+        if os.path.exists(state.vector_store_path):
+            shutil.rmtree(state.vector_store_path)
+        print(f"An error occurred: {str(e)}")
+
+async def chat_with_pdf(message, history, state: SessionState):
+    print("Chat interface called. Checking if database is ready...")
+    if not state or not state.is_db_ready():
+        print("Database is not ready.")
+        yield "Error: Database not ready. Please upload a PDF first."
+        return
+
+    print("Database is ready. Retrieving relevant documents...")
+    retriever = state.db.as_retriever()
+    llm = ChatGoogleGenerativeAI(model=LLM_MODEL, temperature=0.7, google_api_key=google_api_key)
+
+    prompt_template = PromptTemplate(
+        template="""
+        You are a helpful assistant for a PDF document.
+        Answer the user's question based on the following context.
+        If you don't know the answer, just say that you don't know, don't try to make up an answer.
+        ----------------
+        Context: {context}
+        Question: {question}
+        """,
+        input_variables=["context", "question"],
+    )
 
-class SessionState:
-    def __init__(self):
-        self.session_id = str(uuid.uuid4())
-        self.db = None
-        self.vector_store_path = os.path.join(CHROMA_DB_PATH, self.session_id)
+    rag_chain = (
+        {"context": retriever, "question": RunnablePassthrough()}
+        | prompt_template
+        | llm
+        | StrOutputParser()
+    )
 
-    def is_db_ready(self):
-        return self.db is not None
+    response = await rag_chain.ainvoke(message)
+    yield response
 
 with gr.Blocks(title="PDF Chatbot") as demo:
-    chatbot = PDFChatbot()
+    state = gr.State()
 
     gr.Markdown(
         """
@@ -134,22 +119,25 @@ with gr.Blocks(title="PDF Chatbot") as demo:
 
     with gr.Row(visible=False) as chat_row:
         chat_interface = gr.ChatInterface(
-            fn=chatbot.chat_with_pdf,
+            fn=chat_with_pdf,
+            additional_inputs=[state],
             chatbot=gr.Chatbot(type="messages"),
             textbox=gr.Textbox(placeholder="Type your question here...", scale=7),
             examples=[["What is the main topic of the document?"], ["Summarize the key findings."], ["Who are the authors?"]],
             title="Chat Interface",
-            theme="soft"
+            theme="soft",
+            type="messages"
         )
 
     async def process_and_show_chat(file):
-        await chatbot.process_pdf(file)
-        return gr.update(visible=True), gr.update(interactive=False)
+        new_state = SessionState()
+        await process_pdf(file, new_state)
+        return gr.update(visible=True), gr.update(interactive=False), new_state
 
     file_upload_input.upload(
         fn=process_and_show_chat,
         inputs=[file_upload_input],
-        outputs=[chat_row, file_upload_input]
+        outputs=[chat_row, file_upload_input, state]
     )
 
 demo.launch()