avoigt1121
Claude Opus 5.5
fix(drugs): one ChEMBL window for the mechanism lookups, retried round-robin
01b3ee5 | """What every endpoint shares: the client, the call budget, dates, stages and HGNC resolution. | |
| Endpoints are plain functions — strings in, payload dicts out — and no exception escapes one | |
| (the family tool_server contract). A failure inside one source call is a ``source_errors`` entry, | |
| an input that a source does not know is an ``unresolved`` entry, and only an endpoint where | |
| nothing could be looked up at all is an ``error`` payload. | |
| """ | |
| from __future__ import annotations | |
| import threading | |
| import time | |
| from collections.abc import Iterable | |
| from dataclasses import dataclass, field | |
| from typing import Any | |
| from src import config | |
| from src.provenance import stamp | |
| from src.tu_client import SourceError, SourceResult, ToolClient | |
| # One endpoint call may spend at most this long across all its source calls; the orchestrator's | |
| # own call limit is above it (300 s for annotate_genes, its default for the others). | |
| ENDPOINT_BUDGET_SECONDS = 240 | |
| _client: ToolClient | None = None | |
| _client_lock = threading.Lock() | |
| def get_client() -> ToolClient: | |
| """The process-wide client (one cache, one set of rate limits for every endpoint).""" | |
| global _client | |
| with _client_lock: | |
| if _client is None: | |
| _client = ToolClient() | |
| return _client | |
| def set_client(client: ToolClient | None) -> None: | |
| """Tests replay recorded responses through a client of their own.""" | |
| global _client | |
| with _client_lock: | |
| _client = client | |
| class Budget: | |
| """The endpoint's wall-clock budget. Past it, remaining inputs are reported, not looked up.""" | |
| def __init__(self, seconds: float = ENDPOINT_BUDGET_SECONDS, clock=time.monotonic): | |
| self._clock = clock | |
| self.seconds = seconds | |
| self._deadline = clock() + seconds | |
| def remaining(self) -> float: | |
| return max(0.0, self._deadline - self._clock()) | |
| def spent(self) -> bool: | |
| return self.remaining <= 0 | |
| class Collector: | |
| """Accumulates what an endpoint found, failed on and could not resolve.""" | |
| rows: list[dict[str, Any]] = field(default_factory=list) | |
| unresolved: list[dict[str, str]] = field(default_factory=list) | |
| source_errors: list[dict[str, str]] = field(default_factory=list) | |
| sources_used: set[str] = field(default_factory=set) | |
| retrieved: list[float] = field(default_factory=list) | |
| calls: int = 0 | |
| def call( | |
| self, | |
| client: ToolClient, | |
| budget: Budget, | |
| tool: str, | |
| args: dict[str, Any], | |
| *, | |
| unknown: tuple[str, str] | None = None, | |
| ) -> SourceResult | None: | |
| """One source call. None (with the reason recorded) if it failed or the budget ran out. | |
| ``unknown=(marker, input)``: a source error whose detail contains ``marker`` means the | |
| source does not know ``input`` (Reactome's 404, Open Targets' "no disease for ID"). That | |
| is an unresolved input, not a source failure. | |
| """ | |
| if budget.spent: | |
| return None | |
| result, err = self.attempt(client, budget, tool, args) | |
| if err is None: | |
| return result | |
| if unknown is not None and err.reason == "source_unavailable" and unknown[0] in err.detail: | |
| self.unresolved_input(unknown[1], f"unknown to {err.source}: {err.detail}") | |
| else: | |
| self.source_errors.append(err.as_payload()) | |
| return None | |
| def attempt( | |
| self, | |
| client: ToolClient, | |
| budget: Budget, | |
| tool: str, | |
| args: dict[str, Any], | |
| ) -> tuple[SourceResult | None, SourceError | None]: | |
| """One source call whose failure is handed back, not recorded: the caller may retry it. | |
| The call's deadline never runs past ``budget``: a call started with 12 s left gets 12 s. | |
| """ | |
| self.calls += 1 | |
| try: | |
| result = client.call( | |
| tool, args, max_wait=min(10.0, budget.remaining), deadline=budget.remaining | |
| ) | |
| except SourceError as err: | |
| return None, err | |
| self.sources_used.add(result.source) | |
| self.retrieved.append(result.retrieved_at) | |
| return result, None | |
| def unresolved_input(self, value: str, reason: str) -> None: | |
| self.unresolved.append({"input": value, "reason": reason}) | |
| def out_of_budget(self, values: Iterable[str]) -> None: | |
| for value in values: | |
| self.unresolved_input( | |
| value, f"not looked up: the {ENDPOINT_BUDGET_SECONDS} s call budget ran out" | |
| ) | |
| def stamped(row: dict[str, Any], result: SourceResult, *, ref: str, url: str) -> dict[str, Any]: | |
| return stamp( | |
| row, | |
| source=result.source, | |
| source_tool=result.tool, | |
| tooluniverse_version=result.tooluniverse_version, | |
| query=result.args, | |
| source_ref=ref, | |
| source_url=url, | |
| retrieved_at=result.retrieved_at, | |
| cache=result.cache, | |
| ) | |
| def query_dates(retrieved: Iterable[float]) -> str: | |
| """'2026-10-20', or '2026-10-14 to 2026-10-20' when cached answers span days.""" | |
| from src.provenance import iso_utc | |
| days = sorted({iso_utc(t)[:10] for t in retrieved}) | |
| if not days: | |
| return "today" | |
| return days[0] if len(days) == 1 else f"{days[0]} to {days[-1]}" | |
| def according_to(source_ids: Iterable[str], retrieved: Iterable[float]) -> str: | |
| names = [config.source(s)["name"] for s in sorted(set(source_ids))] | |
| joined = names[0] if len(names) == 1 else ", ".join(names[:-1]) + " and " + names[-1] | |
| return f"According to {joined} (queried {query_dates(retrieved)})" | |
| def humanise_list(items: list[str], limit: int = 5) -> str: | |
| shown = items[:limit] | |
| more = len(items) - len(shown) | |
| text = ", ".join(shown) | |
| return f"{text} and {more} more" if more > 0 else text | |
| # Open Targets' clinical-stage vocabulary, highest first. Anything unknown sorts last and is | |
| # shown as Open Targets wrote it. | |
| _STAGE_ORDER = ( | |
| "APPROVAL", | |
| "APPROVED", | |
| "PHASE_4", | |
| "PHASE_3", | |
| "PHASE_2_3", | |
| "PHASE_2", | |
| "PHASE_1_2", | |
| "PHASE_1", | |
| "EARLY_PHASE_1", | |
| "PHASE_0", | |
| "PRECLINICAL", | |
| ) | |
| def stage_rank(stage: str | None) -> int: | |
| try: | |
| return _STAGE_ORDER.index((stage or "").upper()) | |
| except ValueError: | |
| return len(_STAGE_ORDER) | |
| def stage_label(stage: str | None) -> str: | |
| if not stage: | |
| return "stage not stated" | |
| text = stage.upper() | |
| if text in ("APPROVAL", "APPROVED"): | |
| return "approved" | |
| return ( | |
| text.replace("_", " ") | |
| .lower() | |
| .replace("phase 1 2", "phase 1/2") | |
| .replace("phase 2 3", "phase 2/3") | |
| ) | |
| class Gene: | |
| """An input resolved through HGNC.""" | |
| input: str | |
| symbol: str | |
| hgnc_id: str | |
| name: str | |
| locus_type: str | |
| ensembl_gene_id: str | None | |
| uniprot_ids: tuple[str, ...] | |
| provenance_row: dict[str, Any] | |
| HGNC = "HGNC_fetch_gene_by_symbol" | |
| def resolve_symbols( | |
| symbols: list[str], client: ToolClient, budget: Budget, out: Collector | |
| ) -> list[Gene]: | |
| """HGNC symbol → approved record. Unknown symbols go to ``unresolved``.""" | |
| genes: list[Gene] = [] | |
| for i, symbol in enumerate(symbols): | |
| if budget.spent: | |
| out.out_of_budget(symbols[i:]) | |
| break | |
| result = out.call(client, budget, HGNC, {"symbol": symbol}) | |
| if result is None: | |
| continue | |
| data = result.data or {} | |
| if not data.get("symbol"): | |
| out.unresolved_input(symbol, "not a current, previous or alias symbol in HGNC") | |
| continue | |
| row = stamped( | |
| {"input": symbol, "symbol": data["symbol"], "hgnc_id": data["hgnc_id"]}, | |
| result, | |
| ref=data["hgnc_id"], | |
| url=f"https://www.genenames.org/data/gene-symbol-report/#!/hgnc_id/{data['hgnc_id']}", | |
| ) | |
| genes.append( | |
| Gene( | |
| input=symbol, | |
| symbol=data["symbol"], | |
| hgnc_id=data["hgnc_id"], | |
| name=data.get("name") or "", | |
| locus_type=data.get("locus_type") or "", | |
| ensembl_gene_id=data.get("ensembl_gene_id"), | |
| uniprot_ids=tuple(data.get("uniprot_ids") or ()), | |
| provenance_row=row, | |
| ) | |
| ) | |
| return genes | |
| def resolution_rows(genes: list[Gene]) -> list[dict[str, Any]]: | |
| """The identifier mapping each answer rests on, with HGNC's provenance.""" | |
| return [g.provenance_row for g in genes] | |