Update app.py

app.py

@@ -8,6 +8,14 @@ from langchain.chains import ConversationalRetrievalChain, ConversationChain
 from langchain.memory import ConversationBufferMemory
 from langchain.document_loaders import PyPDFLoader
 
+# Initialize session state variables
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+if "chain" not in st.session_state:
+    st.session_state.chain = None
+if "processed_pdfs" not in st.session_state:
+    st.session_state.processed_pdfs = False
+
 def create_sidebar():
     with st.sidebar:
         st.title("PDF Chat")
@@ -36,7 +44,6 @@ def save_uploaded_file(uploaded_file, path='./uploads/'):
         f.write(uploaded_file.getbuffer())
     return file_path
 
-@st.cache_data
 def load_texts_from_papers(papers):
     all_texts = []
     for paper in papers:
@@ -57,77 +64,77 @@ def load_texts_from_papers(papers):
             st.error(f"Error processing {paper.name}: {str(e)}")
     return all_texts
 
-@st.cache_resource
-def initialize_vectorstore(api_key):  # Added api_key parameter
+def initialize_vectorstore(api_key):
     embedding = OpenAIEmbeddings(openai_api_key=api_key)
     vectorstore = Chroma(embedding_function=embedding, persist_directory="db")
     return vectorstore
 
+def process_pdfs(papers, api_key):
+    if papers and not st.session_state.processed_pdfs:
+        with st.spinner("Processing PDFs..."):
+            texts = load_texts_from_papers(papers)
+            if texts:
+                vectorstore = initialize_vectorstore(api_key)
+                vectorstore.add_documents(texts)
+                st.session_state.chain = ConversationalRetrievalChain.from_llm(
+                    ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo", openai_api_key=api_key),
+                    vectorstore.as_retriever(),
+                    memory=ConversationBufferMemory(
+                        memory_key="chat_history",
+                        return_messages=True
+                    )
+                )
+                st.session_state.processed_pdfs = True
+                st.success("PDFs processed successfully!")
+            return texts
+    return []
+
+def handle_chat(prompt, texts):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    
+    try:
+        if texts or st.session_state.processed_pdfs:
+            result = st.session_state.chain({"question": prompt})
+            response = result["answer"]
+        else:
+            response = "Please upload a PDF first."
+        
+        st.session_state.messages.append({"role": "assistant", "content": response})
+        
+    except Exception as e:
+        st.error(f"Error: {str(e)}")
+
 def main():
     st.set_page_config(page_title="PDF Chat", layout="wide")
-
-    # Get API key from sidebar
+    
     api_key = create_sidebar()
-
-    st.title("Chat with PDF")
-    papers = st.file_uploader("Upload PDFs", type=["pdf"], accept_multiple_files=True)
-
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-
+    
     if not api_key:
         st.warning("Please enter your OpenAI API key")
         return
 
-    try:
-        vectorstore = initialize_vectorstore(api_key)
-        texts = load_texts_from_papers(papers) if papers else []
-        
-        if texts:
-            vectorstore.add_documents(texts)
-            qa_chain = ConversationalRetrievalChain.from_llm(
-                ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo", openai_api_key=api_key),  # Added api_key here
-                vectorstore.as_retriever(),
-                memory=ConversationBufferMemory(
-                    memory_key="chat_history",
-                    return_messages=True
-                )
-            )
-            st.success("PDF processed successfully!")
-        else:
-            memory = ConversationBufferMemory(memory_key="chat_history")
-            qa_chain = ConversationChain(
-                llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo", openai_api_key=api_key),  # Added api_key here
-                memory=memory
-            )
-        
-        # Chat interface
+    st.title("Chat with PDF")
+    
+    # File uploader
+    papers = st.file_uploader("Upload PDFs", type=["pdf"], accept_multiple_files=True)
+    
+    # Process PDFs
+    texts = process_pdfs(papers, api_key)
+    
+    # Chat interface container
+    chat_container = st.container()
+    
+    with chat_container:
+        # Display chat messages
         for message in st.session_state.messages:
             with st.chat_message(message["role"]):
                 st.markdown(message["content"])
-        
+    
+        # Chat input
         if prompt := st.chat_input("Ask about your PDFs"):
-            st.session_state.messages.append({"role": "user", "content": prompt})
-            with st.chat_message("user"):
-                st.markdown(prompt)
-            
-            with st.chat_message("assistant"):
-                try:
-                    if texts:
-                        result = qa_chain({"question": prompt})
-                        response = result["answer"]
-                    else:
-                        result = qa_chain.predict(input=prompt)
-                        response = result
-                    
-                    st.session_state.messages.append({"role": "assistant", "content": response})
-                    st.markdown(response)
-                    
-                except Exception as e:
-                    st.error(f"Error: {str(e)}")
-                
-    except Exception as e:
-        st.error(f"Error: {str(e)}")
+            handle_chat(prompt, texts)
+            # Force a rerun to display the new message
+            st.rerun()
 
 if __name__ == "__main__":
     main()