pubmed search

app.py

@@ -22,8 +22,10 @@ from langchain_anthropic import ChatAnthropic
 from langgraph.checkpoint.sqlite import SqliteSaver
 from langchain_community.agent_toolkits import SQLDatabaseToolkit
 from langchain.agents import create_agent
+from langchain_core.tools import tool
 
 from scripts.db import create_database
+from scripts.literature import literature_search
 
 # --- Paths -----------------------------------------------------------
 _DATA_DIR    = "/data" if os.path.isdir("/data") else "."
@@ -51,9 +53,32 @@ def _save_to_registry(name, db_file):
     with open(_REGISTRY_PATH, "w") as f:
         json.dump(registry, f)
 
+# --- Tools -----------------------------------------------------------
+@tool
+def pubmed_search(query: str, max_results: int = 5) -> str:
+    """Search PubMed for published biomedical literature and return titles,
+    authors, journal, year, PMID, URL and abstract for the top matches.
+
+    Use this when the user asks about published research, what is known about a
+    protein or gene, or for evidence supporting a finding in the dataset.
+
+    `query` accepts PubMed search syntax. Prefer gene or protein names over
+    UniProt accessions, and combine concepts with AND, e.g.
+    'Sox2 AND neural stem cell' or 'Tp53[Title] AND apoptosis'.
+    `max_results` is capped at 20.
+    """
+    return literature_search(query, max_results)
+
+
 # --- Agent -----------------------------------------------------------
 _SYSTEM_PROMPT = """
-You are an agent designed to interact with a SQL database.
+You are a proteomics research assistant. You can query a SQL database holding
+the user's proteomics dataset, and you can search PubMed for published
+literature. Decide which tools a question needs: dataset questions need SQL,
+questions about published research need PubMed, and some questions need both.
+
+## Querying the dataset
+
 Given an input question, create a syntactically correct {dialect} query to run,
 then look at the results of the query and return the answer. Unless the user
 specifies a specific number of examples they wish to obtain, always limit your
@@ -69,10 +94,26 @@ executing a query, rewrite the query and try again.
 DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the
 database.
 
-To start you should ALWAYS look at the tables in the database to see what you
-can query. Do NOT skip this step.
+Before running any query you should ALWAYS look at the tables in the database
+to see what you can query. Do NOT skip this step. Then you should query the
+schema of the most relevant tables.
+
+## Searching the literature
+
+Use `pubmed_search` when the user asks what is published or known about a
+protein, gene or biological process, or asks for evidence supporting something
+in the dataset.
+
+Search by gene or protein name rather than UniProt accession, since accessions
+rarely appear in abstracts. If the dataset gives you an accession, look up its
+gene name or description in the database first, then search on that. If a
+search returns nothing, retry with broader or alternative terms before
+concluding there is no literature.
 
-Then you should query the schema of the most relevant tables.
+Ground your answer in the abstracts returned and cite each claim with its PMID
+and PubMed URL. Never invent a citation, PMID or finding that is not in the
+search results. If a question combines the dataset and the literature, report
+the dataset result first, then what the literature says about it.
 """
 
 def _init_agent(db):
@@ -80,14 +121,14 @@ def _init_agent(db):
     toolkit = SQLDatabaseToolkit(db=db, llm=model)
     agent = create_agent(
         model,
-        toolkit.get_tools(),
+        toolkit.get_tools() + [pubmed_search],
         system_prompt=_SYSTEM_PROMPT.format(dialect=db.dialect, top_k=5),
         checkpointer=memory,
     )
 
 # --- UI helpers ------------------------------------------------------
 _msg_locked = dict(interactive=False, placeholder="Select or load a dataset to start chatting")
-_msg_open   = dict(interactive=True,  placeholder="Ask about your proteomics data")
+_msg_open   = dict(interactive=True,  placeholder="Ask about your proteomics data or the PubMed literature")
 
 def _read_tabular(path, nrows=None):
     if path.endswith((".xlsx", ".xls")):

requirements.txt

@@ -2,7 +2,7 @@ langchain
 langchain-community
 gradio
 nltk==3.9.1
-metapub
+requests
 langchain-experimental
 langgraph
 langchain-anthropic 

scripts/literature.py

@@ -1,7 +1,146 @@
-import metapub as mpub
+import os
+import time
+import xml.etree.ElementTree as ET
 
-def literature_search(query):
-    fetch = mpub.PubMedFetcher()
-    ids = fetch.pmids_for_query(query)
+import requests
 
-    return(len(ids))
+_EUTILS = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils'
+_TIMEOUT = 30
+
+# NCBI allows 3 req/s without an API key, 10 req/s with one. Each search costs
+# two requests (esearch + efetch), so consecutive agent calls trip the limit.
+_MIN_INTERVAL = 0.12 if os.environ.get('NCBI_API_KEY') else 0.4
+_RETRIES = 3
+_last_request = 0.0
+
+
+def _params(**kwargs):
+    """Common E-utilities params. NCBI_API_KEY raises the rate limit 3/s -> 10/s."""
+    params = {'db': 'pubmed', 'tool': 'BOTeome', **kwargs}
+    api_key = os.environ.get('NCBI_API_KEY')
+    if api_key:
+        params['api_key'] = api_key
+    email = os.environ.get('NCBI_EMAIL')
+    if email:
+        params['email'] = email
+    return params
+
+
+def _get(endpoint, **params):
+    """Throttled GET against E-utilities, retrying on rate limits."""
+    global _last_request
+
+    for attempt in range(_RETRIES):
+        wait = _MIN_INTERVAL - (time.monotonic() - _last_request)
+        if wait > 0:
+            time.sleep(wait)
+
+        r = requests.get(f'{_EUTILS}/{endpoint}', params=_params(**params), timeout=_TIMEOUT)
+        _last_request = time.monotonic()
+
+        if r.status_code == 429 and attempt < _RETRIES - 1:
+            time.sleep(2 ** attempt)  # 1s, 2s
+            continue
+        r.raise_for_status()
+        return r
+
+    r.raise_for_status()
+    return r
+
+
+def _text(node, path):
+    found = node.find(path)
+    return found.text.strip() if found is not None and found.text else ''
+
+
+def _parse_article(article):
+    medline = article.find('MedlineCitation')
+    if medline is None:
+        return None
+
+    pmid = _text(medline, 'PMID')
+    art = medline.find('Article')
+    if art is None:
+        return None
+
+    # Abstracts are split into labelled sections (BACKGROUND, METHODS, ...).
+    sections = []
+    for chunk in art.findall('Abstract/AbstractText'):
+        body = ''.join(chunk.itertext()).strip()
+        if not body:
+            continue
+        label = chunk.get('Label')
+        sections.append(f'{label}: {body}' if label else body)
+    abstract = ' '.join(sections) or 'no abstract available'
+
+    authors = []
+    for author in art.findall('AuthorList/Author'):
+        last = _text(author, 'LastName')
+        initials = _text(author, 'Initials')
+        if last:
+            authors.append(f'{last} {initials}'.strip())
+    if len(authors) > 3:
+        authors = authors[:3] + ['et al.']
+
+    year = (
+        _text(art, 'Journal/JournalIssue/PubDate/Year')
+        or _text(art, 'Journal/JournalIssue/PubDate/MedlineDate')[:4]
+        or 'n.d.'
+    )
+
+    return {
+        'pmid': pmid,
+        'title': ''.join(art.find('ArticleTitle').itertext()).strip()
+                 if art.find('ArticleTitle') is not None else 'untitled',
+        'journal': _text(art, 'Journal/ISOAbbreviation') or _text(art, 'Journal/Title'),
+        'year': year,
+        'authors': ', '.join(authors) or 'unknown authors',
+        'abstract': abstract,
+        'url': f'https://pubmed.ncbi.nlm.nih.gov/{pmid}/',
+    }
+
+
+def search_pubmed(query, max_results=5):
+    """Search PubMed and return a list of article dicts, most relevant first.
+
+    `query` accepts PubMed syntax, e.g. 'Tp53[Title] AND apoptosis'.
+    Raises requests.HTTPError if NCBI is unreachable or rejects the request.
+    """
+    max_results = max(1, min(int(max_results), 20))
+
+    r = _get('esearch.fcgi', term=query, retmax=max_results, retmode='json', sort='relevance')
+    pmids = r.json().get('esearchresult', {}).get('idlist', [])
+    if not pmids:
+        return []
+
+    r = _get('efetch.fcgi', id=','.join(pmids), retmode='xml')
+
+    root = ET.fromstring(r.content)
+    articles = [_parse_article(a) for a in root.findall('PubmedArticle')]
+    return [a for a in articles if a]
+
+
+def format_results(articles, abstract_chars=1200):
+    """Render articles as plain text for the agent to read."""
+    if not articles:
+        return 'No PubMed articles found for that query.'
+
+    blocks = []
+    for i, a in enumerate(articles, 1):
+        abstract = a['abstract']
+        if len(abstract) > abstract_chars:
+            abstract = abstract[:abstract_chars].rstrip() + '...'
+        blocks.append(
+            f"[{i}] {a['title']}\n"
+            f"{a['authors']}. {a['journal']} ({a['year']}). PMID {a['pmid']} — {a['url']}\n"
+            f"{abstract}"
+        )
+    return '\n\n'.join(blocks)
+
+
+def literature_search(query, max_results=5):
+    """Search PubMed and return formatted text. Errors are returned, not raised."""
+    try:
+        return format_results(search_pubmed(query, max_results))
+    except Exception as e:
+        return f'PubMed search failed: {e}'