Update app.py

app.py

@@ -38,19 +38,105 @@ index = None
 # PDF FUNCTIONS
 # =========================
 
-def download_arxiv_pdf(arxiv_id):
+def doi_to_pdf_url(doi):
+    """
+    Given a DOI, try multiple strategies to find a downloadable PDF:
+    1. Unpaywall API  — finds legal open-access PDFs for any DOI (no key needed)
+    2. Direct arXiv   — if the DOI belongs to an arXiv paper
+    3. Europe PMC     — broad biomedical / life-science coverage
+    Returns (pdf_url, paper_title).
+    """
+    doi = doi.strip()
+    # Strip common DOI prefixes so bare DOI always works
+    for prefix in ("https://doi.org/", "http://doi.org/", "doi:", "DOI:"):
+        if doi.startswith(prefix):
+            doi = doi[len(prefix):]
+            break
+
+    title = None
+
+    # ── Strategy 1: Unpaywall (free, ~85% OA coverage) ──────────────
     try:
-        arxiv_id = arxiv_id.strip()
-        url = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
-        response = requests.get(url, timeout=30)
-        response.raise_for_status()
-        file_path = f"/tmp/{arxiv_id}.pdf"
+        email = "research@assistant.app"
+        r = requests.get(
+            f"https://api.unpaywall.org/v2/{doi}?email={email}", timeout=15
+        )
+        if r.status_code == 200:
+            data = r.json()
+            title = data.get("title")
+            best = data.get("best_oa_location")
+            if best:
+                pdf_url = best.get("url_for_pdf") or best.get("url")
+                if pdf_url:
+                    return pdf_url, title
+            for loc in data.get("oa_locations", []):
+                pdf_url = loc.get("url_for_pdf") or loc.get("url")
+                if pdf_url:
+                    return pdf_url, title
+    except Exception as e:
+        print(f"Unpaywall error: {e}")
+
+    # ── Strategy 2: arXiv DOI pattern ───────────────────────────────
+    try:
+        arxiv_match = re.search(r"arXiv[\.:](\d{4}\.\d{4,5})", doi, re.IGNORECASE)
+        if arxiv_match:
+            arxiv_id = arxiv_match.group(1)
+            return f"https://arxiv.org/pdf/{arxiv_id}.pdf", title
+    except Exception as e:
+        print(f"arXiv DOI parse error: {e}")
+
+    # ── Strategy 3: Europe PMC ───────────────────────────────────────
+    try:
+        r = requests.get(
+            f"https://www.ebi.ac.uk/europepmc/webservices/rest/search"
+            f"?query=DOI:{doi}&format=json&resultType=core",
+            timeout=15,
+        )
+        if r.status_code == 200:
+            results = r.json().get("resultList", {}).get("result", [])
+            if results:
+                item = results[0]
+                title = title or item.get("title")
+                pmcid = item.get("pmcid")
+                if pmcid:
+                    pdf_url = (
+                        f"https://europepmc.org/backend/ptpmcrender.fcgi"
+                        f"?accid={pmcid}&blobtype=pdf"
+                    )
+                    return pdf_url, title
+    except Exception as e:
+        print(f"Europe PMC error: {e}")
+
+    return None, title
+
+
+def download_pdf_from_doi(doi):
+    """Resolve DOI → PDF URL → download to /tmp. Returns (file_path, paper_title)."""
+    try:
+        pdf_url, title = doi_to_pdf_url(doi)
+        if not pdf_url:
+            return None, title
+
+        safe_name = re.sub(r"[^\w\-]", "_", doi)[:60]
+        file_path = f"/tmp/{safe_name}.pdf"
+
+        headers = {"User-Agent": "Mozilla/5.0 (ResearchAssistant/1.0)"}
+        r = requests.get(pdf_url, timeout=40, headers=headers, allow_redirects=True)
+        r.raise_for_status()
+
+        # Verify it's actually a PDF
+        if b"%PDF" not in r.content[:16]:
+            print(f"Response is not a PDF from {pdf_url}")
+            return None, title
+
         with open(file_path, "wb") as f:
-            f.write(response.content)
-        return file_path
+            f.write(r.content)
+
+        return file_path, title
+
     except Exception as e:
         print(f"PDF download error: {e}")
-        return None
+        return None, None
 
 def extract_text_from_pdf(pdf_path):
     doc = fitz.open(pdf_path)
@@ -98,19 +184,24 @@ def build_vector_store(sections_dict):
     index = faiss.IndexFlatL2(dim)
     index.add(embeddings)
 
-def load_paper(arxiv_id):
+def load_paper(doi):
     global sections
-    if not arxiv_id or not arxiv_id.strip():
-        return gr.update(choices=[]), "⚠️ Please enter a valid arXiv ID"
-    pdf_path = download_arxiv_pdf(arxiv_id)
+    if not doi or not doi.strip():
+        return gr.update(choices=[]), "⚠️ Please enter a valid DOI"
+    pdf_path, title = download_pdf_from_doi(doi)
     if pdf_path is None:
-        return gr.update(choices=[]), "❌ Could not download paper. Check the arXiv ID and try again."
+        return gr.update(choices=[]), (
+            "❌ Could not find or download a PDF for this DOI.\n\n"
+            "**Tips:** Make sure the paper is open-access. "
+            "Try formats like `10.1038/s41586-021-03819-2` or paste the full `https://doi.org/...` URL."
+        )
     text = extract_text_from_pdf(pdf_path)
     sections = extract_sections(text)
     if not sections:
-        return gr.update(choices=[]), "⚠️ Paper loaded but no sections could be detected."
+        return gr.update(choices=[]), "⚠️ PDF downloaded but no sections could be detected."
     build_vector_store(sections)
-    return gr.update(choices=list(sections.keys())), f"✅ Paper loaded — {len(sections)} sections indexed and ready"
+    label = f"✅ **{title}**\n\n{len(sections)} sections indexed and ready" if title else f"✅ Paper loaded — {len(sections)} sections indexed and ready"
+    return gr.update(choices=list(sections.keys())), label
 
 def summarize_section(section_title):
     try:
@@ -552,7 +643,7 @@ with gr.Blocks(title="ArXiv Research Assistant") as demo:
             <span class="badge">🧠 LLaMA 3.3 70B</span>
             <span class="badge">🎙️ Whisper ASR</span>
             <span class="badge">🔍 FAISS Semantic Search</span>
-            <span class="badge">📄 PyMuPDF</span>
+            <span class="badge">🔗 DOI Resolver</span>
         </div>
     </div>
     """)
@@ -562,14 +653,14 @@ with gr.Blocks(title="ArXiv Research Assistant") as demo:
         gr.HTML('<div class="sec-label">01 &nbsp; Load Paper</div>')
         with gr.Row(equal_height=True):
             arxiv_input = gr.Textbox(
-                label="arXiv Paper ID",
-                placeholder="e.g.  2305.10601  ·  2312.00752  ·  1706.03762",
+                label="Paper DOI",
+                placeholder="e.g.  10.1038/s41586-021-03819-2  or  https://doi.org/10.48550/arXiv.1706.03762",
                 scale=5,
             )
             load_btn = gr.Button("🚀  Load Paper", variant="primary", scale=1, min_width=150)
 
         status = gr.Markdown(
-            value="*Enter an arXiv ID above and click **Load Paper** to begin.*",
+            value="*Enter a paper DOI above and click **Load Paper** to begin.*",
             elem_id="status-box",
         )