tooluniverse.rxclass_tool 源代码

"""
RxClass Tool

Drug classification tools using the NLM RxClass API (part of RxNav):
  - get_drug_classes:   Look up ATC/EPC/MoA/VA drug classes for a drug name or RXCUI
  - get_class_members:  List drugs that belong to a given class ID
  - find_classes:       Search for drug classes by name keyword

API base: https://rxnav.nlm.nih.gov/REST/rxclass
No authentication required. Free public NLM API.
"""

import requests
from typing import Dict, Any
from .base_tool import BaseTool
from .tool_registry import register_tool

RXCLASS_BASE = "https://rxnav.nlm.nih.gov/REST/rxclass"
RXNORM_BASE = "https://rxnav.nlm.nih.gov/REST"

# Supported relaSource values for byDrugName endpoint
RELA_SOURCES = {
    "ATC": "WHO Anatomical Therapeutic Chemical classification",
    "FDASPL": "FDA Pharmacologic Class (EPC, MoA, PE)",
    "MESH": "MeSH pharmacological actions",
    "VA": "VA Drug Classification",
    "DAILYMED": "DailyMed drug classification",
}


[文档] @register_tool("RxClassTool") class RxClassTool(BaseTool): """ Drug classification via NLM RxClass API. Operations: - get_drug_classes: Get ATC, EPC, MoA, VA drug classes for a drug - get_class_members: List drugs in a specified drug class - find_classes: Search drug classes by keyword """
[文档] def __init__(self, tool_config: Dict[str, Any]): super().__init__(tool_config) self.timeout = tool_config.get("timeout", 30) self.operation = tool_config.get("fields", {}).get( "operation", "get_drug_classes" )
[文档] def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]: op = self.operation if op == "get_drug_classes": return self._get_drug_classes(arguments) if op == "get_class_members": return self._get_class_members(arguments) if op == "find_classes": return self._find_classes(arguments) return {"status": "error", "error": f"Unknown operation: {op}"}
# ------------------------------------------------------------------ # operation: get_drug_classes # ------------------------------------------------------------------
[文档] def _get_drug_classes(self, arguments: Dict[str, Any]) -> Dict[str, Any]: drug_name = arguments.get("drug_name") or arguments.get("name") rxcui = arguments.get("rxcui") rela_source = arguments.get("rela_source", "ATC") limit = arguments.get("limit", 20) if not drug_name and not rxcui: return {"status": "error", "error": "Provide 'drug_name' or 'rxcui'."} if rela_source not in RELA_SOURCES and rela_source != "ALL": rela_source = "ATC" try: if rxcui: url = f"{RXCLASS_BASE}/class/byRxcui.json" params: Dict[str, Any] = {"rxcui": str(rxcui).strip()} if rela_source != "ALL": params["relaSource"] = rela_source else: url = f"{RXCLASS_BASE}/class/byDrugName.json" params = {"drugName": drug_name.strip()} if rela_source != "ALL": params["relaSource"] = rela_source resp = requests.get(url, params=params, timeout=self.timeout) resp.raise_for_status() data = resp.json() except requests.exceptions.Timeout: return { "status": "error", "error": "RxClass API timeout", "retryable": True, } except requests.exceptions.HTTPError as e: sc = e.response.status_code return { "status": "error", "error": f"RxClass HTTP {sc}", "retryable": sc in (408, 429, 500, 502, 503, 504), } except ValueError: ct = resp.headers.get("content-type", "") return { "status": "error", "error": "RxClass returned non-JSON response", "content_type": ct, "response_snippet": resp.text[:200], "retryable": "text/html" in ct or resp.text.lstrip().startswith("<"), "suggestion": "RxClass may be under maintenance. Retry in a few minutes.", } except Exception as e: return {"status": "error", "error": str(e), "retryable": False} items = data.get("rxclassDrugInfoList", {}).get("rxclassDrugInfo", []) if not items: query_str = rxcui if rxcui else drug_name return { "status": "success", "data": [], "metadata": { "query": query_str, "rela_source": rela_source, "count": 0, "note": f"No drug classes found for '{query_str}' in source '{rela_source}'. Try rela_source='ALL' or a different source.", }, } classes = [] seen = set() for item in items: mc = item.get("rxclassMinConceptItem", {}) drug_mc = item.get("minConcept", {}) class_id = mc.get("classId", "") class_key = (class_id, drug_mc.get("rxcui", "")) if class_key in seen: continue seen.add(class_key) classes.append( { "classId": class_id, "className": mc.get("className", ""), "classType": mc.get("classType", ""), "rxcui": drug_mc.get("rxcui", ""), "drugName": drug_mc.get("name", ""), "tty": drug_mc.get("tty", ""), "rela": item.get("rela", ""), "relaSource": item.get("relaSource", rela_source), } ) classes = classes[:limit] return { "status": "success", "data": classes, "metadata": { "query": rxcui if rxcui else drug_name, "rela_source": rela_source, "count": len(classes), "available_sources": list(RELA_SOURCES.keys()), }, }
# ------------------------------------------------------------------ # operation: get_class_members # ------------------------------------------------------------------
[文档] def _get_class_members(self, arguments: Dict[str, Any]) -> Dict[str, Any]: class_id = arguments.get("class_id") or arguments.get("classId") rela_source = arguments.get("rela_source", "ATC") ttys = arguments.get("ttys", "IN") limit = arguments.get("limit", 50) if not class_id: return { "status": "error", "error": "Provide 'class_id' (e.g., 'M01AE', 'N02BA').", } try: url = f"{RXCLASS_BASE}/classMembers.json" params: Dict[str, Any] = { "classId": str(class_id).strip(), "relaSource": rela_source, "ttys": ttys, } resp = requests.get(url, params=params, timeout=self.timeout) resp.raise_for_status() data = resp.json() except requests.exceptions.Timeout: return { "status": "error", "error": "RxClass API timeout", "retryable": True, } except requests.exceptions.HTTPError as e: sc = e.response.status_code return { "status": "error", "error": f"RxClass HTTP {sc}", "retryable": sc in (408, 429, 500, 502, 503, 504), } except ValueError: ct = resp.headers.get("content-type", "") return { "status": "error", "error": "RxClass returned non-JSON response", "content_type": ct, "response_snippet": resp.text[:200], "retryable": "text/html" in ct or resp.text.lstrip().startswith("<"), "suggestion": "RxClass may be under maintenance. Retry in a few minutes.", } except Exception as e: return {"status": "error", "error": str(e), "retryable": False} members = data.get("drugMemberGroup", {}).get("drugMember", []) if not members: return { "status": "success", "data": [], "metadata": { "class_id": class_id, "rela_source": rela_source, "count": 0, "note": f"No drug members found for class '{class_id}'. Verify class ID and rela_source.", }, } drugs = [] for m in members[:limit]: mc = m.get("minConcept", {}) drugs.append( { "rxcui": mc.get("rxcui", ""), "name": mc.get("name", ""), "tty": mc.get("tty", ""), } ) return { "status": "success", "data": drugs, "metadata": { "class_id": class_id, "rela_source": rela_source, "ttys": ttys, "count": len(drugs), }, }
# ------------------------------------------------------------------ # operation: find_classes # ------------------------------------------------------------------
[文档] def _find_classes(self, arguments: Dict[str, Any]) -> Dict[str, Any]: query = arguments.get("query") or arguments.get("keyword") class_type = arguments.get("class_type", "") limit = arguments.get("limit", 20) if not query: return { "status": "error", "error": "Provide 'query' keyword to search drug classes.", } # classSearch.json is not available in current RxClass API version. # Use allClasses.json and filter client-side by class name keyword. params: Dict[str, Any] = {} if class_type: params["classTypes"] = class_type try: resp = requests.get( f"{RXCLASS_BASE}/allClasses.json", params=params, timeout=self.timeout ) resp.raise_for_status() data = resp.json() except requests.exceptions.Timeout: return { "status": "error", "error": "RxClass API timeout", "retryable": True, } except requests.exceptions.HTTPError as e: sc = e.response.status_code return { "status": "error", "error": f"RxClass HTTP {sc}", "retryable": sc in (408, 429, 500, 502, 503, 504), } except Exception as e: return {"status": "error", "error": str(e), "retryable": False} all_classes = data.get("rxclassMinConceptList", {}).get("rxclassMinConcept", []) kw = query.strip().lower() matches = [ { "classId": c.get("classId", ""), "className": c.get("className", ""), "classType": c.get("classType", ""), } for c in all_classes if kw in c.get("className", "").lower() ][:limit] return { "status": "success", "data": matches, "metadata": { "query": query, "class_type": class_type or "all", "count": len(matches), }, }