Source code for tooluniverse.pdbe_kb_tool
# pdbe_kb_tool.py
"""
PDBe-KB (PDBe Knowledge Base) Graph API tool for ToolUniverse.
PDBe-KB is an aggregated knowledge base that integrates structural data from
PDB with functional annotations from 30+ partner resources (UniProt, CATH,
SCOP, Pfam, etc.). The Graph API provides access to ligand binding sites,
protein-protein interaction interfaces, structural summaries, and more.
API: https://www.ebi.ac.uk/pdbe/graph-api/
No authentication required.
"""
import requests
from typing import Dict, Any
from .base_tool import BaseTool
from .tool_registry import register_tool
PDBE_KB_BASE_URL = "https://www.ebi.ac.uk/pdbe/graph-api"
def _residue_row(
residue: dict, pdb_entry_offset: int, pdb_entry_limit: int
) -> dict:
"""One residue of a PDBe-KB site, with the residue in it.
The row used to be `{start, end, pdb_entries, ...}`: a position and no
amino acid, in tools whose descriptions promise "binding site residues"
and "interface residues". `startCode`/`endCode` are the residue --
P04637's first cofactor-binding residue is LYS 382 -- and `indexType`
says which numbering frame the index is in, without which 382 could be a
PDB position rather than a UniProt one.
`allPDBEntries` is every PDB entry covering the residue; the paged
`pdb_entries` field has always been that list. `interactingPDBEntries` is
the subset where this residue actually contacts the partner, with the
chains, and is the narrower and more useful claim, so it is returned
whole beside it.
Both handlers in this module build this row; they are the ligand and the
protein-protein faces of one endpoint family and had no reason to differ.
"""
all_pdb_entries = residue.get("allPDBEntries", []) or []
pdb_entry_page = all_pdb_entries[
pdb_entry_offset : pdb_entry_offset + pdb_entry_limit
]
pdb_entries_truncated = pdb_entry_offset + len(pdb_entry_page) < len(
all_pdb_entries
)
return {
"start": residue.get("startIndex"),
"end": residue.get("endIndex"),
"start_code": residue.get("startCode"),
"end_code": residue.get("endCode"),
"index_type": residue.get("indexType"),
"interacting_pdb_entries": [
e for e in (residue.get("interactingPDBEntries") or []) if e
],
"pdb_entries": pdb_entry_page,
"pdb_entries_total": len(all_pdb_entries),
"pdb_entries_returned": len(pdb_entry_page),
"pdb_entries_truncated": pdb_entries_truncated,
"next_pdb_entry_offset": (
pdb_entry_offset + len(pdb_entry_page) if pdb_entries_truncated else None
),
}
[docs]
@register_tool("PDBe_KB_Tool")
class PDBe_KB_Tool(BaseTool):
"""
Tool for querying PDBe-KB (Knowledge Base) Graph API.
PDBe-KB aggregates structural biology knowledge including:
- Ligand binding sites mapped to UniProt positions
- Protein-protein interaction interfaces
- Structural coverage statistics
- Superposition clusters and best chain coverage
Data is indexed by UniProt accession and provides residue-level
annotations using UniProt numbering.
No authentication required.
"""
[docs]
def __init__(self, tool_config: Dict[str, Any]):
super().__init__(tool_config)
self.timeout = tool_config.get("timeout", 30)
fields = tool_config.get("fields", {})
self.endpoint = fields.get("endpoint", "summary_stats")
[docs]
def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Execute the PDBe-KB Graph API call."""
try:
return self._query(arguments)
except requests.exceptions.Timeout:
return {
"status": "error",
"error": f"PDBe-KB API timed out after {self.timeout}s",
}
except requests.exceptions.ConnectionError:
return {"status": "error", "error": "Failed to connect to PDBe-KB API"}
except requests.exceptions.HTTPError as e:
if e.response.status_code == 404:
acc = arguments.get("uniprot_accession", "unknown")
return {"status": "error", "error": f"No PDBe-KB data found for {acc}"}
return {
"status": "error",
"error": f"PDBe-KB API HTTP error: {e.response.status_code}",
}
except Exception as e:
return {
"status": "error",
"error": f"Unexpected error querying PDBe-KB: {str(e)}",
}
[docs]
def _query(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Route to appropriate PDBe-KB endpoint."""
if self.endpoint == "summary_stats":
return self._get_summary_stats(arguments)
elif self.endpoint == "ligand_sites":
return self._get_ligand_sites(arguments)
elif self.endpoint == "interface_residues":
return self._get_interface_residues(arguments)
elif self.endpoint == "superposition":
return self._get_superposition(arguments)
else:
return {"status": "error", "error": f"Unknown endpoint: {self.endpoint}"}
[docs]
def _get_summary_stats(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Get aggregated structural summary statistics for a protein."""
acc = arguments.get("uniprot_accession", "")
if not acc:
return {
"status": "error",
"error": "uniprot_accession parameter is required",
}
url = f"{PDBE_KB_BASE_URL}/uniprot/summary_stats/{acc}"
response = requests.get(url, timeout=self.timeout)
response.raise_for_status()
data = response.json()
if acc not in data:
return {"status": "error", "error": f"No summary data for {acc}"}
stats = data[acc]
return {
"status": "success",
"data": {
"uniprot_accession": acc,
"pdbs": stats.get("pdbs"),
"ligands": stats.get("ligands"),
"interaction_partners": stats.get("interaction_partners"),
"annotations": stats.get("annotations"),
"similar_proteins": stats.get("similar_proteins"),
},
"metadata": {
"source": "PDBe-KB (PDBe Knowledge Base) Graph API",
},
}
[docs]
def _get_ligand_sites(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Get ligand binding site residues for a protein."""
acc = arguments.get("uniprot_accession", "")
if not acc:
return {
"status": "error",
"error": "uniprot_accession parameter is required",
}
url = f"{PDBE_KB_BASE_URL}/uniprot/ligand_sites/{acc}"
response = requests.get(url, timeout=self.timeout)
response.raise_for_status()
data = response.json()
if acc not in data:
return {"status": "error", "error": f"No ligand binding data for {acc}"}
protein_data = data[acc]
ligands = []
all_ligand_entries = protein_data.get("data", [])
ligand_offset = max(int(arguments.get("ligand_offset") or 0), 0)
max_ligands = arguments.get("max_ligands")
# Check for None explicitly, then require a positive page size so the
# advertised next offset always advances.
max_ligands = 100 if max_ligands is None else max(1, int(max_ligands))
residue_offset = max(int(arguments.get("residue_offset") or 0), 0)
residue_limit = min(max(int(arguments.get("residue_limit") or 20), 1), 500)
pdb_entry_offset = max(int(arguments.get("pdb_entry_offset") or 0), 0)
pdb_entry_limit = min(max(int(arguments.get("pdb_entry_limit") or 10), 1), 500)
ligand_page = all_ligand_entries[ligand_offset : ligand_offset + max_ligands]
for entry in ligand_page:
all_residues = entry.get("residues", [])
residue_page = all_residues[residue_offset : residue_offset + residue_limit]
binding_residues = []
for res in residue_page:
binding_residues.append(
_residue_row(res, pdb_entry_offset, pdb_entry_limit)
)
additional = entry.get("additionalData", {})
residues_truncated = residue_offset + len(residue_page) < len(all_residues)
ligands.append(
{
"name": entry.get("name"),
"accession": entry.get("accession"),
"binding_residues": binding_residues,
"binding_residues_total": len(all_residues),
"binding_residues_returned": len(residue_page),
"binding_residues_truncated": residues_truncated,
"next_residue_offset": (
residue_offset + len(residue_page)
if residues_truncated
else None
),
"is_cofactor": bool(additional.get("coFactorId")),
"cofactor_id": additional.get("coFactorId") or None,
"is_solvent": additional.get("isSolvent", False),
"chembl_id": additional.get("chemblId") or None,
"drugbank_id": additional.get("drugBankId") or None,
# What the ligand is, and how strongly PDBe-KB rates the
# site: the SMILES of the scaffold, its atom count, and
# the significance the ligand list is ordered by.
"scaffold_smiles": additional.get("scaffoldId") or None,
"atom_count": additional.get("numAtoms"),
"significance": additional.get("significance"),
"ligand_type": additional.get("type") or None,
"reaction_id": additional.get("reactionId") or None,
"target_uniprot_accessions": [
t for t in (additional.get("targetUniProts") or []) if t
],
"ligand_pdb_entries": [
e for e in (additional.get("pdbEntries") or []) if e
],
}
)
result_data = {
"uniprot_accession": acc,
"protein_length": protein_data.get("length"),
"ligands": ligands,
"total_ligands": len(all_ligand_entries),
"ligand_offset": ligand_offset,
"ligands_returned": len(ligands),
}
ligands_truncated = ligand_offset + len(ligand_page) < len(all_ligand_entries)
result_data["truncated"] = ligands_truncated
result_data["next_ligand_offset"] = (
ligand_offset + len(ligand_page) if ligands_truncated else None
)
if ligands_truncated:
result_data["truncated"] = True
result_data["truncation_note"] = (
f"Returning {len(ligands)} ligands from offset {ligand_offset} of "
f"{len(all_ligand_entries)} total. Continue with "
f"ligand_offset={ligand_offset + len(ligand_page)}. "
"The API does not guarantee ordering by clinical relevance, so "
"page through every ligand when looking for a specific compound."
)
return {
"status": "success",
"data": result_data,
"metadata": {
"source": "PDBe-KB (PDBe Knowledge Base) Graph API",
},
}
[docs]
def _get_interface_residues(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Get protein-protein interaction interface residues."""
acc = arguments.get("uniprot_accession", "")
if not acc:
return {
"status": "error",
"error": "uniprot_accession parameter is required",
}
url = f"{PDBE_KB_BASE_URL}/uniprot/interface_residues/{acc}"
response = requests.get(url, timeout=self.timeout)
response.raise_for_status()
data = response.json()
if acc not in data:
return {"status": "error", "error": f"No interface data for {acc}"}
protein_data = data[acc]
partners = []
all_partner_entries = protein_data.get("data", [])
partner_offset = max(int(arguments.get("partner_offset") or 0), 0)
max_partners = arguments.get("max_partners")
# Check for None explicitly, then require a positive page size so the
# advertised next offset always advances.
max_partners = 60 if max_partners is None else max(1, int(max_partners))
residue_offset = max(int(arguments.get("residue_offset") or 0), 0)
residue_limit = min(max(int(arguments.get("residue_limit") or 30), 1), 500)
pdb_entry_offset = max(int(arguments.get("pdb_entry_offset") or 0), 0)
pdb_entry_limit = min(max(int(arguments.get("pdb_entry_limit") or 10), 1), 500)
partner_page = all_partner_entries[
partner_offset : partner_offset + max_partners
]
for entry in partner_page:
all_residues = entry.get("residues", [])
residue_page = all_residues[residue_offset : residue_offset + residue_limit]
interface_residues = []
for res in residue_page:
interface_residues.append(
_residue_row(res, pdb_entry_offset, pdb_entry_limit)
)
residues_truncated = residue_offset + len(residue_page) < len(all_residues)
partners.append(
{
"partner_name": entry.get("name"),
"partner_accession": entry.get("accession"),
"interface_residues": interface_residues,
"interface_residues_total": len(all_residues),
"interface_residues_returned": len(residue_page),
"interface_residues_truncated": residues_truncated,
"next_residue_offset": (
residue_offset + len(residue_page)
if residues_truncated
else None
),
}
)
result_data = {
"uniprot_accession": acc,
"protein_length": protein_data.get("length"),
"interaction_partners": partners,
"total_partners": len(all_partner_entries),
"partner_offset": partner_offset,
"partners_returned": len(partners),
}
partners_truncated = partner_offset + len(partner_page) < len(
all_partner_entries
)
result_data["truncated"] = partners_truncated
result_data["next_partner_offset"] = (
partner_offset + len(partner_page) if partners_truncated else None
)
if partners_truncated:
result_data["truncated"] = True
result_data["truncation_note"] = (
f"Returning {len(partners)} interaction partners from offset "
f"{partner_offset} of {len(all_partner_entries)} total. Continue "
f"with partner_offset={partner_offset + len(partner_page)}."
)
return {
"status": "success",
"data": result_data,
"metadata": {
"source": "PDBe-KB (PDBe Knowledge Base) Graph API",
},
}
[docs]
def _get_superposition(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Get structural superposition clusters for a protein.
Returns clusters of structurally superposed PDB chains grouped by
protein segments. Each cluster contains a representative structure
and aligned member structures.
"""
acc = arguments.get("uniprot_accession", "")
if not acc:
return {
"status": "error",
"error": "uniprot_accession parameter is required (e.g., 'P04637' for TP53, 'P00533' for EGFR).",
}
url = f"{PDBE_KB_BASE_URL}/uniprot/superposition/{acc}"
response = requests.get(url, timeout=self.timeout)
response.raise_for_status()
data = response.json()
if acc not in data:
return {"status": "error", "error": f"No superposition data for {acc}"}
segments_data = data[acc]
segments = []
segment_offset = max(int(arguments.get("segment_offset") or 0), 0)
segment_limit = min(max(int(arguments.get("segment_limit") or 10), 1), 100)
cluster_offset = max(int(arguments.get("cluster_offset") or 0), 0)
cluster_limit = min(max(int(arguments.get("cluster_limit") or 10), 1), 100)
member_offset = max(int(arguments.get("member_offset") or 0), 0)
member_limit = min(max(int(arguments.get("member_limit") or 20), 1), 500)
segment_page = segments_data[segment_offset : segment_offset + segment_limit]
for seg in segment_page:
all_clusters = seg.get("clusters", [])
cluster_page = all_clusters[cluster_offset : cluster_offset + cluster_limit]
cluster_results = []
for cluster in cluster_page:
members = []
representative_member = next(
(m for m in cluster if m.get("is_representative", False)), None
)
representative = (
{
"pdb_id": representative_member.get("pdb_id"),
"auth_asym_id": representative_member.get("auth_asym_id"),
"struct_asym_id": representative_member.get("struct_asym_id"),
"entity_id": representative_member.get("entity_id"),
"is_representative": True,
}
if representative_member is not None
else None
)
member_page = cluster[member_offset : member_offset + member_limit]
for member in member_page:
entry = {
"pdb_id": member.get("pdb_id"),
"auth_asym_id": member.get("auth_asym_id"),
"struct_asym_id": member.get("struct_asym_id"),
"entity_id": member.get("entity_id"),
"is_representative": member.get("is_representative", False),
}
members.append(entry)
members_truncated = member_offset + len(member_page) < len(cluster)
cluster_results.append(
{
"representative": representative,
"total_members": len(cluster),
"member_offset": member_offset,
"members_returned": len(member_page),
"members_truncated": members_truncated,
"next_member_offset": (
member_offset + len(member_page)
if members_truncated
else None
),
"members": members,
}
)
clusters_truncated = cluster_offset + len(cluster_page) < len(all_clusters)
segments.append(
{
"segment_start": seg.get("segment_start"),
"segment_end": seg.get("segment_end"),
"num_clusters": len(all_clusters),
"cluster_offset": cluster_offset,
"clusters_returned": len(cluster_page),
"clusters_truncated": clusters_truncated,
"next_cluster_offset": (
cluster_offset + len(cluster_page)
if clusters_truncated
else None
),
"clusters": cluster_results,
}
)
segments_truncated = segment_offset + len(segment_page) < len(segments_data)
return {
"status": "success",
"data": {
"uniprot_accession": acc,
"segments": segments,
"total_segments": len(segments_data),
"segment_offset": segment_offset,
"segments_returned": len(segment_page),
"segments_truncated": segments_truncated,
"next_segment_offset": (
segment_offset + len(segment_page) if segments_truncated else None
),
},
"metadata": {
"source": "PDBe-KB (PDBe Knowledge Base) Graph API",
},
}