Update app.py

app.py

@@ -11,7 +11,7 @@ import os
 # Set your OpenAI API key
 os.environ["OPENAI_API_KEY"] = "sk-proj-vFPqdrr801blzZCRBjztT3BlbkFJJJeQVcc62PA40cQ1S9Zv"
 
-# Initialize global variables
+# Global variables
 selected_paper = None
 qa_chain = None
 papers = []
@@ -28,7 +28,6 @@ def main(message: str):
     global selected_paper, qa_chain, papers, state
 
     if state == "SEARCH":
-        # Search for papers
         search = arxiv.Search(
             query=message,
             max_results=5,
@@ -36,12 +35,11 @@ def main(message: str):
         )
 
         papers = list(search.results())
-        
+
         if not papers:
             cl.Message(content="No papers found. Please try another search query.").send()
             return
 
-        # Create a numbered list of papers with links
         paper_list = "\n".join([f"{i+1}. {paper.title} - {paper.authors[0]}\nLink: {paper.entry_id}" for i, paper in enumerate(papers)])
         cl.Message(content=f"Please select a paper by entering its number:\n\n{paper_list}\n\nEnter the number of the paper you want to select:").send()
         state = "SELECT"
@@ -58,7 +56,6 @@ def main(message: str):
             cl.Message(content="Invalid input. Please enter a number.").send()
             return
 
-        # Download and process the selected paper
         paper_text = selected_paper.summary
 
         # Split the text into chunks
@@ -75,7 +72,7 @@ def main(message: str):
             ChatOpenAI(temperature=0),
             vectorstore.as_retriever(),
             memory=memory,
-            return_source_documents=True  # Ensure source documents are returned
+            return_source_documents=True
         )
 
         cl.Message(content=f"Selected paper: {selected_paper.title}\nLink: {selected_paper.entry_id}\nYou can now ask questions about this paper. Type 'new search' when you want to search for a different paper.").send()
@@ -92,16 +89,16 @@ def main(message: str):
             # Answer questions about the selected paper
             response = qa_chain({"question": message})
             answer = response["answer"]
-            source_documents = response.get("source_documents", [])
 
-            # Create a message with the answer and sources
-            if source_documents:
-                sources = "\n".join([f"- {doc.metadata['source']}" for doc in source_documents])
+            # Handling the sources and formatting the response
+            sources = "\n".join([f"- {doc.metadata['source']}" for doc in response.get("source_documents", [])])
+            if sources:
                 answer += f"\n\nSources:\n{sources}"
 
+            # Send the response with sources
             cl.Message(content=answer).send()
 
-            # Store only the answer in memory to prevent errors
+            # Store the chat history in memory without storing the sources
             qa_chain.memory.save_context({"question": message}, {"answer": answer})
 
 if __name__ == "__main__":