"""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 @property def remaining(self) -> float: return max(0.0, self._deadline - self._clock()) @property def spent(self) -> bool: return self.remaining <= 0 @dataclass 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") ) @dataclass(frozen=True) 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]