Update app.py

app.py

@@ -1,6 +1,8 @@
 import chainlit as cl
 import arxiv
-from typing import List  # <— add this
+import requests
+from typing import List
+
 from langchain.chat_models import ChatOpenAI
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
@@ -18,45 +20,49 @@ class ArxivResearchAssistant:
         self.papers: List[arxiv.Result] = []
         self.state = "SEARCH"
 
-        # NEW: modern client that uses HTTPS + retry/backoff
+        # ---- NEW: custom session with UA (no 'user_agent' kwarg) ----
+        sess = requests.Session()
+        sess.headers.update({
+            "User-Agent": f"arxiv-chainlit-app/1.0 (mailto:{os.getenv('CONTACT_EMAIL','noreply@example.com')})"
+        })
+        # If you’re behind a proxy or want requests to use env vars:
+        sess.trust_env = True
+
+        # ArXiv client (retries + small delay)
         self.client = arxiv.Client(
-            page_size=50,
+            page_size=5,
             delay_seconds=3,
             num_retries=3,
-            user_agent="chainlit-arxiv-app/1.0 (mailto:your-email@example.com)"
+            http_session=sess
         )
 
     async def search_papers(self, query: str):
-        query = (query or "").strip()
-        if not query:
-            await cl.Message(content="Please enter a non-empty search query.").send()
-            return None
-
+        # Use arxiv.Search, then fetch with our client to leverage the session/retries
         search = arxiv.Search(
             query=query,
             max_results=5,
             sort_by=arxiv.SortCriterion.Relevance
         )
-
-        # CHANGED: use the client to fetch results (handles HTTPS correctly)
-        self.papers = list(self.client.results(search))
+        try:
+            self.papers = list(self.client.results(search))
+        except Exception as e:
+            await cl.Message(
+                content=f"Error talking to arXiv: {e}\nTry again in a moment or tweak your query."
+            ).send()
+            return None
 
         if not self.papers:
-            await cl.Message(content="No papers found. Please try another search query.").send()
+            await cl.Message(
+                content="No papers found. Please try another search query."
+            ).send()
             return None
 
-        paper_list = "\n".join(
-            [
-                f"{i+1}. {paper.title} - {paper.authors[0]}\nLink: {paper.entry_id}"
-                for i, paper in enumerate(self.papers)
-            ]
-        )
-
+        paper_list = "\n".join([
+            f"{i+1}. {paper.title} - {paper.authors[0]}\nLink: {paper.entry_id}"
+            for i, paper in enumerate(self.papers)
+        ])
         await cl.Message(
-            content=(
-                f"Please select a paper by entering its number:\n\n{paper_list}\n\n"
-                "Enter the number of the paper you want to select:"
-            )
+            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()
         self.state = "SELECT"
         return self.papers
@@ -73,7 +79,11 @@ class ArxivResearchAssistant:
             await cl.Message(content="Invalid input. Please enter a number.").send()
             return None
 
-        paper_text = f"{self.selected_paper.title}\n\n{self.selected_paper.summary}\n\n{self.selected_paper.comment or ''}"
+        paper_text = (
+            f"{self.selected_paper.title}\n\n"
+            f"{self.selected_paper.summary}\n\n"
+            f"{self.selected_paper.comment or ''}"
+        )
 
         text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
         chunks = text_splitter.split_text(paper_text)
@@ -82,13 +92,11 @@ class ArxivResearchAssistant:
         vectorstore = FAISS.from_texts(
             chunks,
             embeddings,
-            metadatas=[
-                {
-                    "title": self.selected_paper.title,
-                    "link": self.selected_paper.entry_id,
-                    "chunk": f"Chunk {i+1}/{len(chunks)}"
-                } for i in range(len(chunks))
-            ]
+            metadatas=[{
+                "title": self.selected_paper.title,
+                "link": self.selected_paper.entry_id,
+                "chunk": f"Chunk {i+1}/{len(chunks)}"
+            } for i in range(len(chunks))]
         )
 
         memory = ConversationBufferMemory(
@@ -107,8 +115,9 @@ class ArxivResearchAssistant:
         await cl.Message(
             content=(
                 f"Selected paper: {self.selected_paper.title}\n"
-                f"Link: {self.selected_paper.entry_id}\n"
-                "You can now ask questions about this paper. Type 'new search' when you want to search for a different paper."
+                f"Link: {self.selected_paper.entry_id}\n\n"
+                f"You can now ask questions about this paper. "
+                f"Type 'new search' when you want to search for a different paper."
             )
         ).send()
         self.state = "QA"
@@ -123,14 +132,11 @@ class ArxivResearchAssistant:
         response = self.qa_chain({"question": message})
         answer = response["answer"]
 
-        sources = "\n".join(
-            [
-                f"- {doc.metadata.get('title', 'Unknown title')} "
-                f"({doc.metadata.get('link', 'No link')}) - "
-                f"{doc.metadata.get('chunk', 'No chunk info')}"
-                for doc in response.get("source_documents", [])
-            ]
-        )
+        sources = "\n".join([
+            f"- {doc.metadata.get('title','Unknown title')} "
+            f"({doc.metadata.get('link','No link')}) - {doc.metadata.get('chunk','No chunk info')}"
+            for doc in response.get("source_documents", [])
+        ])
         if sources:
             answer += f"\n\nSources:\n{sources}"
 
@@ -146,12 +152,10 @@ assistant = ArxivResearchAssistant()
 
 @cl.on_chat_start
 async def start():
-    await cl.Message(
-        content=(
-            "Welcome! This tool helps you search for papers on arXiv, pick one, and ask questions about its content.\n\n"
-            "Please enter a topic to search for on arXiv papers.\n\n"
-        )
-    ).send()
+    await cl.Message(content=(
+        "Welcome! This tool helps you search for papers on arXiv, pick one, and ask questions about its content.\n\n"
+        "Please enter a topic to search for on arXiv papers."
+    )).send()
 
 @cl.on_message
 async def main(message: cl.Message):