Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import streamlit as st
 from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.vectorstores import FAISS  # Changed to FAISS for in-memory storage
+from langchain.vectorstores import FAISS
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.chat_models import ChatOpenAI
 from langchain.chains import ConversationalRetrievalChain
@@ -39,19 +39,15 @@ def process_pdfs(papers, api_key):
         
     with st.spinner("Processing PDFs..."):
         try:
-            # Create embeddings instance
             embeddings = OpenAIEmbeddings(openai_api_key=api_key)
-            
-            # Process all PDFs
             all_texts = []
+            
             for paper in papers:
-                # Save and load PDF
                 file_path = os.path.join('./uploads', paper.name)
                 os.makedirs('./uploads', exist_ok=True)
                 with open(file_path, "wb") as f:
                     f.write(paper.getbuffer())
                 
-                # Load and split the PDF
                 loader = PyPDFLoader(file_path)
                 documents = loader.load()
                 text_splitter = RecursiveCharacterTextSplitter(
@@ -60,24 +56,22 @@ def process_pdfs(papers, api_key):
                 )
                 texts = text_splitter.split_documents(documents)
                 all_texts.extend(texts)
-                
-                # Cleanup
                 os.remove(file_path)
             
-            # Create vectorstore
             vectorstore = FAISS.from_documents(all_texts, embeddings)
             
-            # Create chain
+            memory = ConversationBufferMemory(
+                memory_key="chat_history",
+                return_messages=True,
+                output_key="answer"
+            )
+            
             st.session_state.chain = ConversationalRetrievalChain.from_llm(
                 llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo", openai_api_key=api_key),
-                retriever=vectorstore.as_retriever(
-                    search_kwargs={"k": 3}
-                ),
-                memory=ConversationBufferMemory(
-                    memory_key="chat_history",
-                    return_messages=True,
-                ),
-                return_source_documents=True,
+                retriever=vectorstore.as_retriever(),
+                memory=memory,
+                return_source_documents=False,
+                chain_type="stuff"
             )
             
             st.success(f"Processed {len(papers)} PDF(s) successfully!")
@@ -90,7 +84,6 @@ def process_pdfs(papers, api_key):
 def main():
     st.set_page_config(page_title="PDF Chat")
     
-    # Sidebar with API key input
     api_key = create_sidebar()
     
     if not api_key:
@@ -99,47 +92,29 @@ def main():
 
     st.title("Chat with PDF")
     
-    # File uploader
     papers = st.file_uploader("Upload PDFs", type=["pdf"], accept_multiple_files=True)
     
-    # Process PDFs button
     if papers:
         if st.button("Process PDFs"):
             process_pdfs(papers, api_key)
     
-    # Display chat messages from history
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
     
-    # Accept user input
     if prompt := st.chat_input("Ask about your PDFs"):
-        # Add user message to chat history
         st.session_state.messages.append({"role": "user", "content": prompt})
         
-        # Display user message
         with st.chat_message("user"):
             st.markdown(prompt)
             
-        # Generate and display assistant response
         with st.chat_message("assistant"):
             if st.session_state.chain is None:
                 response = "Please upload and process a PDF first."
             else:
                 with st.spinner("Thinking..."):
-                    # Get response with source documents
                     result = st.session_state.chain({"question": prompt})
                     response = result["answer"]
-                    
-                    # Optionally show sources
-                    if "source_documents" in result:
-                        sources = result["source_documents"]
-                        if sources:
-                            response += "\n\nSources:"
-                            for i, doc in enumerate(sources, 1):
-                                # Add page numbers if available
-                                page_info = f" (Page {doc.metadata['page'] + 1})" if 'page' in doc.metadata else ""
-                                response += f"\n{i}.{page_info} {doc.page_content[:200]}..."
             
             st.markdown(response)
             st.session_state.messages.append({"role": "assistant", "content": response})