attempt to fix keyerror:7 on global state object

The additional_inputs parameter in gr.ChatInterface doesn't directly support gr.State objects the way it is being used
Instead, we need to encapsulate state and functions within a class

app.py

@@ -15,132 +15,105 @@ import tempfile
 # Constants
 LLM_MODEL = "gemini-1.5-flash"
 EMBEDDING_MODEL = "BAAI/bge-large-en-v1.5"
-CHROMA_DB_PATH = tempfile.gettempdir()  + "/chroma_db"
+CHROMA_DB_PATH = tempfile.gettempdir() + "/chroma_db"
 
-# Set the Google API key from environment variables
-# This is the recommended way to handle secrets in Hugging Face Spaces
-if "GOOGLE_API_KEY" not in os.environ:
-    gr.Error("Please set the GOOGLE_API_KEY environment variable in your Hugging Face Space secrets.")
-else:
-    os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")
-
-# Global state to hold session data
-class SessionState:
+class PDFChatbot:
     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)
-
-    def is_db_ready(self):
-        return self.db is not None
+        self.state = SessionState()
+
+    def process_pdf(self, pdf_file):
+        try:
+            if self.state.is_db_ready():
+                return (
+                    gr.update(interactive=False),
+                    gr.update(visible=True)
+                )
+
+            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 (
+                    gr.update(interactive=True),
+                    gr.update(visible=False)
+                )
+
+            self.state = SessionState()
+            doc = fitz.open(pdf_file.name)
+            text = ""
+            for page in doc:
+                text += page.get_text()
+            doc.close()
+
+            text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+            docs = text_splitter.create_documents([text])
+
+            embeddings = GoogleGenerativeAIEmbeddings(model=EMBEDDING_MODEL)
+            self.state.db = Chroma.from_documents(
+                documents=docs,
+                embedding=embeddings,
+                persist_directory=self.state.vector_store_path
+            )
 
-# Helper function to generate a new session state
-def new_session():
-    return SessionState()
+            gr.Info("PDF processed successfully! You can now ask questions about the document.")
 
-# Function to handle PDF upload and ingestion
-def process_pdf(pdf_file, state):
-    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(visible=True),
-                state
+                gr.update(visible=True)
             )
-
-        # File size validation
-        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.")
+        except Exception as e:
+            if os.path.exists(self.state.vector_store_path):
+                shutil.rmtree(self.state.vector_store_path)
+            gr.Error(f"An error occurred: {str(e)}")
             return (
                 gr.update(interactive=True),
-                gr.update(visible=False),
-                state
+                gr.update(visible=False)
             )
 
-        # If a file is uploaded, a new session should be started
-        # The new session object is returned to update the state
-        new_state = new_session()
-
-        # Extract text from the PDF using PyMuPDF (fitz)
-        doc = fitz.open(pdf_file.name)
-        text = ""
-        for page in doc:
-            text += doc.get_text()
-        doc.close()
-
-        # Split text into chunks
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        docs = text_splitter.create_documents([text])
-
-        # Create a ChromaDB vector store from the documents
-        embeddings = GoogleGenerativeAIEmbeddings(model=EMBEDDING_MODEL)
-        new_state.db = Chroma.from_documents(
-            documents=docs,
-            embedding=embeddings,
-            persist_directory=new_state.vector_store_path
+    def chat_with_pdf(self, message, history):
+        if not self.state.is_db_ready():
+            yield "Please upload a PDF first to begin the conversation."
+            return
+
+        retriever = self.state.db.as_retriever()
+        llm = ChatGoogleGenerativeAI(model=LLM_MODEL, temperature=0.7)
+
+        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"],
         )
 
-        gr.Info("PDF processed successfully! You can now ask questions about the document.")
-
-        # Return updates to the UI components and the new state
-        return (
-            gr.update(interactive=False),
-            gr.update(visible=True),
-            new_state
-        )
-    except Exception as e:
-        # Clean up the directory in case of an error
-        if state and os.path.exists(state.vector_store_path):
-            shutil.rmtree(state.vector_store_path)
-        gr.Error(f"An error occurred: {str(e)}")
-        # Re-enable the file upload in case of error
-        return (
-            gr.update(interactive=True),
-            gr.update(visible=False),
-            state
+        rag_chain = (
+            {"context": retriever, "question": RunnablePassthrough()}
+            | prompt_template
+            | llm
+            | StrOutputParser()
         )
 
-# Function to handle user queries
-def chat_with_pdf(message, history, state):
-    # Add a defensive check for the state object itself
-    if not state or not state.is_db_ready():
-        yield "Please upload a PDF first to begin the conversation."
-        return
-
-    # Use the ChromaDB instance from the session state
-    retriever = state.db.as_retriever()
-    
-    # Set up the RAG chain
-    llm = ChatGoogleGenerativeAI(model=LLM_MODEL, temperature=0.7)
-    
-    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"],
-    )
+        response = rag_chain.invoke(message)
+        yield response
 
-    rag_chain = (
-        {"context": retriever, "question": RunnablePassthrough()}
-        | prompt_template
-        | llm
-        | StrOutputParser()
-    )
+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)
 
-    response = rag_chain.invoke(message)
-    yield response
+    def is_db_ready(self):
+        return self.db is not None
+
+# Set the Google API key from environment variables
+if "GOOGLE_API_KEY" not in os.environ:
+    raise Exception("Please set the GOOGLE_API_KEY environment variable.")
 
-# Gradio Interface
 with gr.Blocks(title="PDF Chatbot") as demo:
-    # Corrected: Initialize gr.State with the object returned by the function
-    state = gr.State(new_session())
+    chatbot = PDFChatbot()
 
     gr.Markdown(
         """
@@ -148,34 +121,28 @@ with gr.Blocks(title="PDF Chatbot") as demo:
         Upload a PDF to start a conversation with your document.
         """
     )
-    
+
     with gr.Row():
         file_upload_input = gr.File(
             file_types=[".pdf"],
             label="Upload your PDF document",
             interactive=True
         )
-    
-    # Use gr.ChatInterface as a top-level component that wraps the chat logic
+
     chat_interface = gr.ChatInterface(
-        fn=chat_with_pdf,
+        fn=chatbot.chat_with_pdf,
         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",
-        # Fix: Add the state as an additional input to the ChatInterface
-        additional_inputs=[state]
+        visible=False
     )
 
-    # Initially hide the chat interface until a file is processed
-    chat_interface.visible = False
-    
-    # Event handler for file upload
     file_upload_input.upload(
-        fn=process_pdf,
-        inputs=[file_upload_input, state],
-        outputs=[file_upload_input, chat_interface, state]
+        fn=chatbot.process_pdf,
+        inputs=[file_upload_input],
+        outputs=[file_upload_input, chat_interface]
     )
 
-demo.launch()
+demo.launch()