diff --git a/pyproject.toml b/pyproject.toml index f5f6b62..8747787 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,9 +75,14 @@ cli = [ "stack[bib]", "stack[llm]", "stack[mail]", + "stack[nlm]", "typer>=0.24.1", "uvicorn>=0.41.0", ] +nlm = [ + "stack[conf]", + "httpx>=0.28.1", +] mail = [ "httpx>=0.28.1", "pydo>=0.29.0", @@ -172,6 +177,7 @@ all = [ "stack[rec]", "stack[rex]", "stack[lake]", + "stack[nlm]", ] [dependency-groups] @@ -239,6 +245,6 @@ link-mode = "copy" python-preference = "managed" [tool.uv.build-backend] -module-name = ["aco", "api", "bcda", "bib", "bls", "ccw", "cli", "cms", "conf", "llm", "mail", "opps", "perf", "pfs", "prisma", "qpp", "rec", "rex", "sem"] +module-name = ["aco", "api", "bcda", "bib", "bls", "ccw", "cli", "cms", "conf", "llm", "mail", "nlm", "opps", "perf", "pfs", "prisma", "qpp", "rec", "rex", "sem"] namespace = true source-exclude = ["compose.yml","infra/**","data/**","notebooks/**","tuva/**","assets/**","docs/**","dev/**","bundle/**","cloud/**"] diff --git a/src/cli/__init__.py b/src/cli/__init__.py index b1cd7af..40e36a5 100644 --- a/src/cli/__init__.py +++ b/src/cli/__init__.py @@ -24,6 +24,7 @@ from cli.lake import app as lake_app from cli.llm import app as llm_app from cli.load import app as load_app from cli.mail import app as mail_app +from cli.nlm import app as nlm_app from cli.perf import app as perf_app from cli.pfs import app as pfs_app from cli.prisma import app as prisma_app @@ -71,6 +72,7 @@ app.add_typer( app.add_typer(zot_app, name="zot", help="Zotero database maintenance.") app.add_typer(prisma_app, name="prisma", help="LLM-driven PRISMA screening workflow.") app.add_typer(mail_app, name="mail", help="Maddy mail server (DO + Resend smarthost).") +app.add_typer(nlm_app, name="nlm", help="NLM terminology releases and VSAC value sets.") def main() -> None: diff --git a/src/cli/nlm.py b/src/cli/nlm.py new file mode 100644 index 0000000..de19289 --- /dev/null +++ b/src/cli/nlm.py @@ -0,0 +1,130 @@ +"""stack nlm — NLM terminology releases and VSAC value sets (UMLS licence). + +uv run stack nlm releases # current file of every stream +uv run stack nlm fetch # download all of them +uv run stack nlm fetch --stream snomed-ct-us-edition --stream rxnorm-full-monthly-release +uv run stack nlm vsac # every definition + expansion +uv run stack nlm vsac --definitions-only + +Needs ``UMLS_API_KEY`` in the environment (``set -a; . ./.env``). +Files land under ``conf.path("nlm.dir")`` unless ``--dest`` says otherwise; +the content is licensed — never publish it. +""" + +from __future__ import annotations + +import gzip +import json +from pathlib import Path + +import typer + +app = typer.Typer(no_args_is_help=True) + + +def _key() -> str: + from conf import secret + + key = secret("nlm.api_key", "UMLS_API_KEY") + if not key: + typer.echo("UMLS_API_KEY is not set (UTS profile → API key; add it to .env).") + raise typer.Exit(2) + return key + + +def _root(dest: Path | None) -> Path: + if dest is not None: + return dest + from conf import path + + return path("nlm.dir") + + +def _client(): + import httpx + + return httpx.Client(timeout=httpx.Timeout(60.0, read=600.0)) + + +@app.command() +def releases( + stream: list[str] = typer.Option(None, "--stream", help="Only these streams."), +) -> None: + """List the current file of every release stream.""" + from nlm.releases import current + + with _client() as client: + for rel in current(client, stream or None): + typer.echo(f"{rel.stream:52} {rel.version:12} {rel.file_name}") + + +@app.command() +def fetch( + stream: list[str] = typer.Option(None, "--stream", help="Only these streams."), + dest: Path = typer.Option(None, "--dest", help="Root directory (default nlm.dir)."), +) -> None: + """Download the current file of every (or each named) release stream.""" + from nlm.releases import current + from nlm.releases import fetch as fetch_release + + key, root = _key(), _root(dest) + failed = 0 + with _client() as client: + for rel in current(client, stream or None): + try: + path, status = fetch_release(client, rel, root, key) + size = path.stat().st_size / 1e6 + typer.echo(f"{status:10} {size:10.1f} MB {path.relative_to(root)}") + except RuntimeError as e: + failed += 1 + typer.echo(f"failed {e}") + if failed: + raise typer.Exit(1) + + +@app.command() +def vsac( + dest: Path = typer.Option(None, "--dest", help="Root directory (default nlm.dir)."), + definitions_only: bool = typer.Option(False, "--definitions-only"), + workers: int = typer.Option(4, "--workers", min=1, max=16), + limit: int = typer.Option(0, "--limit", help="Expand at most N (0 = all)."), +) -> None: + """Save every VSAC value-set definition, then expand each one.""" + from nlm import vsac as v + + key = _key() + root = _root(dest) / "vsac" + with v.make_client(key) as client: + n = v.save_definitions(client, root) + typer.echo(f"definitions: {n}") + if definitions_only: + return + ids = [ + json.loads(line)["id"] + for line in (root / "index.jsonl").read_text(encoding="utf-8").splitlines() + if line + ] + if limit: + ids = ids[:limit] + + def progress(done: int, total: int) -> None: + if done % 500 == 0 or done == total: + typer.echo(f" expanded {done}/{total}") + + stats = v.hoard( + lambda: v.make_client(key), + root, + ids, + api_key=key, + workers=workers, + progress=progress, + ) + typer.echo( + f"expansions: {stats['expanded']} new, {stats['present']} present, " + f"{stats['failed']} failed" + ) + codes = 0 + for p in (root / "expansions").glob("*.json.gz"): + with gzip.open(p, "rt", encoding="utf-8") as f: + codes += len(json.load(f).get("expansion", {}).get("contains", [])) + typer.echo(f"codes on disk: {codes:,}") diff --git a/src/nlm/__init__.py b/src/nlm/__init__.py new file mode 100644 index 0000000..31b5586 --- /dev/null +++ b/src/nlm/__init__.py @@ -0,0 +1,29 @@ +"""NLM terminology: UMLS, SNOMED CT and RxNorm releases, and VSAC value sets. + +Everything here needs a UMLS licence API key (UTS profile → API key), +read with ``conf.secret("nlm.api_key", "UMLS_API_KEY")``. The content is +licensed, not public: it lives under ``conf.path("nlm.dir")`` (gitignored +``data/nlm``) and must never be published — not to the docs site, the +package registry, Zotero or a commit. + + releases.release_types(client) → every UTS release stream + releases.current(client, slugs=…) → the current file of each stream + releases.fetch(client, rel, root, key) → the zip on disk + sha256 sidecar + vsac.iter_definitions(client) → every ValueSet resource (paged) + vsac.hoard(make_client, root, ids, …) → one $expand per value set, on disk + +Secrets: the key travels as a query parameter (UTS) or basic auth (VSAC); +:func:`scrub` strips it from any message before it can reach a log. +""" + +from __future__ import annotations + +import re + +_KEY_PARAM = re.compile(r"(apiKey=)[^&\s]+", re.I) + + +def scrub(text: str, key: str = "") -> str: + """Remove an API key from *text* — URLs in httpx errors carry it.""" + text = _KEY_PARAM.sub(r"\1***", text) + return text.replace(key, "***") if key else text diff --git a/src/nlm/releases.py b/src/nlm/releases.py new file mode 100644 index 0000000..25dce5d --- /dev/null +++ b/src/nlm/releases.py @@ -0,0 +1,140 @@ +"""UTS release downloads: UMLS, SNOMED CT and RxNorm. + +The release API (``uts-ws.nlm.nih.gov/releases``) is public and lists +every stream (``rxnorm-full-monthly-release``, ``snomed-ct-us-edition``, +``umls-full-release``, …) with its files; ``current=true`` narrows each +to the latest. Downloads go through ``/download?url=&apiKey=``, +which redirects to the file. + +On disk: ``///`` plus ``.json`` (what +it is, where it came from, bytes, sha256). A file whose sidecar exists +is done; an interrupted download resumes from its ``.part`` with a +Range request. +""" + +from __future__ import annotations + +import hashlib +import json +import zipfile +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from urllib.parse import parse_qs, urlparse + +import httpx + +from nlm import scrub + +RELEASES_URL = "https://uts-ws.nlm.nih.gov/releases" +DOWNLOAD_URL = "https://uts-ws.nlm.nih.gov/download" + + +@dataclass(frozen=True) +class Release: + product: str + stream: str # the releaseType slug, e.g. "snomed-ct-us-edition" + name: str # the human release type, e.g. "SNOMED CT US Edition" + version: str + date: str + file_name: str + download_url: str + + +def release_types(client: httpx.Client) -> list[dict[str, str]]: + """Every release stream: ``{product, name, stream, endpoint}``.""" + r = client.get(RELEASES_URL) + r.raise_for_status() + out = [] + for t in r.json()["releaseTypes"]: + slug = parse_qs(urlparse(t["endpoint"]).query)["releaseType"][0] + out.append( + { + "product": t["product"], + "name": t["releaseType"], + "stream": slug, + "endpoint": t["endpoint"], + } + ) + return out + + +def current(client: httpx.Client, streams: list[str] | None = None) -> list[Release]: + """The current file of each stream (all streams when *streams* is empty).""" + out = [] + for t in release_types(client): + if streams and t["stream"] not in streams: + continue + r = client.get( + RELEASES_URL, params={"releaseType": t["stream"], "current": "true"} + ) + r.raise_for_status() + for f in r.json(): + out.append( + Release( + product=t["product"], + stream=t["stream"], + name=t["name"], + version=str(f.get("releaseVersion") or f.get("releaseDate") or ""), + date=str(f.get("releaseDate") or ""), + file_name=f["fileName"], + download_url=f["downloadUrl"], + ) + ) + return out + + +def target(root: Path, rel: Release) -> Path: + return root / rel.stream / rel.version / rel.file_name + + +def _sha256(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as f: + for block in iter(lambda: f.read(1 << 20), b""): + h.update(block) + return h.hexdigest() + + +def fetch( + client: httpx.Client, rel: Release, root: Path, api_key: str +) -> tuple[Path, str]: + """Download *rel* under *root*; return ``(path, status)``. + + *status* is ``"present"`` (sidecar already written) or + ``"downloaded"``. Raises ``RuntimeError`` — with the key scrubbed — + on an HTTP failure or a file that is not the zip it claims to be. + """ + path = target(root, rel) + sidecar = path.with_name(path.name + ".json") + if path.exists() and sidecar.exists(): + return path, "present" + path.parent.mkdir(parents=True, exist_ok=True) + part = path.with_name(path.name + ".part") + have = part.stat().st_size if part.exists() else 0 + headers = {"Range": f"bytes={have}-"} if have else {} + params = {"url": rel.download_url, "apiKey": api_key} + try: + with client.stream( + "GET", DOWNLOAD_URL, params=params, headers=headers, follow_redirects=True + ) as r: + r.raise_for_status() + mode = "ab" if have and r.status_code == 206 else "wb" + with part.open(mode) as out: + for chunk in r.iter_bytes(1 << 20): + out.write(chunk) + except httpx.HTTPError as e: + raise RuntimeError(f"{rel.file_name}: {scrub(str(e), api_key)}") from None + if rel.file_name.lower().endswith(".zip") and not zipfile.is_zipfile(part): + head = part.read_bytes()[:200].decode("utf-8", "replace") + part.unlink() + raise RuntimeError(f"{rel.file_name}: not a zip ({scrub(head, api_key)!r})") + part.rename(path) + meta = { + **asdict(rel), + "bytes": path.stat().st_size, + "sha256": _sha256(path), + "fetched_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), + } + sidecar.write_text(json.dumps(meta, indent=1) + "\n", encoding="utf-8") + return path, "downloaded" diff --git a/src/nlm/vsac.py b/src/nlm/vsac.py new file mode 100644 index 0000000..63aa1b1 --- /dev/null +++ b/src/nlm/vsac.py @@ -0,0 +1,190 @@ +"""VSAC value sets over the FHIR API (``cts.nlm.nih.gov/fhir``). + +Auth is HTTP basic with user ``apikey`` and the UMLS key as password. +Every value-set *version* is its own resource (id ``-``), +18.6k of them in October 2026. + +On disk under ````: + +- ``definitions/.json.gz`` — the ValueSet resource (metadata and + ``compose``), from the paged search; +- ``index.jsonl`` — one line per definition (id, url, version, title, + status, publisher, date), rewritten on each definitions pass; +- ``expansions/.json.gz`` — the ``$expand`` result with every code + (large expansions arrive in pages and are joined); +- ``errors.jsonl`` — value sets that would not expand, with the reason. + +Existing files are skipped, so a hoard resumes where it stopped. +""" + +from __future__ import annotations + +import gzip +import json +import threading +import time +from collections.abc import Callable, Iterable, Iterator +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import httpx + +from nlm import scrub + +FHIR_URL = "https://cts.nlm.nih.gov/fhir" +_RETRY = {429, 500, 502, 503, 504} + + +def make_client(api_key: str, *, timeout: float = 120.0) -> httpx.Client: + return httpx.Client( + base_url=FHIR_URL, + auth=("apikey", api_key), + headers={"Accept": "application/fhir+json"}, + timeout=timeout, + ) + + +def _get( + client: httpx.Client, + url: str, + params: dict | None = None, + *, + tries: int = 5, + sleep: Callable[[float], None] = time.sleep, +) -> dict: + """GET with backoff on throttling / transient server errors.""" + for attempt in range(tries): + r = client.get(url, params=params) + if r.status_code in _RETRY and attempt < tries - 1: + wait = float(r.headers.get("Retry-After") or 2**attempt) + sleep(wait) + continue + r.raise_for_status() + return r.json() + raise AssertionError("unreachable") # pragma: no cover + + +def _write_gz(path: Path, payload: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + with gzip.open(tmp, "wt", encoding="utf-8") as f: + json.dump(payload, f) + tmp.rename(path) + + +def iter_definitions(client: httpx.Client, *, page_size: int = 500) -> Iterator[dict]: + """Every ValueSet resource, following the search's ``next`` links.""" + url: str | None = "/ValueSet" + params: dict | None = {"_count": page_size} + while url: + bundle = _get(client, url, params) + for entry in bundle.get("entry", []): + yield entry["resource"] + nxt = next( + ( + link["url"] + for link in bundle.get("link", []) + if link["relation"] == "next" + ), + None, + ) + # VSAC hands back http:// next links; the API itself is https-only. + url = nxt.replace("http://", "https://", 1) if nxt else None + params = None + + +def save_definitions(client: httpx.Client, root: Path, *, page_size: int = 500) -> int: + """Write every definition plus ``index.jsonl``; return how many.""" + index = [] + for vs in iter_definitions(client, page_size=page_size): + path = root / "definitions" / f"{vs['id']}.json.gz" + if not path.exists(): + _write_gz(path, vs) + index.append( + { + k: vs.get(k, "") + for k in ( + "id", + "url", + "version", + "title", + "status", + "publisher", + "date", + ) + } + ) + root.mkdir(parents=True, exist_ok=True) + (root / "index.jsonl").write_text( + "".join(json.dumps(row) + "\n" for row in index), encoding="utf-8" + ) + return len(index) + + +def expand(client: httpx.Client, vs_id: str, *, page: int = 1000) -> dict: + """``$expand`` one value set, joining paged expansions into one.""" + params = {"count": page, "offset": 0} + result = _get(client, f"/ValueSet/{vs_id}/$expand", params) + expansion = result.setdefault("expansion", {}) + contains = expansion.setdefault("contains", []) + total = expansion.get("total", len(contains)) + while len(contains) < total: + params["offset"] = len(contains) + more = _get(client, f"/ValueSet/{vs_id}/$expand", params) + batch = more.get("expansion", {}).get("contains", []) + if not batch: + break + contains.extend(batch) + expansion.pop("offset", None) + return result + + +def hoard( + client_factory: Callable[[], httpx.Client], + root: Path, + ids: Iterable[str], + *, + api_key: str = "", + workers: int = 4, + progress: Callable[[int, int], None] | None = None, +) -> dict[str, int]: + """Expand every id not already on disk, *workers* at a time. + + Returns ``{"expanded", "present", "failed"}``. A failure is logged to + ``errors.jsonl`` (key scrubbed) and never stops the run. + """ + todo = [] + stats = {"expanded": 0, "present": 0, "failed": 0} + for vs_id in ids: + if (root / "expansions" / f"{vs_id}.json.gz").exists(): + stats["present"] += 1 + else: + todo.append(vs_id) + + local = threading.local() + lock = threading.Lock() + errors = root / "errors.jsonl" + + def one(vs_id: str) -> None: + if not hasattr(local, "client"): + local.client = client_factory() + try: + _write_gz( + root / "expansions" / f"{vs_id}.json.gz", expand(local.client, vs_id) + ) + key = "expanded" + except (httpx.HTTPError, ValueError) as e: + key = "failed" + line = json.dumps({"id": vs_id, "error": scrub(str(e), api_key)}) + with lock: + errors.parent.mkdir(parents=True, exist_ok=True) + with errors.open("a", encoding="utf-8") as f: + f.write(line + "\n") + with lock: + stats[key] += 1 + if progress: + progress(stats["expanded"] + stats["failed"], len(todo)) + + with ThreadPoolExecutor(max_workers=workers) as pool: + list(pool.map(one, todo)) + return stats diff --git a/stack.toml b/stack.toml index f283404..4df2ea8 100644 --- a/stack.toml +++ b/stack.toml @@ -239,6 +239,12 @@ output_schema = "_monitoring" [default] target = "local" # duckdb | databricks | trino | lake +[nlm] +# UMLS / SNOMED CT / RxNorm releases and VSAC value sets. UMLS-licensed: +# never publish (docs site, package registry, Zotero, commits). +dir = "data/nlm" +api_key = "" # override via UMLS_API_KEY env var + [api] host = "0.0.0.0" port = 8000 diff --git a/tests/cli/test_nlm_cli.py b/tests/cli/test_nlm_cli.py new file mode 100644 index 0000000..cc2bcab --- /dev/null +++ b/tests/cli/test_nlm_cli.py @@ -0,0 +1,140 @@ +"""stack nlm — releases / fetch / vsac.""" + +from __future__ import annotations + +import gzip +import json + +import pytest +from typer.testing import CliRunner + +import cli.nlm as nlm_cli +import nlm.releases as releases_mod +import nlm.vsac as vsac_mod +from cli import app +from nlm.releases import Release + +runner = CliRunner() + +REL = Release( + product="UMLS", + stream="umls-full-release", + name="UMLS Full Release", + version="2026AA", + date="2026-05-03", + file_name="umls-2026AA-full.zip", + download_url="https://download.nlm.nih.gov/umls/kss/2026AA/umls-2026AA-full.zip", +) + + +@pytest.fixture +def key(monkeypatch): + monkeypatch.setenv("UMLS_API_KEY", "k-123") + + +def test_missing_key_exits_with_a_hint(monkeypatch, tmp_path): + monkeypatch.delenv("UMLS_API_KEY", raising=False) + monkeypatch.setattr("conf.secret", lambda *a: "") + r = runner.invoke(app, ["nlm", "fetch", "--dest", str(tmp_path)]) + assert r.exit_code == 2 and "UMLS_API_KEY" in r.output + + +def test_releases_lists_current_files(monkeypatch): + monkeypatch.setattr(releases_mod, "current", lambda client, streams: [REL]) + r = runner.invoke(app, ["nlm", "releases"]) + assert r.exit_code == 0 and "umls-2026AA-full.zip" in r.output + + +def test_fetch_reports_each_file_and_fails_when_one_did(monkeypatch, tmp_path, key): + other = Release( + **{ + **REL.__dict__, + "stream": "rxnorm-full-monthly-release", + "file_name": "rx.zip", + } + ) + monkeypatch.setattr(releases_mod, "current", lambda client, streams: [REL, other]) + + def fake_fetch(client, rel, root, api_key): + assert api_key == "k-123" + if rel.file_name == "rx.zip": + raise RuntimeError("rx.zip: 401 Unauthorized") + p = root / rel.stream / rel.version / rel.file_name + p.parent.mkdir(parents=True) + p.write_bytes(b"x" * 2_000_000) + return p, "downloaded" + + monkeypatch.setattr(releases_mod, "fetch", fake_fetch) + r = runner.invoke(app, ["nlm", "fetch", "--dest", str(tmp_path)]) + assert r.exit_code == 1 + assert "downloaded" in r.output and "2.0 MB" in r.output + assert "failed rx.zip: 401" in r.output + + +def test_fetch_defaults_to_the_configured_dir(monkeypatch, tmp_path, key): + monkeypatch.setattr(releases_mod, "current", lambda client, streams: []) + monkeypatch.setattr("conf.path", lambda k: tmp_path / k) + assert nlm_cli._root(None) == tmp_path / "nlm.dir" + r = runner.invoke(app, ["nlm", "fetch"]) + assert r.exit_code == 0 + + +def _stub_vsac(monkeypatch, root, ids): + monkeypatch.setattr(vsac_mod, "make_client", lambda k: _Dummy()) + + def save(client, r): + r.mkdir(parents=True, exist_ok=True) + (r / "index.jsonl").write_text( + "".join(json.dumps({"id": i}) + "\n" for i in ids) + "\n" + ) + return len(ids) + + monkeypatch.setattr(vsac_mod, "save_definitions", save) + + +class _Dummy: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + +def test_vsac_definitions_only(monkeypatch, tmp_path, key): + _stub_vsac(monkeypatch, tmp_path, ["a", "b"]) + monkeypatch.setattr( + vsac_mod, "hoard", lambda *a, **k: pytest.fail("should not expand") + ) + r = runner.invoke( + app, ["nlm", "vsac", "--dest", str(tmp_path), "--definitions-only"] + ) + assert r.exit_code == 0 and "definitions: 2" in r.output + + +def test_vsac_expands_with_limit_and_counts_codes(monkeypatch, tmp_path, key): + _stub_vsac(monkeypatch, tmp_path, [f"id{i}" for i in range(600)]) + seen = {} + + def fake_hoard(factory, root, ids, *, api_key, workers, progress): + seen.update(ids=ids, workers=workers, api_key=api_key, client=factory()) + exp = root / "expansions" + exp.mkdir(parents=True) + with gzip.open(exp / "id0.json.gz", "wt") as f: + json.dump({"expansion": {"contains": [{"code": "1"}, {"code": "2"}]}}, f) + for done in range(1, 501): + progress(done, 500) + return {"expanded": 500, "present": 0, "failed": 0} + + monkeypatch.setattr(vsac_mod, "hoard", fake_hoard) + r = runner.invoke( + app, + ["nlm", "vsac", "--dest", str(tmp_path), "--limit", "500", "--workers", "2"], + ) + assert r.exit_code == 0, r.output + assert ( + len(seen["ids"]) == 500 and seen["workers"] == 2 and seen["api_key"] == "k-123" + ) + assert isinstance(seen["client"], _Dummy) + assert "expanded 500/500" in r.output + assert "expansions: 500 new, 0 present, 0 failed" in r.output + assert "codes on disk: 2" in r.output diff --git a/tests/nlm/__init__.py b/tests/nlm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/nlm/test_releases.py b/tests/nlm/test_releases.py new file mode 100644 index 0000000..bd5dfd6 --- /dev/null +++ b/tests/nlm/test_releases.py @@ -0,0 +1,203 @@ +"""nlm.releases — UTS release listing and resumable downloads.""" + +from __future__ import annotations + +import io +import json +import zipfile + +import httpx +import pytest + +from nlm import scrub +from nlm.releases import Release, current, fetch, release_types, target + +KEY = "0123-secret-key" + + +def _zip_bytes() -> bytes: + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as z: + z.writestr("readme.txt", "hello") + return buf.getvalue() + + +def _rel(**kw) -> Release: + base = dict( + product="SNOMED CT", + stream="snomed-ct-us-edition", + name="SNOMED CT US Edition", + version="2026-09-01", + date="2026-09-01", + file_name="Snomed.zip", + download_url="https://download.nlm.nih.gov/umls/kss/Snomed.zip", + ) + return Release(**{**base, **kw}) + + +TYPES = { + "releaseTypes": [ + { + "product": "RxNorm", + "releaseType": "RxNorm Full Monthly Release", + "endpoint": "https://uts-ws.nlm.nih.gov/releases?releaseType=rxnorm-full-monthly-release", + }, + { + "product": "UMLS", + "releaseType": "UMLS Full Release", + "endpoint": "https://uts-ws.nlm.nih.gov/releases?releaseType=umls-full-release", + }, + ] +} + + +def _listing_client() -> httpx.Client: + def handler(request: httpx.Request) -> httpx.Response: + stream = request.url.params.get("releaseType") + if stream is None: + return httpx.Response(200, json=TYPES) + if stream == "umls-full-release": + return httpx.Response( + 200, + json=[ + { + "fileName": "umls-2026AA-full.zip", + "releaseVersion": "2026AA", + "releaseDate": "2026-05-03", + "downloadUrl": "https://download.nlm.nih.gov/umls/kss/2026AA/umls-2026AA-full.zip", + } + ], + ) + return httpx.Response( + 200, + json=[ + { + "fileName": "RxNorm_full_10052026.zip", + "releaseDate": "2026-10-05", + "downloadUrl": "https://download.nlm.nih.gov/rxnorm/RxNorm_full_10052026.zip", + } + ], + ) + + return httpx.Client(transport=httpx.MockTransport(handler)) + + +def test_scrub_removes_key_param_and_literal(): + msg = f"GET https://x/download?url=f.zip&apiKey={KEY} failed; key {KEY}" + out = scrub(msg, KEY) + assert KEY not in out + assert "apiKey=***" in out + assert scrub("apiKey=abc") == "apiKey=***" + + +def test_release_types_reads_the_slug_from_each_endpoint(): + with _listing_client() as c: + types = release_types(c) + assert [t["stream"] for t in types] == [ + "rxnorm-full-monthly-release", + "umls-full-release", + ] + assert types[1]["name"] == "UMLS Full Release" + + +def test_current_lists_every_stream_and_falls_back_to_the_date_for_version(): + with _listing_client() as c: + rels = current(c) + by_stream = {r.stream: r for r in rels} + assert by_stream["umls-full-release"].version == "2026AA" + assert by_stream["rxnorm-full-monthly-release"].version == "2026-10-05" + assert by_stream["rxnorm-full-monthly-release"].product == "RxNorm" + + +def test_current_filters_to_named_streams(): + with _listing_client() as c: + rels = current(c, ["umls-full-release"]) + assert [r.file_name for r in rels] == ["umls-2026AA-full.zip"] + + +def test_fetch_writes_the_zip_and_a_sidecar_then_reports_present(tmp_path): + body = _zip_bytes() + seen = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request) + return httpx.Response(200, content=body) + + rel = _rel() + with httpx.Client(transport=httpx.MockTransport(handler)) as c: + path, status = fetch(c, rel, tmp_path, KEY) + again, status2 = fetch(c, rel, tmp_path, KEY) + + assert status == "downloaded" and status2 == "present" + assert path == target(tmp_path, rel) == again + assert path.read_bytes() == body + meta = json.loads(path.with_name(path.name + ".json").read_text()) + assert meta["bytes"] == len(body) and len(meta["sha256"]) == 64 + assert meta["download_url"] == rel.download_url # the file URL, never the key + assert KEY not in json.dumps(meta) + assert len(seen) == 1 + assert seen[0].url.params["url"] == rel.download_url + + +def test_fetch_resumes_a_partial_download_with_a_range_request(tmp_path): + body = _zip_bytes() + rel = _rel() + path = target(tmp_path, rel) + path.parent.mkdir(parents=True) + part = path.with_name(path.name + ".part") + part.write_bytes(body[:10]) + + def handler(request: httpx.Request) -> httpx.Response: + assert request.headers["Range"] == "bytes=10-" + return httpx.Response(206, content=body[10:]) + + with httpx.Client(transport=httpx.MockTransport(handler)) as c: + out, status = fetch(c, rel, tmp_path, KEY) + assert status == "downloaded" and out.read_bytes() == body + + +def test_fetch_restarts_when_the_server_ignores_the_range(tmp_path): + body = _zip_bytes() + rel = _rel() + path = target(tmp_path, rel) + path.parent.mkdir(parents=True) + path.with_name(path.name + ".part").write_bytes(b"stale bytes") + + with httpx.Client( + transport=httpx.MockTransport(lambda r: httpx.Response(200, content=body)) + ) as c: + out, _ = fetch(c, rel, tmp_path, KEY) + assert out.read_bytes() == body + + +def test_fetch_rejects_a_non_zip_and_never_echoes_the_key(tmp_path): + html = f"login required for apiKey={KEY}".encode() + rel = _rel() + with httpx.Client( + transport=httpx.MockTransport(lambda r: httpx.Response(200, content=html)) + ) as c: + with pytest.raises(RuntimeError) as err: + fetch(c, rel, tmp_path, KEY) + assert "not a zip" in str(err.value) and KEY not in str(err.value) + assert not any(tmp_path.rglob("*.part")) + + +def test_fetch_http_error_is_scrubbed(tmp_path): + rel = _rel() + with httpx.Client( + transport=httpx.MockTransport(lambda r: httpx.Response(401, text="no")) + ) as c: + with pytest.raises(RuntimeError) as err: + fetch(c, rel, tmp_path, KEY) + assert "401" in str(err.value) and KEY not in str(err.value) + + +def test_fetch_keeps_non_zip_files_as_is(tmp_path): + rel = _rel( + file_name="notes.txt", download_url="https://download.nlm.nih.gov/notes.txt" + ) + with httpx.Client( + transport=httpx.MockTransport(lambda r: httpx.Response(200, content=b"plain")) + ) as c: + path, status = fetch(c, rel, tmp_path, KEY) + assert status == "downloaded" and path.read_bytes() == b"plain" diff --git a/tests/nlm/test_vsac.py b/tests/nlm/test_vsac.py new file mode 100644 index 0000000..85aafa0 --- /dev/null +++ b/tests/nlm/test_vsac.py @@ -0,0 +1,195 @@ +"""nlm.vsac — paged definitions, paged expansions, resumable hoard.""" + +from __future__ import annotations + +import gzip +import json + +import httpx +import pytest + +from nlm import vsac + +KEY = "0123-secret-key" + + +def _client(handler) -> httpx.Client: + return httpx.Client(base_url=vsac.FHIR_URL, transport=httpx.MockTransport(handler)) + + +def _vs(i: int) -> dict: + return {"resourceType": "ValueSet", "id": f"1.2.{i}-20260101", "title": f"VS {i}"} + + +def _read_gz(path): + with gzip.open(path, "rt", encoding="utf-8") as f: + return json.load(f) + + +def test_make_client_uses_basic_auth_and_fhir_json(): + c = vsac.make_client(KEY) + try: + assert str(c.base_url).rstrip("/") == vsac.FHIR_URL + assert c.headers["Accept"] == "application/fhir+json" + assert isinstance(c.auth, httpx.BasicAuth) + finally: + c.close() + + +def test_get_backs_off_on_throttling_then_succeeds(): + calls, slept = [], [] + + def handler(request): + calls.append(request) + if len(calls) == 1: + return httpx.Response(429, headers={"Retry-After": "3"}) + if len(calls) == 2: + return httpx.Response(503) + return httpx.Response(200, json={"ok": True}) + + with _client(handler) as c: + assert vsac._get(c, "/x", sleep=slept.append) == {"ok": True} + assert slept == [3.0, 2.0] + + +def test_get_gives_up_after_the_last_try(): + with _client(lambda r: httpx.Response(503)) as c: + with pytest.raises(httpx.HTTPStatusError): + vsac._get(c, "/x", tries=2, sleep=lambda s: None) + + +def test_iter_definitions_follows_next_links_over_https(): + pages = { + 0: { + "entry": [{"resource": _vs(1)}, {"resource": _vs(2)}], + "link": [ + { + "relation": "next", + "url": "http://cts.nlm.nih.gov/fhir/res/ValueSet?_offset=2", + } + ], + }, + 2: { + "entry": [{"resource": _vs(3)}], + "link": [{"relation": "self", "url": "x"}], + }, + } + + def handler(request): + assert request.url.scheme == "https" + return httpx.Response( + 200, json=pages[int(request.url.params.get("_offset", 0))] + ) + + with _client(handler) as c: + ids = [vs["id"] for vs in vsac.iter_definitions(c, page_size=2)] + assert ids == ["1.2.1-20260101", "1.2.2-20260101", "1.2.3-20260101"] + + +def test_save_definitions_writes_files_and_index_and_keeps_existing(tmp_path): + existing = tmp_path / "definitions" / "1.2.1-20260101.json.gz" + existing.parent.mkdir(parents=True) + with gzip.open(existing, "wt") as f: + json.dump({"id": "kept"}, f) + + with _client( + lambda r: httpx.Response( + 200, json={"entry": [{"resource": _vs(1)}, {"resource": _vs(2)}]} + ) + ) as c: + n = vsac.save_definitions(c, tmp_path) + + assert n == 2 + assert _read_gz(existing) == {"id": "kept"} + assert ( + _read_gz(tmp_path / "definitions" / "1.2.2-20260101.json.gz")["title"] == "VS 2" + ) + index = [ + json.loads(line) for line in (tmp_path / "index.jsonl").read_text().splitlines() + ] + assert [row["id"] for row in index] == ["1.2.1-20260101", "1.2.2-20260101"] + assert index[0]["status"] == "" + + +def _expanding(codes: list[str], *, total: bool = True, stall_at: int | None = None): + def handler(request): + offset = int(request.url.params["offset"]) + count = int(request.url.params["count"]) + page = ( + [] + if stall_at is not None and offset >= stall_at + else codes[offset : offset + count] + ) + exp = {"contains": [{"system": "s", "code": c} for c in page], "offset": offset} + if total: + exp["total"] = len(codes) + return httpx.Response(200, json={"resourceType": "ValueSet", "expansion": exp}) + + return handler + + +def test_expand_joins_pages(): + codes = [str(i) for i in range(7)] + with _client(_expanding(codes)) as c: + out = vsac.expand(c, "1.2.3", page=3) + assert [x["code"] for x in out["expansion"]["contains"]] == codes + assert "offset" not in out["expansion"] + + +def test_expand_stops_when_a_page_comes_back_empty(): + with _client(_expanding([str(i) for i in range(7)], stall_at=3)) as c: + out = vsac.expand(c, "1.2.3", page=3) + assert len(out["expansion"]["contains"]) == 3 + + +def test_expand_without_total_takes_one_page(): + with _client(_expanding(["a", "b"], total=False)) as c: + out = vsac.expand(c, "1.2.3") + assert [x["code"] for x in out["expansion"]["contains"]] == ["a", "b"] + + +def test_hoard_skips_present_expands_new_and_logs_failures(tmp_path): + present = tmp_path / "expansions" / "done.json.gz" + present.parent.mkdir(parents=True) + present.write_bytes(b"") + + def handler(request): + if "/ValueSet/bad/" in request.url.path: + return httpx.Response(404, text=f"no such value set apiKey={KEY}") + return httpx.Response( + 200, json={"expansion": {"total": 1, "contains": [{"code": "x"}]}} + ) + + made, ticks = [], [] + + def factory(): + made.append(1) + return _client(handler) + + stats = vsac.hoard( + factory, + tmp_path, + ["done", "good", "bad"], + api_key=KEY, + workers=2, + progress=lambda done, total: ticks.append((done, total)), + ) + assert stats == {"expanded": 1, "present": 1, "failed": 1} + assert _read_gz(tmp_path / "expansions" / "good.json.gz")["expansion"][ + "contains" + ] == [{"code": "x"}] + err = [ + json.loads(line) + for line in (tmp_path / "errors.jsonl").read_text().splitlines() + ] + assert err[0]["id"] == "bad" and KEY not in err[0]["error"] + assert sorted(ticks) == [(1, 2), (2, 2)] + assert 1 <= len(made) <= 2 + + +def test_hoard_without_progress_callback(tmp_path): + def handler(request): + return httpx.Response(200, json={"expansion": {"contains": []}}) + + stats = vsac.hoard(lambda: _client(handler), tmp_path, ["a"], workers=1) + assert stats["expanded"] == 1 diff --git a/uv.lock b/uv.lock index 3928fde..82cf771 100644 --- a/uv.lock +++ b/uv.lock @@ -14,7 +14,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-10-20T14:49:41.97878101Z" +exclude-newer = "2026-10-20T16:53:32.422989299Z" exclude-newer-span = "P7D" [[package]] @@ -4663,6 +4663,11 @@ mail = [ { name = "pydo" }, { name = "resend" }, ] +nlm = [ + { name = "duckdb" }, + { name = "httpx" }, + { name = "pydantic" }, +] opps = [ { name = "duckdb" }, { name = "fsspec" }, @@ -4777,6 +4782,7 @@ requires-dist = [ { name = "httpx", marker = "extra == 'conf'", specifier = ">=0.28.1" }, { name = "httpx", marker = "extra == 'llm'", specifier = ">=0.28.1" }, { name = "httpx", marker = "extra == 'mail'", specifier = ">=0.28.1" }, + { name = "httpx", marker = "extra == 'nlm'", specifier = ">=0.28.1" }, { name = "httpx", extras = ["socks"], marker = "extra == 'prisma'", specifier = ">=0.28.1" }, { name = "langchain-core", marker = "extra == 'llm'", specifier = ">=0.3.0" }, { name = "langchain-ollama", marker = "extra == 'llm'", specifier = ">=0.2.0" }, @@ -4839,6 +4845,7 @@ requires-dist = [ { name = "stack", extras = ["conf"], marker = "extra == 'bcda'" }, { name = "stack", extras = ["conf"], marker = "extra == 'bib'" }, { name = "stack", extras = ["conf"], marker = "extra == 'llm'" }, + { name = "stack", extras = ["conf"], marker = "extra == 'nlm'" }, { name = "stack", extras = ["conf"], marker = "extra == 'opps'" }, { name = "stack", extras = ["conf"], marker = "extra == 'perf'" }, { name = "stack", extras = ["conf"], marker = "extra == 'pfs'" }, @@ -4850,6 +4857,8 @@ requires-dist = [ { name = "stack", extras = ["llm"], marker = "extra == 'all'" }, { name = "stack", extras = ["llm"], marker = "extra == 'cli'" }, { name = "stack", extras = ["mail"], marker = "extra == 'cli'" }, + { name = "stack", extras = ["nlm"], marker = "extra == 'all'" }, + { name = "stack", extras = ["nlm"], marker = "extra == 'cli'" }, { name = "stack", extras = ["perf"], marker = "extra == 'all'" }, { name = "stack", extras = ["pfs"], marker = "extra == 'all'" }, { name = "stack", extras = ["pfs"], marker = "extra == 'rec'" }, @@ -4865,7 +4874,7 @@ requires-dist = [ { name = "uvicorn", marker = "extra == 'llm'", specifier = ">=0.41.0" }, { name = "xlrd", specifier = ">=2.0.2" }, ] -provides-extras = ["conf", "aco", "api", "bcda", "bib", "bls", "llm", "ccw", "cli", "mail", "cms", "opps", "pfs", "qpp", "rec", "rex", "prisma", "perf", "sem", "lake", "aws", "gcp", "azure", "all"] +provides-extras = ["conf", "aco", "api", "bcda", "bib", "bls", "llm", "ccw", "cli", "nlm", "mail", "cms", "opps", "pfs", "qpp", "rec", "rex", "prisma", "perf", "sem", "lake", "aws", "gcp", "azure", "all"] [package.metadata.requires-dev] dev = [