Update app.py

app.py

@@ -75,7 +75,7 @@ def main(message: str):
             ChatOpenAI(temperature=0),
             vectorstore.as_retriever(),
             memory=memory,
-            return_source_documents=True  # This enables returning the source of the answer
+            return_source_documents=True  # Ensure source documents are returned
         )
 
         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,13 +92,17 @@ def main(message: str):
             # Answer questions about the selected paper
             response = qa_chain({"question": message})
             answer = response["answer"]
-            sources = "\n".join([f"- {doc.metadata['source']}" for doc in response["source_documents"]])
-            
-            # Store only the answer in memory
-            memory.save_context({"question": message}, {"answer": answer})
-            
-            # Send the answer and sources as a message
-            cl.Message(content=f"{answer}\n\nSources:\n{sources}").send()
+            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])
+                answer += f"\n\nSources:\n{sources}"
+
+            cl.Message(content=answer).send()
+
+            # Store only the answer in memory to prevent errors
+            qa_chain.memory.save_context({"question": message}, {"answer": answer})
 
 if __name__ == "__main__":
     cl.run()