Source code for tooluniverse.uniprot_idmapping_tool
# uniprot_idmapping_tool.py
"""
UniProt ID Mapping Service tool for ToolUniverse.
Provides cross-database protein/gene identifier conversion using the
canonical UniProt ID Mapping REST API. Supports 100+ databases including
UniProtKB, Gene Names, Ensembl, RefSeq, PDB, ChEMBL, and more.
The service is asynchronous: a mapping job is submitted via POST,
its status is polled, and results are retrieved when complete.
API: https://rest.uniprot.org/idmapping/
No authentication required.
"""
import time
import requests
from typing import Dict, Any
from .base_tool import BaseTool
from .tool_registry import register_tool
UNIPROT_BASE_URL = "https://rest.uniprot.org"
def _entry_summary(entry: Dict[str, Any]) -> Dict[str, Any]:
"""What the mapped-to UniProtKB entry is, when the response carries it.
`to` is a bare accession for most target databases and the whole entry
for UniProtKB, where these are the fields that let a caller pick between
several accessions for one gene name.
"""
# The keys are always present so the row shape does not change with the
# endpoint: /idmapping/results/{job} answers a bare accession string while
# /idmapping/uniprotkb/results/{job} answers the entry, and UniProt
# redirects between them.
if not isinstance(entry, dict):
entry = {}
protein = entry.get("proteinDescription") or {}
recommended = (protein.get("recommendedName") or {}).get("fullName") or {}
submitted = protein.get("submissionNames") or []
if not recommended and submitted:
recommended = (submitted[0].get("fullName") or {}) if isinstance(
submitted[0], dict
) else {}
organism = entry.get("organism") or {}
entry_type = entry.get("entryType")
return {
"entry_type": entry_type,
# "unreviewed" contains "reviewed": the substring test calls every
# TrEMBL entry reviewed, so test for the negative string.
"reviewed": ("unreviewed" not in entry_type.lower())
if isinstance(entry_type, str) and entry_type
else None,
"protein_name": recommended.get("value")
if isinstance(recommended, dict)
else None,
"gene_names": [
(gene.get("geneName") or {}).get("value")
for gene in (entry.get("genes") or [])
if isinstance(gene, dict) and (gene.get("geneName") or {}).get("value")
],
"organism": organism.get("scientificName"),
"organism_taxon_id": organism.get("taxonId"),
"sequence_length": (entry.get("sequence") or {}).get("length"),
"protein_existence": entry.get("proteinExistence"),
"annotation_score": entry.get("annotationScore"),
}
[docs]
@register_tool("UniProtIDMappingTool")
class UniProtIDMappingTool(BaseTool):
"""
Tool for converting identifiers between databases using
the UniProt ID Mapping service.
Handles the async submit -> poll -> results workflow automatically.
No authentication required.
"""
[docs]
def __init__(self, tool_config: Dict[str, Any]):
super().__init__(tool_config)
self.timeout = tool_config.get("timeout", 30)
self.endpoint_type = tool_config.get("fields", {}).get(
"endpoint_type", "convert"
)
[docs]
def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Execute the UniProt ID Mapping API call."""
try:
return self._dispatch(arguments)
except requests.exceptions.Timeout:
return {
"status": "error",
"error": f"UniProt ID Mapping API timed out after {self.timeout}s",
}
except requests.exceptions.ConnectionError:
return {
"status": "error",
"error": "Failed to connect to UniProt ID Mapping API",
}
except requests.exceptions.HTTPError as e:
status = e.response.status_code if e.response else "unknown"
body = ""
try:
body = e.response.json().get("messages", [""])[0]
except Exception:
pass
return {
"status": "error",
"error": f"UniProt ID Mapping HTTP error {status}: {body}",
}
except Exception as e:
return {"status": "error", "error": f"Unexpected error: {str(e)}"}
[docs]
def _dispatch(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Route to appropriate endpoint."""
if self.endpoint_type == "convert":
return self._convert(arguments)
elif self.endpoint_type == "to_pdb":
return self._to_pdb(arguments)
elif self.endpoint_type == "gene_to_uniprot":
return self._gene_to_uniprot(arguments)
elif self.endpoint_type == "list_databases":
return self._list_databases(arguments)
return {
"status": "error",
"error": f"Unknown endpoint_type: {self.endpoint_type}",
}
[docs]
def _submit_and_poll(
self, from_db: str, to_db: str, ids: str, tax_id: int = None
) -> Dict:
"""Submit a mapping job and poll for results."""
# Submit job
data = {
"from": from_db,
"to": to_db,
"ids": ids,
}
if tax_id:
data["taxId"] = tax_id
submit_resp = requests.post(
f"{UNIPROT_BASE_URL}/idmapping/run",
data=data,
timeout=self.timeout,
)
submit_resp.raise_for_status()
job_id = submit_resp.json().get("jobId")
if not job_id:
return {
"status": "error",
"error": "Failed to get job ID from UniProt ID Mapping",
}
# Poll for completion (max 60 seconds). The status endpoint sometimes
# 303-redirects straight into a results page instead of ever
# reporting jobStatus=="FINISHED" -- that redirected page uses its
# own default page size (25), not the size=500 this tool needs, so
# treat its presence as "job done" and always re-fetch the full,
# paginated result set via _fetch_all_results() below rather than
# returning that first small page directly. Confirmed live:
# P00533 -> PDB has 354 real mappings, but the redirected status
# response only ever carries the first 25 of them.
max_polls = 20
for _ in range(max_polls):
status_resp = requests.get(
f"{UNIPROT_BASE_URL}/idmapping/status/{job_id}",
timeout=self.timeout,
)
status_data = status_resp.json()
# For jobs that complete almost instantly, the status endpoint can
# return results embedded directly instead of a "FINISHED"
# jobStatus (confirmed live) -- either way, treat it as done and
# fall through to the paginated results fetch below (size=500).
# Returning this embedded copy directly previously silently
# truncated to UniProt's unpaginated default page size (25),
# even when the job actually had hundreds of matches.
if (
status_data.get("jobStatus") == "FINISHED"
or "results" in status_data
or "failedIds" in status_data
):
break
if status_data.get("jobStatus") == "ERROR":
msg = status_data.get("errorMessage", "Unknown error")
return {
"status": "error",
"error": f"UniProt mapping job failed: {msg}",
}
time.sleep(1.5)
else:
return {
"status": "error",
"error": "UniProt ID mapping job did not complete within timeout",
}
# A UniProtKB target has a richer result view; anything else does not.
stream = "uniprotkb" if str(to_db).startswith("UniProtKB") else ""
results, failed_ids = self._fetch_all_results(job_id, stream=stream)
return {
"status": "success",
"results": results,
"job_id": job_id,
"failed_ids": failed_ids,
}
[docs]
def _fetch_all_results(self, job_id: str, max_pages: int = 20, stream: str = ""):
"""Fetch every page of a completed mapping job's results, following
the Link header's rel="next" URL so large result sets (e.g. a
heavily-studied protein with hundreds of PDB structures) aren't
silently truncated to a single page.
`stream` selects the result view. `/idmapping/results/{job}` answers
each mapping as a bare target id; `/idmapping/uniprotkb/results/{job}`
answers the whole UniProtKB entry, which is what tells two accessions
for one gene name apart. Measured on Gene_Name TP53 -> UniProtKB: the
first returns "P04637", the second the entry it belongs to.
"""
results = []
failed_ids = []
path = f"{stream}/results" if stream else "results"
url = f"{UNIPROT_BASE_URL}/idmapping/{path}/{job_id}"
params = {"size": 500}
for _ in range(max_pages):
resp = requests.get(url, params=params, timeout=self.timeout)
resp.raise_for_status()
page = resp.json()
results.extend(page.get("results", []))
failed_ids.extend(page.get("failedIds", []))
next_url = self._parse_next_link(resp.headers.get("Link"))
if not next_url:
break
url, params = next_url, None
return results, failed_ids
[docs]
@staticmethod
def _parse_next_link(link_header):
"""Extract the rel="next" URL from an RFC 5988 Link header."""
if not link_header:
return None
for part in link_header.split(","):
segments = part.split(";")
if len(segments) >= 2 and 'rel="next"' in segments[1]:
return segments[0].strip().strip("<>")
return None
[docs]
@staticmethod
def _paginate_parsed_results(parsed, arguments):
"""Apply recoverable local pagination after all upstream pages are fetched."""
limit = max(1, min(int(arguments.get("limit", 500)), 5000))
offset = max(0, int(arguments.get("offset", 0)))
page = parsed[offset : offset + limit]
has_more = offset + len(page) < len(parsed)
return {
"results": page,
"returned": len(page),
"limit": limit,
"offset": offset,
"has_more": has_more,
"next_offset": offset + len(page) if has_more else None,
"truncated": offset > 0 or has_more,
}
[docs]
def _convert(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Generic ID conversion between any supported databases."""
ids = arguments.get("ids", "")
from_db = arguments.get("from_db", "")
to_db = arguments.get("to_db", "UniProtKB")
tax_id = arguments.get("tax_id")
if not ids:
return {
"status": "error",
"error": "ids parameter is required (e.g., 'TP53,BRCA1')",
}
if not from_db:
return {
"status": "error",
"error": "from_db parameter is required (e.g., 'Gene_Name')",
}
result = self._submit_and_poll(from_db, to_db, ids, tax_id)
if "error" in result:
return result
raw_results = result.get("results", [])
# Parse results - handle both simple and complex formats
parsed = []
for r in raw_results:
to_val = r.get("to", "")
# Some results have nested objects for 'to'
if isinstance(to_val, dict):
to_val = to_val.get("primaryAccession", to_val.get("id", str(to_val)))
parsed.append({"from": r.get("from", ""), "to": str(to_val)})
page_data = self._paginate_parsed_results(parsed, arguments)
return {
"status": "success",
"data": {
"from_db": from_db,
"to_db": to_db,
"result_count": len(parsed),
**page_data,
"failed_ids": result.get("failed_ids", []),
},
"metadata": {
"source": "UniProt ID Mapping Service",
"job_id": result.get("job_id", ""),
"endpoint": "idmapping",
},
}
[docs]
def _to_pdb(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Convert UniProt accessions to PDB IDs."""
uniprot_ids = arguments.get("uniprot_ids", "")
if not uniprot_ids:
return {
"status": "error",
"error": "uniprot_ids is required (e.g., 'P04637')",
}
result = self._submit_and_poll("UniProtKB_AC-ID", "PDB", uniprot_ids)
if "error" in result:
return result
raw_results = result.get("results", [])
parsed = [
{"from": r.get("from", ""), "to": str(r.get("to", ""))} for r in raw_results
]
page_data = self._paginate_parsed_results(parsed, arguments)
return {
"status": "success",
"data": {
"query_ids": uniprot_ids,
"result_count": len(parsed),
**page_data,
},
"metadata": {
"source": "UniProt ID Mapping Service",
"endpoint": "idmapping (UniProtKB_AC-ID -> PDB)",
},
}
[docs]
def _gene_to_uniprot(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Convert gene names to UniProt accessions."""
gene_names = arguments.get("gene_names", "")
tax_id = arguments.get("tax_id", 9606)
reviewed_only = arguments.get("reviewed_only", False)
if not gene_names:
return {
"status": "error",
"error": "gene_names is required (e.g., 'TP53,BRCA1')",
}
to_db = "UniProtKB-Swiss-Prot" if reviewed_only else "UniProtKB"
result = self._submit_and_poll("Gene_Name", to_db, gene_names, tax_id)
if "error" in result:
return result
raw_results = result.get("results", [])
parsed = []
for r in raw_results:
to_val = r.get("to", "")
entry = to_val if isinstance(to_val, dict) else {}
if isinstance(to_val, dict):
to_val = to_val.get("primaryAccession", str(to_val))
# A gene name resolves to several accessions and the row was the
# accession string alone, so nothing distinguished the reviewed
# human entry from an unreviewed isoform of it. The entry that
# answered the mapping is already in the response.
parsed.append(
{
"from": r.get("from", ""),
"to": str(to_val),
**_entry_summary(entry),
}
)
page_data = self._paginate_parsed_results(parsed, arguments)
return {
"status": "success",
"data": {
"gene_names": gene_names,
"species_taxid": tax_id,
"result_count": len(parsed),
**page_data,
},
"metadata": {
"source": "UniProt ID Mapping Service",
"endpoint": f"idmapping (Gene_Name -> {to_db})",
},
}
[docs]
def _list_databases(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""List available databases for ID mapping."""
url = f"{UNIPROT_BASE_URL}/configure/idmapping/fields"
response = requests.get(url, timeout=self.timeout)
response.raise_for_status()
raw = response.json()
groups_raw = raw.get("groups", [])
groups = []
for g in groups_raw:
dbs = []
for item in g.get("items", []):
dbs.append(
{
"name": item.get("name", ""),
"display_name": item.get("displayName", ""),
"from_supported": item.get("from", False),
# Mapping has two directions and only one was
# reported: 95 databases each say whether they can be
# mapped *from* and *to*, and a caller choosing a
# target database had nothing to choose on.
"to_supported": item.get("to", False),
# How to resolve an id in that database, which is what
# a mapped id is for. 90 of the 95 carry one.
"uri_template": item.get("uriLink"),
"rule_id": item.get("ruleId"),
}
)
groups.append(
{
"group_name": g.get("groupName", ""),
"databases": dbs,
}
)
return {
"status": "success",
"data": {
"group_count": len(groups),
"groups": groups,
},
"metadata": {
"source": "UniProt ID Mapping Service",
"endpoint": "configure/idmapping/fields",
},
}