update

app.py

@@ -13,6 +13,22 @@ import requests
 import seaborn as sns
 from tqdm import tqdm
 
+import arxiv
+
+
+def get_paper_info(paper_id):
+    # Create a search query with the arXiv ID
+    search = arxiv.Search(id_list=[paper_id])
+
+    # Fetch the paper using its arXiv ID
+    paper = next(search.results(), None)
+
+    if paper is not None:
+        # Return the paper's title and abstract
+        return paper.title, paper.summary
+    else:
+        return None, None
+
 
 def download_arxiv_source(paper_id):
     url = f"https://arxiv.org/e-print/{paper_id}"
@@ -74,10 +90,8 @@ class ContextualQA:
         self.context = text
 
     def ask_question(self, question):
-        leading_prompt = "Consider the text document below:"
-        trailing_prompt = (
-            "Now answer the following question, use Markdown to format your answer."
-        )
+        leading_prompt = "Consider the following paper:"
+        trailing_prompt = "Now, answer the following question using Markdown syntax:"
         prompt = f"{anthropic.HUMAN_PROMPT} {leading_prompt}\n\n{self.context}\n\n{trailing_prompt}\n\n{anthropic.HUMAN_PROMPT} {question} {anthropic.AI_PROMPT}"
         response = self.client.completion_stream(
             prompt=prompt,
@@ -107,16 +121,24 @@ class ContextualQA:
 
 
 def load_context(paper_id):
-    latex_source = download_arxiv_source(paper_id)
+    try:
+        latex_source = download_arxiv_source(paper_id)
+    except Exception as e:
+        return None, [(f"Error loading paper with id {paper_id}.", str(e))]
+
     client = anthropic.Client(api_key=os.environ["ANTHROPIC_API_KEY"])
     model = ContextualQA(client, model="claude-v1.3-100k")
     model.load_text(latex_source)
+
+    # Usage
+    title, abstract = get_paper_info(paper_id)
+
     return (
         model,
         [
             (
                 f"Load the paper with id {paper_id}.",
-                "Paper loaded, You can now ask questions.",
+                f"\n**Title**: {title}\n\n**Abstract**: {abstract}\n\nPaper loaded, You can now ask questions.",
             )
         ],
     )
@@ -131,7 +153,11 @@ def answer_fn(model, question, chat_history):
     client = anthropic.Client(api_key=os.environ["ANTHROPIC_API_KEY"])
     model.client = client
 
-    response = model.ask_question(question)
+    try:
+        response = model.ask_question(question)
+    except Exception as e:
+        chat_history.append(("Error Asking Question", str(e)))
+        return model, chat_history, ""
 
     chat_history.append((question, response[0]["completion"]))
     return model, chat_history, ""