File size: 8,499 Bytes
3ed6b66
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
01b3ee5
3ed6b66
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
01b3ee5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3ed6b66
 
01b3ee5
 
 
3ed6b66
01b3ee5
3ed6b66
 
01b3ee5
3ed6b66
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
"""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]