Update app.py

app.py

@@ -1,63 +1,57 @@
+import os
+from typing import List
+
 import chainlit as cl
 import arxiv
-from langchain.chat_models import ChatOpenAI
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
 from langchain.text_splitter import CharacterTextSplitter
-from langchain.embeddings import OpenAIEmbeddings
 from langchain.vectorstores import FAISS
-import os
 from dotenv import load_dotenv
-load_dotenv()
-
 
-# Initialize global variables
-selected_paper = None
-qa_chain = None
-papers = []
-state = "SEARCH"  # Possible states: SEARCH, SELECT, QA
+load_dotenv()
 
-@cl.on_chat_start
-def start():
-    global state
-    state = "SEARCH"
-    cl.Message(content="Welcome! Please enter a search query for arXiv papers.").send()
+class ArxivResearchAssistant:
+    def __init__(self):
+        self.selected_paper = None
+        self.qa_chain = None
+        self.papers: List[arxiv.Result] = []
+        self.state = "SEARCH"
 
-@cl.on_message
-def main(message: str):
-    global selected_paper, qa_chain, papers, state
-
-    if state == "SEARCH":
+    async def search_papers(self, query: str):
         search = arxiv.Search(
-            query=message,
+            query=query,
             max_results=5,
             sort_by=arxiv.SortCriterion.Relevance
         )
 
-        papers = list(search.results())
+        self.papers = list(search.results())
         
-        if not papers:
-            cl.Message(content="No papers found. Please try another search query.").send()
-            return
+        if not self.papers:
+            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(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"
+        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\nEnter the number of the paper you want to select:").send()
+        self.state = "SELECT"
+        return self.papers
 
-    elif state == "SELECT":
+    async def select_paper(self, selection: str):
         try:
-            selected_index = int(message) - 1
-            if 0 <= selected_index < len(papers):
-                selected_paper = papers[selected_index]
+            selected_index = int(selection) - 1
+            if 0 <= selected_index < len(self.papers):
+                self.selected_paper = self.papers[selected_index]
             else:
-                cl.Message(content="Invalid selection. Please try again.").send()
-                return
+                await cl.Message(content="Invalid selection. Please try again.").send()
+                return None
         except ValueError:
-            cl.Message(content="Invalid input. Please enter a number.").send()
-            return
+            await cl.Message(content="Invalid input. Please enter a number.").send()
+            return None
 
         # Download the entire paper content (if available)
-        paper_text = f"{selected_paper.title}\n\n{selected_paper.summary}\n\n{selected_paper.comment}"
+        paper_text = f"{self.selected_paper.title}\n\n{self.selected_paper.summary}\n\n{self.selected_paper.comment or ''}"
 
         # Split the text into chunks
         text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
@@ -69,8 +63,8 @@ def main(message: str):
             chunks,
             embeddings,
             metadatas=[{
-                "title": selected_paper.title,
-                "link": selected_paper.entry_id,
+                "title": self.selected_paper.title,
+                "link": self.selected_paper.entry_id,
                 "chunk": f"Chunk {i+1}/{len(chunks)}"
             } for i in range(len(chunks))]
         )
@@ -82,35 +76,59 @@ def main(message: str):
             output_key="answer"
         )
         
-        qa_chain = ConversationalRetrievalChain.from_llm(
+        self.qa_chain = ConversationalRetrievalChain.from_llm(
             ChatOpenAI(temperature=0, model="gpt-4o-mini"),
             vectorstore.as_retriever(),
             memory=memory,
             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()
-        state = "QA"
+        await cl.Message(content=f"Selected paper: {self.selected_paper.title}\nLink: {self.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()
+        self.state = "QA"
+        return self.selected_paper
 
-    elif state == "QA":
+    async def process_question(self, message: str):
         if message.lower() == "new search":
-            state = "SEARCH"
-            selected_paper = None
-            qa_chain = None
-            papers = []
-            cl.Message(content="Sure! Please enter a new search query for arXiv papers.").send()
-        else:
-            # Answer questions about the selected paper
-            response = qa_chain({"question": message})
-            answer = response["answer"]
-
-            # Handling the sources with chunk-specific metadata
-            sources = "\n".join([f"- {doc.metadata.get('title', 'Unknown title')} ({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}"
-
-            # Send the response with sources
-            cl.Message(content=answer).send()
+            self.reset()
+            await cl.Message(content="Sure! Please enter a new search query for arXiv papers.").send()
+            return None
+
+        response = self.qa_chain({"question": message})
+        answer = response["answer"]
+
+        # Handling the sources with chunk-specific metadata
+        sources = "\n".join([f"- {doc.metadata.get('title', 'Unknown title')} ({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}"
+
+        return answer
+
+    def reset(self):
+        self.selected_paper = None
+        self.qa_chain = None
+        self.papers = []
+        self.state = "SEARCH"
+
+# Global assistant instance
+assistant = ArxivResearchAssistant()
+
+@cl.on_chat_start
+async def start():
+    await cl.Message(content="Welcome! Please enter a search query for arXiv papers.").send()
+
+@cl.on_message
+async def main(message: cl.Message):
+    # Route the message based on the current state
+    if assistant.state == "SEARCH":
+        await assistant.search_papers(message.content)
+    
+    elif assistant.state == "SELECT":
+        await assistant.select_paper(message.content)
+    
+    elif assistant.state == "QA":
+        answer = await assistant.process_question(message.content)
+        if answer:
+            await cl.Message(content=answer).send()
 
 if __name__ == "__main__":
-    cl.run()
+    cl.run()