Update app.py

app.py

@@ -203,7 +203,7 @@ class PDFQAAssistant:
             Dictionary with the answer and source documents
         """
         if self.conversation_chain is None:
-            return {"answer": "Please load a document first before asking questions."}
+            return {"answer": "Please load a document first before asking questions.", "sources": []}
         
         try:
             result = self.conversation_chain({"question": question})
@@ -224,7 +224,7 @@ class PDFQAAssistant:
             
         except Exception as e:
             st.error(f"Error processing question: {e}")
-            return {"answer": f"Error processing your question: {e}"}
+            return {"answer": f"Error processing your question: {e}", "sources": []}
     
     def clear_memory(self) -> None:
         """Clear the conversation memory."""
@@ -255,14 +255,23 @@ def main():
     with st.sidebar:
         st.header("Settings")
         
-        # Option to use HF token from environment or manual entry
-        use_env_token = st.checkbox("Use HF_TOKEN from environment", value=True)
-        
-        if use_env_token:
-            hf_token = os.environ.get("HF_TOKEN", None)
-            if not hf_token:
-                st.warning("HF_TOKEN not found in environment variables.")
+        # Get HF_TOKEN from secrets or environment
+        if "HF_TOKEN" in st.secrets:
+            hf_token = st.secrets["HF_TOKEN"]
+            token_source = "Using HF_TOKEN from app secrets"
+        elif os.environ.get("HF_TOKEN"):
+            hf_token = os.environ.get("HF_TOKEN")
+            token_source = "Using HF_TOKEN from environment variables"
         else:
+            hf_token = None
+            token_source = "No HF_TOKEN found"
+        
+        st.info(token_source)
+        
+        # Option to manually enter token if needed
+        use_manual_token = st.checkbox("Enter token manually", value=not hf_token)
+        
+        if use_manual_token:
             hf_token = st.text_input("Enter Hugging Face API Token:", type="password")
         
         # Model selection
@@ -291,21 +300,27 @@ def main():
         if uploaded_files:
             process_btn = st.button("Process Documents")
             if process_btn:
-                # Initialize the assistant
-                assistant = PDFQAAssistant(
-                    hf_token=hf_token,
-                    model_name=model_name,
-                    embedding_model_name=embedding_model
-                )
-                
-                # Process each uploaded file
-                for pdf_file in uploaded_files:
-                    file_name = pdf_file.name
-                    st.session_state.file_names.append(file_name)
-                    assistant.process_pdf(pdf_file, file_name)
-                
-                # Store the assistant in session state
-                st.session_state.assistant = assistant
+                if not hf_token:
+                    st.error("Please provide a valid Hugging Face API token.")
+                else:
+                    # Initialize the assistant
+                    try:
+                        assistant = PDFQAAssistant(
+                            hf_token=hf_token,
+                            model_name=model_name,
+                            embedding_model_name=embedding_model
+                        )
+                        
+                        # Process each uploaded file
+                        for pdf_file in uploaded_files:
+                            file_name = pdf_file.name
+                            st.session_state.file_names.append(file_name)
+                            assistant.process_pdf(pdf_file, file_name)
+                        
+                        # Store the assistant in session state
+                        st.session_state.assistant = assistant
+                    except Exception as e:
+                        st.error(f"Error initializing assistant: {e}")
         
         # Document management
         if st.session_state.document_processed:
@@ -352,7 +367,7 @@ def main():
                 st.chat_message("user").write(message["content"])
             else:
                 st.chat_message("assistant").write(message["content"])
-                if "sources" in message:
+                if message.get("sources"):  # Use .get() with default to avoid KeyError
                     with st.expander("View Sources"):
                         for source in message["sources"]:
                             st.write(f"- {source['source']} (chunk {source['chunk']})")
@@ -371,21 +386,25 @@ def main():
             # Get the answer
             with st.chat_message("assistant"):
                 with st.spinner("Thinking..."):
-                    result = st.session_state.assistant.ask(question)
-                    st.write(result["answer"])
-                    
-                    # Show sources if available
-                    if result["sources"]:
-                        with st.expander("View Sources"):
-                            for source in result["sources"]:
-                                st.write(f"- {source['source']} (chunk {source['chunk']})")
-                    
-                    # Add assistant response to chat history
-                    st.session_state.chat_history.append({
-                        "role": "assistant",
-                        "content": result["answer"],
-                        "sources": result["sources"]
-                    })
+                    try:
+                        result = st.session_state.assistant.ask(question)
+                        
+                        st.write(result["answer"])
+                        
+                        # Show sources if available
+                        if result.get("sources"):  # Use .get() with default to avoid KeyError
+                            with st.expander("View Sources"):
+                                for source in result["sources"]:
+                                    st.write(f"- {source['source']} (chunk {source['chunk']})")
+                        
+                        # Add assistant response to chat history
+                        st.session_state.chat_history.append({
+                            "role": "assistant",
+                            "content": result["answer"],
+                            "sources": result.get("sources", [])  # Use .get() with default to avoid KeyError
+                        })
+                    except Exception as e:
+                        st.error(f"Error getting response: {e}")
 
 if __name__ == "__main__":
     main()