Update app.py

app.py

@@ -1,140 +1,107 @@
 import streamlit as st
 import google.generativeai as genai
 from langchain.document_loaders import PyPDFDirectoryLoader
+from dotenv import load_dotenv
 import os
-import shutil
 
-# Configuration
-GOOGLE_API_KEY = st.secrets["GOOGLE_API_KEY"]
+# Load API keys from .env
+load_dotenv()
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+# HF_API_KEY = os.getenv("HF_API_KEY")  # Not used directly here, but loaded if needed
 
 # Page configuration
 st.set_page_config(page_title="Chat with PDFs", page_icon="📚")
 
 def initialize_session_state():
-    """Initialize session state variables"""
-    session_state_vars = {
+    session_vars = {
         "messages": [],
         "loaded_files": False,
         "pdf_content": None,
         "chat": None
     }
-    
-    for var, value in session_state_vars.items():
-        if var not in st.session_state:
-            st.session_state[var] = value
+    for key, value in session_vars.items():
+        if key not in st.session_state:
+            st.session_state[key] = value
 
-def load_pdfs(pdf_folder):
-    """Load PDFs and return their content"""
-    if not os.path.exists(pdf_folder):
-        os.makedirs(pdf_folder)
-        
-    loader = PyPDFDirectoryLoader(pdf_folder)
+def load_pdfs(folder_path):
+    if not os.path.exists(folder_path):
+        os.makedirs(folder_path)
+    loader = PyPDFDirectoryLoader(folder_path)
     documents = loader.load()
-    
-    # Concatenate all documents content
-    content = "\n\n".join([doc.page_content for doc in documents])
-    return content
+    return "\n\n".join(doc.page_content for doc in documents)
 
 def initialize_chat(pdf_content):
-    """Initialize Gemini chat with PDF content"""
     genai.configure(api_key=GOOGLE_API_KEY)
-    
     generation_config = {
         "temperature": 0.7,
         "top_p": 0.95,
         "top_k": 40,
         "max_output_tokens": 8192,
     }
-    
+
     model = genai.GenerativeModel(
         model_name="gemini-1.5-pro",
         generation_config=generation_config,
     )
-    
-    # Start chat with context
-    context_prompt = f"""You are a helpful assistant that answers questions based on the following document content:
-{pdf_content}
-Please use this content to answer user questions. If the answer cannot be found in the content, say so."""
 
+    context_prompt = f"""You are a helpful assistant. Use the following PDF content to answer user questions:
+{pdf_content}
+If the answer isn't in the content, reply accordingly."""
+    
     chat = model.start_chat(history=[])
-    # Send initial context
     chat.send_message(context_prompt)
     return chat
 
 def main():
     initialize_session_state()
-    
     st.title("💬 Chat with PDFs")
-    
-    # Sidebar for PDF upload
+
     with st.sidebar:
-        st.header("Upload Documents")
-        uploaded_files = st.file_uploader(
-            "Upload your PDFs",
-            type=["pdf"],
-            accept_multiple_files=True
-        )
+        st.header("Upload PDFs")
+        uploaded_files = st.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True)
 
         if uploaded_files and not st.session_state.loaded_files:
-            # Create pdfs directory if it doesn't exist
             if not os.path.exists("pdfs"):
                 os.makedirs("pdfs")
-            
-            # Clean up old PDF files
             for file in os.listdir("pdfs"):
                 os.remove(os.path.join("pdfs", file))
-            
-            # Save uploaded files
             for file in uploaded_files:
                 with open(f"pdfs/{file.name}", "wb") as f:
                     f.write(file.getvalue())
-            
-            # Load PDF content
+
             with st.spinner("Processing PDFs..."):
                 try:
-                    pdf_content = load_pdfs("pdfs")
-                    st.session_state.pdf_content = pdf_content
+                    pdf_text = load_pdfs("pdfs")
+                    st.session_state.pdf_content = pdf_text
                     st.session_state.loaded_files = True
-                    
-                    # Initialize chat with content
-                    st.session_state.chat = initialize_chat(pdf_content)
+                    st.session_state.chat = initialize_chat(pdf_text)
                 except Exception as e:
-                    st.error(f"Error processing PDFs: {str(e)}")
+                    st.error(f"Error loading PDFs: {str(e)}")
                     return
 
-    # Main chat interface
     if st.session_state.loaded_files:
-        # Display chat messages
-        for message in st.session_state.messages:
-            with st.chat_message(message["role"]):
-                st.markdown(message["content"])
+        for msg in st.session_state.messages:
+            with st.chat_message(msg["role"]):
+                st.markdown(msg["content"])
 
-        # Chat input
         if prompt := st.chat_input("Ask a question about your PDFs:"):
-            # Add user message to chat history
             st.session_state.messages.append({"role": "user", "content": prompt})
             with st.chat_message("user"):
                 st.markdown(prompt)
 
             with st.chat_message("assistant"):
-                response_placeholder = st.empty()
+                placeholder = st.empty()
                 try:
-                    # Get response from Gemini
                     if not st.session_state.chat:
                         st.session_state.chat = initialize_chat(st.session_state.pdf_content)
-                    
                     response = st.session_state.chat.send_message(prompt)
                     response_text = response.text
-                    
-                    response_placeholder.markdown(response_text)
-                    
-                    # Add assistant response to chat history
+                    placeholder.markdown(response_text)
                     st.session_state.messages.append({"role": "assistant", "content": response_text})
                 except Exception as e:
-                    response_placeholder.error(f"Error generating response: {str(e)}")
-
+                    placeholder.error(f"Error generating response: {str(e)}")
     else:
-        st.info("Please upload PDFs to start chatting.")
+        st.info("Please upload PDFs to begin.")
 
 if __name__ == "__main__":
-    main()
+    main()