Skill Tracer
Skills › Scientific Agent Skills › ncats-arax › arax_client.py

arax_client.py

2109 lines · skills/ncats-arax/scripts/arax_client.py at 5f2bc03

1#!/usr/bin/env python3
2"""Bounded, provenance-preserving client for the NCATS Translator ARAX API."""
3
4from __future__ import annotations
5
6import argparse
7import email.utils
8import hashlib
9import ipaddress
10import json
11import os
12import re
13import socket
14import sys
15import tempfile
16import time
17import urllib.error
18import urllib.parse
19import urllib.request
20import uuid
21from dataclasses import dataclass
22from datetime import datetime, timezone
23from pathlib import Path
24from typing import Any, Callable, Iterable, Mapping, Sequence
27CLIENT_NAME = "ncats-arax"
28CLIENT_VERSION = "1.1"
29USER_AGENT = "scientific-agent-skills-ncats-arax/1.1"
30SUBMITTER = "scientific-agent-skills-ncats-arax"
31PRODUCTION_BASE_URL = "https://arax.transltr.io/api/arax/v1.4"
32TESTED_ARAX_VERSION = "1.5.4"
33TESTED_TRAPI_VERSION = "1.5.0"
34SUPPORTED_TRAPI_SERIES = {"1.5", "1.6"}
36GET_TIMEOUT_SECONDS = 30
37LOOKUP_TIMEOUT_SECONDS = 120
38FEDERATED_TIMEOUT_SECONDS = 180
39KP_TIMEOUT_SECONDS = 30
40LOOKUP_RESULT_LIMIT = 20
41FEDERATED_RESULT_LIMIT = 50
42MAX_RESULT_LIMIT = 50
43MAX_PROVIDERS = 5
44MIN_FEDERATED_PROVIDERS = 2
45MAX_PREDICATES = 5
46MAX_QUALIFIERS = 6
47MAX_IDENTIFIER_LENGTH = 200
48MAX_RESPONSE_BYTES = 25 * 1024 * 1024
49MAX_WARNING_MESSAGE = 500
50MAX_DISPLAY_TEXT = 500
51PUBLICATION_PREVIEW = 10
53RETRYABLE_GET_STATUS = {429, 502, 503, 504}
54CURIE_RE = re.compile(r"^[A-Za-z][A-Za-z0-9._-]*:[^\s]+$")
55BIOLINK_RE = re.compile(r"^biolink:[A-Za-z][A-Za-z0-9._-]*$")
56PROVIDER_RE = re.compile(r"^infores:[A-Za-z0-9._-]+$")
57TRAPI_VERSION_RE = re.compile(r"\bTRAPI\s+(\d+\.\d+(?:\.\d+)?)\b", re.IGNORECASE)
58EXPAND_RE = re.compile(
59 r"^expand\(edge_key=(e[01]),kp=(.+),kp_timeout=30,"
60 r"return_minimal_metadata=false\)$"
62FILTER_RE = re.compile(
63 r"^filter_results\(action=limit_number_of_results,max_results=(\d+),"
64 r"prune_kg=true\)$"
67WARNING_CODES = {
68 "PUBLIC_QUERY",
69 "NORMALIZATION_REQUIRES_CONFIRMATION",
70 "NORMALIZATION_CATEGORY_MISMATCH",
71 "NO_RESULTS",
72 "NO_PUBLICATIONS_RETURNED",
73 "NO_PRIMARY_SOURCE_RETURNED",
74 "UNSCORED_RESPONSE_ORDER",
75 "RESULT_LIMIT_REACHED",
76 "INTERNAL_PRUNING_DETECTED",
77 "KP_TIMEOUT",
78 "KP_ERROR",
79 "MALFORMED_KP_RESPONSE",
80 "MISSING_AUXILIARY_GRAPH",
81 "REVERSED_EDGE_BINDING",
82 "UNTESTED_SERVICE_VERSION",
83 "NONPRODUCTION_ENDPOINT",
87class AraxClientError(Exception):
88 """Base error carrying the documented CLI exit classification."""
90 def __init__(
91 self,
92 message: str,
93 *,
94 exit_code: int,
95 kind: str,
96 http_status: int | None = None,
97 response_body: bytes | None = None,
98 attempts: int = 0,
99 ) -> None:
100 super().__init__(message)
101 self.exit_code = exit_code
102 self.kind = kind
103 self.http_status = http_status
104 self.response_body = response_body
105 self.attempts = attempts
108class UsageError(AraxClientError):
109 def __init__(self, message: str) -> None:
110 super().__init__(message, exit_code=2, kind="invalid_input")
113class PreflightError(AraxClientError):
114 def __init__(self, message: str, **kwargs: Any) -> None:
115 super().__init__(message, exit_code=3, kind="preflight", **kwargs)
118class NormalizationError(AraxClientError):
119 def __init__(self, message: str) -> None:
120 super().__init__(message, exit_code=4, kind="normalization_no_result")
123class TransportError(AraxClientError):
124 def __init__(self, message: str, **kwargs: Any) -> None:
125 super().__init__(message, exit_code=5, kind="transport_or_http", **kwargs)
128class ResponseError(AraxClientError):
129 def __init__(self, message: str, **kwargs: Any) -> None:
130 super().__init__(message, exit_code=6, kind="invalid_response", **kwargs)
133@dataclass(frozen=True)
134class HttpResult:
135 status: int
136 body: bytes
137 headers: dict[str, str]
138 elapsed_ms: int
139 attempts: int
142@dataclass(frozen=True)
143class ServiceInfo:
144 base_url: str
145 openapi_url: str
146 arax_version: str | None
147 trapi_version: str | None
148 warnings: tuple[dict[str, Any], ...]
150 def as_dict(self, *, biolink_version: str | None = None) -> dict[str, Any]:
151 return {
152 "base_url": self.base_url,
153 "openapi_url": self.openapi_url,
154 "arax_version": self.arax_version,
155 "trapi_version": self.trapi_version,
156 "biolink_version": biolink_version,
157 }
160@dataclass(frozen=True)
161class QueryContract:
162 kind: str
163 mode: str
164 provider_ids: tuple[str, ...]
165 expand_order: str | None
166 qnode_ids: dict[str, list[str]]
167 result_limit: int
168 query_edges: dict[str, tuple[str, str]]
171class RejectRepeatedValueAction(argparse.Action):
172 """Reject repeated scalar flags instead of silently accepting the last value."""
174 def __call__(
175 self,
176 parser: argparse.ArgumentParser,
177 namespace: argparse.Namespace,
178 values: Any,
179 option_string: str | None = None,
180 ) -> None:
181 if getattr(namespace, self.dest, None) is not None:
182 parser.error(f"{option_string} may be supplied only once")
183 setattr(namespace, self.dest, values)
186def _sanitize_text(value: Any, limit: int = MAX_DISPLAY_TEXT) -> str:
187 text = str(value)
188 text = "".join(" " if ord(char) < 32 or ord(char) == 127 else char for char in text)
189 text = " ".join(text.split())
190 if len(text) > limit:
191 return text[: max(0, limit - 1)] + "…"
192 return text
195def make_warning(
196 code: str,
197 message: str,
198 context: Mapping[str, Any] | None = None,
199) -> dict[str, Any]:
200 if code not in WARNING_CODES:
201 raise ValueError(f"unknown warning code: {code}")
202 safe_context: dict[str, Any] = {}
203 for key, value in (context or {}).items():
204 if value is None or isinstance(value, (bool, int, float)):
205 safe_context[str(key)] = value
206 else:
207 safe_context[str(key)] = _sanitize_text(value, 200)
208 return {
209 "code": code,
210 "message": _sanitize_text(message, MAX_WARNING_MESSAGE),
211 "context": safe_context,
212 }
215def _append_warning(
216 warnings: list[dict[str, Any]],
217 code: str,
218 message: str,
219 context: Mapping[str, Any] | None = None,
220) -> None:
221 warning = make_warning(code, message, context)
222 identity = (warning["code"], json.dumps(warning["context"], sort_keys=True), warning["message"])
223 existing = {
224 (item.get("code"), json.dumps(item.get("context", {}), sort_keys=True), item.get("message"))
225 for item in warnings
226 }
227 if identity not in existing:
228 warnings.append(warning)
231def _has_control(value: str) -> bool:
232 return any(ord(char) < 32 or ord(char) == 127 for char in value)
235def validate_curie(value: str) -> str:
236 if not isinstance(value, str) or not value or len(value) > MAX_IDENTIFIER_LENGTH:
237 raise UsageError("CURIE must be a nonempty string of at most 200 characters")
238 if _has_control(value) or not CURIE_RE.fullmatch(value):
239 raise UsageError(f"invalid CURIE: {_sanitize_text(value, 80)}")
240 return value
243def _looks_like_curie(value: str) -> bool:
244 try:
245 validate_curie(value)
246 except UsageError:
247 return False
248 return True
251def validate_biolink_term(value: str, label: str = "Biolink term") -> str:
252 if not isinstance(value, str) or len(value) > MAX_IDENTIFIER_LENGTH:
253 raise UsageError(f"{label} must be at most 200 characters")
254 if _has_control(value) or not BIOLINK_RE.fullmatch(value):
255 raise UsageError(f"invalid {label.lower()}: {_sanitize_text(value, 80)}")
256 return value
259def validate_provider_id(value: str) -> str:
260 if not isinstance(value, str) or len(value) > MAX_IDENTIFIER_LENGTH:
261 raise UsageError("provider identifier must be at most 200 characters")
262 if _has_control(value) or not PROVIDER_RE.fullmatch(value):
263 raise UsageError(f"invalid provider identifier: {_sanitize_text(value, 80)}")
264 if value.lower() == "infores:all":
265 raise UsageError("all-provider federation is not supported")
266 return value
269def _unique(values: Sequence[str], label: str) -> list[str]:
270 if len(set(values)) != len(values):
271 raise UsageError(f"duplicate {label} values are not allowed")
272 return list(values)
275def validate_predicates(values: Sequence[str]) -> list[str]:
276 if not 1 <= len(values) <= MAX_PREDICATES:
277 raise UsageError("each query edge requires one to five predicates")
278 return _unique([validate_biolink_term(value, "Biolink predicate") for value in values], "predicate")
281def parse_qualifiers(values: Sequence[str]) -> list[tuple[str, str]]:
282 if len(values) > MAX_QUALIFIERS:
283 raise UsageError("each query edge accepts at most six qualifiers")
284 parsed: list[tuple[str, str]] = []
285 seen_types: set[str] = set()
286 for raw in values:
287 if "=" not in raw:
288 raise UsageError("qualifiers must use TYPE=VALUE")
289 qualifier_type, qualifier_value = raw.split("=", 1)
290 validate_biolink_term(qualifier_type, "Biolink qualifier type")
291 if (
292 not qualifier_value
293 or len(qualifier_value) > MAX_IDENTIFIER_LENGTH
294 or _has_control(qualifier_value)
295 ):
296 raise UsageError("qualifier values must be nonempty, control-free, and at most 200 characters")
297 if qualifier_type in seen_types:
298 raise UsageError("a qualifier type may appear only once per edge")
299 seen_types.add(qualifier_type)
300 parsed.append((qualifier_type, qualifier_value))
301 return parsed
304def validate_result_limit(value: int) -> int:
305 if not 1 <= value <= MAX_RESULT_LIMIT:
306 raise UsageError("result limit must be between 1 and 50")
307 return value
310def validate_base_url(value: str, allow_nonproduction: bool) -> tuple[str, list[dict[str, Any]]]:
311 raw = value.rstrip("/")
312 parts = urllib.parse.urlsplit(raw)
313 try:
314 parts.port
315 except ValueError as exc:
316 raise UsageError("base URL contains an invalid port") from exc
317 if parts.scheme.lower() != "https" or not parts.hostname:
318 raise UsageError("base URL must be an absolute HTTPS URL")
319 if parts.username or parts.password or parts.query or parts.fragment:
320 raise UsageError("base URL may not contain credentials, a query string, or a fragment")
321 hostname = parts.hostname.lower()
322 if hostname == "localhost" or hostname.endswith(".localhost"):
323 raise UsageError("localhost endpoints are not allowed")
324 try:
325 address = ipaddress.ip_address(hostname)
326 except ValueError:
327 address = None
328 if address is not None and not address.is_global:
329 raise UsageError("private, loopback, link-local, and reserved endpoint addresses are not allowed")
330 normalized = urllib.parse.urlunsplit(("https", parts.netloc, parts.path.rstrip("/"), "", ""))
331 warnings: list[dict[str, Any]] = []
332 if normalized != PRODUCTION_BASE_URL:
333 if not allow_nonproduction:
334 raise PreflightError(
335 "nonproduction endpoint requires --allow-nonproduction-endpoint"
336 )
337 _append_warning(
338 warnings,
339 "NONPRODUCTION_ENDPOINT",
340 "An explicitly selected nonproduction ARAX endpoint is in use.",
341 {"base_url": normalized},
342 )
343 return normalized, warnings
346def resolve_mode(
347 mode: str,
348 provider_values: Sequence[str],
349 result_limit: int | None,
350) -> tuple[list[str], int]:
351 if mode == "lookup":
352 if provider_values:
353 raise UsageError("--kp is allowed only with --mode federated")
354 providers = ["infores:rtx-kg2"]
355 limit = LOOKUP_RESULT_LIMIT if result_limit is None else result_limit
356 elif mode == "federated":
357 providers = _unique([validate_provider_id(value) for value in provider_values], "provider")
358 if not MIN_FEDERATED_PROVIDERS <= len(providers) <= MAX_PROVIDERS:
359 raise UsageError("federated mode requires two to five explicit providers")
360 limit = FEDERATED_RESULT_LIMIT if result_limit is None else result_limit
361 else:
362 raise UsageError(f"unsupported mode: {mode}")
363 return providers, validate_result_limit(limit)
366def build_query_edge(
367 subject: str,
368 object_: str,
369 predicates: Sequence[str],
370 qualifiers: Sequence[tuple[str, str]],
371) -> dict[str, Any]:
372 edge: dict[str, Any] = {
373 "subject": subject,
374 "object": object_,
375 "predicates": list(predicates),
376 }
377 if qualifiers:
378 edge["qualifier_constraints"] = [
379 {
380 "qualifier_set": [
381 {"qualifier_type_id": qualifier_type, "qualifier_value": qualifier_value}
382 for qualifier_type, qualifier_value in qualifiers
383 ]
384 }
385 ]
386 return edge
389def _provider_parameter(mode: str, provider_ids: Sequence[str]) -> str:
390 if mode == "lookup":
391 return "infores:rtx-kg2"
392 return "[" + ",".join(provider_ids) + "]"
395def build_operations(
396 edge_keys: Sequence[str],
397 mode: str,
398 provider_ids: Sequence[str],
399 result_limit: int,
400) -> dict[str, list[str]]:
401 kp_value = _provider_parameter(mode, provider_ids)
402 actions = [
403 f"expand(edge_key={edge_key},kp={kp_value},kp_timeout={KP_TIMEOUT_SECONDS},"
404 "return_minimal_metadata=false)"
405 for edge_key in edge_keys
406 ]
407 actions.extend(
408 [
409 "scoreless_resultify(ignore_edge_direction=true)",
410 f"filter_results(action=limit_number_of_results,max_results={result_limit},prune_kg=true)",
411 "return(response=true,store=false)",
412 ]
413 )
414 return {"actions": actions}
417def build_one_hop_query(
418 *,
419 subject_id: str | None,
420 subject_category: str,
421 predicates: Sequence[str],
422 object_id: str | None,
423 object_category: str,
424 qualifiers: Sequence[tuple[str, str]],
425 mode: str,
426 provider_ids: Sequence[str],
427 result_limit: int,
428) -> dict[str, Any]:
429 if subject_id is None and object_id is None:
430 raise UsageError("one-hop queries require at least one pinned endpoint")
431 nodes: dict[str, dict[str, Any]] = {
432 "n0": {"categories": [validate_biolink_term(subject_category, "Biolink category")]},
433 "n1": {"categories": [validate_biolink_term(object_category, "Biolink category")]},
434 }
435 if subject_id is not None:
436 nodes["n0"]["ids"] = [validate_curie(subject_id)]
437 if object_id is not None:
438 nodes["n1"]["ids"] = [validate_curie(object_id)]
439 body = {
440 "message": {
441 "query_graph": {
442 "nodes": nodes,
443 "edges": {
444 "e0": build_query_edge("n0", "n1", validate_predicates(predicates), qualifiers)
445 },
446 }
447 },
448 "operations": build_operations(["e0"], mode, provider_ids, result_limit),
449 "stream_progress": False,
450 "submitter": SUBMITTER,
451 }
452 return body
455def build_two_hop_query(
456 *,
457 subject_id: str,
458 subject_category: str,
459 predicates_1: Sequence[str],
460 intermediate_category: str,
461 predicates_2: Sequence[str],
462 object_id: str,
463 object_category: str,
464 qualifiers_1: Sequence[tuple[str, str]],
465 qualifiers_2: Sequence[tuple[str, str]],
466 mode: str,
467 provider_ids: Sequence[str],
468 expand_order: str,
469 result_limit: int,
470) -> dict[str, Any]:
471 if expand_order not in {"right-first", "left-first"}:
472 raise UsageError("expand order must be right-first or left-first")
473 nodes = {
474 "n0": {
475 "ids": [validate_curie(subject_id)],
476 "categories": [validate_biolink_term(subject_category, "Biolink category")],
477 },
478 "n1": {
479 "categories": [validate_biolink_term(intermediate_category, "Biolink category")]
480 },
481 "n2": {
482 "ids": [validate_curie(object_id)],
483 "categories": [validate_biolink_term(object_category, "Biolink category")],
484 },
485 }
486 edge_order = ["e1", "e0"] if expand_order == "right-first" else ["e0", "e1"]
487 return {
488 "message": {
489 "query_graph": {
490 "nodes": nodes,
491 "edges": {
492 "e0": build_query_edge(
493 "n0", "n1", validate_predicates(predicates_1), qualifiers_1
494 ),
495 "e1": build_query_edge(
496 "n1", "n2", validate_predicates(predicates_2), qualifiers_2
497 ),
498 },
499 }
500 },
501 "operations": build_operations(edge_order, mode, provider_ids, result_limit),
502 "stream_progress": False,
503 "submitter": SUBMITTER,
504 }
507def serialize_request(body: Mapping[str, Any]) -> bytes:
508 try:
509 return json.dumps(
510 body,
511 sort_keys=True,
512 separators=(",", ":"),
513 ensure_ascii=False,
514 allow_nan=False,
515 ).encode("utf-8")
516 except (TypeError, ValueError) as exc:
517 raise UsageError("request contains a value that cannot be represented as strict JSON") from exc
520def _origin(url: str) -> tuple[str, str | None, int | None]:
521 parts = urllib.parse.urlsplit(url)
522 try:
523 port = parts.port
524 except ValueError as exc:
525 raise ValueError("URL contains an invalid port") from exc
526 if port is None and parts.scheme.lower() == "https":
527 port = 443
528 return parts.scheme.lower(), parts.hostname.lower() if parts.hostname else None, port
531class SameOriginHttpsRedirectHandler(urllib.request.HTTPRedirectHandler):
532 def redirect_request(
533 self,
534 req: urllib.request.Request,
535 fp: Any,
536 code: int,
537 msg: str,
538 headers: Any,
539 newurl: str,
540 ) -> urllib.request.Request | None:
541 try:
542 original_origin = _origin(req.full_url)
543 redirect_origin = _origin(newurl)
544 except ValueError:
545 original_origin = _origin(req.full_url)
546 redirect_origin = ("", None, None)
547 if original_origin != redirect_origin or redirect_origin[0] != "https":
548 raise urllib.error.HTTPError(
549 newurl,
550 code,
551 "cross-origin or protocol-downgrade redirect rejected",
552 headers,
553 fp,
554 )
555 return super().redirect_request(req, fp, code, msg, headers, newurl)
558def _open_url(request: urllib.request.Request, timeout: int) -> Any:
559 opener = urllib.request.build_opener(SameOriginHttpsRedirectHandler())
560 return opener.open(request, timeout=timeout)
563def read_bounded_response(response: Any, byte_limit: int = MAX_RESPONSE_BYTES) -> bytes:
564 body = response.read(byte_limit + 1)
565 if len(body) > byte_limit:
566 raise ResponseError("response exceeded the 25 MiB byte limit")
567 return body
570def _response_status(response: Any) -> int:
571 status = getattr(response, "status", None)
572 return int(status if status is not None else response.getcode())
575def _headers_dict(headers: Any) -> dict[str, str]:
576 if headers is None:
577 return {}
578 try:
579 return {str(key): str(value) for key, value in headers.items()}
580 except AttributeError:
581 return {}
584def _retry_delay(headers: Mapping[str, str], now: Callable[[], datetime] | None = None) -> float:
585 raw = next((value for key, value in headers.items() if key.lower() == "retry-after"), None)
586 if raw is None:
587 return 1.0
588 try:
589 return min(10.0, max(0.0, float(raw)))
590 except ValueError:
591 try:
592 parsed = email.utils.parsedate_to_datetime(raw)
593 current = (now or (lambda: datetime.now(timezone.utc)))()
594 if parsed.tzinfo is None:
595 parsed = parsed.replace(tzinfo=timezone.utc)
596 return min(10.0, max(0.0, (parsed - current).total_seconds()))
597 except (TypeError, ValueError, OverflowError):
598 return 1.0
601def _is_timeout_error(error: BaseException) -> bool:
602 if isinstance(error, (TimeoutError, socket.timeout)):
603 return True
604 return isinstance(error, urllib.error.URLError) and isinstance(
605 getattr(error, "reason", None), (TimeoutError, socket.timeout)
606 )
609def request_get_with_retry(
610 url: str,
611 *,
612 timeout: int = GET_TIMEOUT_SECONDS,
613 sleep: Callable[[float], None] = time.sleep,
614) -> HttpResult:
615 request = urllib.request.Request(
616 url,
617 method="GET",
618 headers={
619 "Accept": "application/json",
620 "Accept-Encoding": "identity",
621 "User-Agent": USER_AGENT,
622 },
623 )
624 started = time.monotonic()
625 for attempt in (1, 2):
626 try:
627 with _open_url(request, timeout) as response:
628 try:
629 body = read_bounded_response(response)
630 except ResponseError as exc:
631 exc.attempts = attempt
632 exc.http_status = _response_status(response)
633 raise
634 return HttpResult(
635 status=_response_status(response),
636 body=body,
637 headers=_headers_dict(getattr(response, "headers", None)),
638 elapsed_ms=round((time.monotonic() - started) * 1000),
639 attempts=attempt,
640 )
641 except urllib.error.HTTPError as exc:
642 headers = _headers_dict(exc.headers)
643 if exc.code in RETRYABLE_GET_STATUS and attempt == 1:
644 try:
645 exc.close()
646 finally:
647 sleep(_retry_delay(headers))
648 continue
649 try:
650 body = read_bounded_response(exc)
651 except ResponseError as oversized:
652 oversized.http_status = exc.code
653 oversized.attempts = attempt
654 raise oversized from exc
655 raise TransportError(
656 f"GET failed with HTTP {exc.code}",
657 http_status=exc.code,
658 response_body=body,
659 attempts=attempt,
660 ) from exc
661 except (urllib.error.URLError, TimeoutError, socket.timeout) as exc:
662 if _is_timeout_error(exc) and attempt == 1:
663 sleep(1.0)
664 continue
665 raise TransportError(f"GET transport failure: {_sanitize_text(exc)}", attempts=attempt) from exc
666 raise AssertionError("unreachable")
669def post_query(url: str, request_bytes: bytes, *, timeout: int) -> HttpResult:
670 request = urllib.request.Request(
671 url,
672 data=request_bytes,
673 method="POST",
674 headers={
675 "Accept": "application/json",
676 "Accept-Encoding": "identity",
677 "Content-Type": "application/json",
678 "User-Agent": USER_AGENT,
679 },
680 )
681 started = time.monotonic()
682 try:
683 with _open_url(request, timeout) as response:
684 try:
685 body = read_bounded_response(response)
686 except ResponseError as exc:
687 exc.attempts = 1
688 exc.http_status = _response_status(response)
689 raise
690 return HttpResult(
691 status=_response_status(response),
692 body=body,
693 headers=_headers_dict(getattr(response, "headers", None)),
694 elapsed_ms=round((time.monotonic() - started) * 1000),
695 attempts=1,
696 )
697 except urllib.error.HTTPError as exc:
698 try:
699 body = read_bounded_response(exc)
700 except ResponseError as oversized:
701 oversized.http_status = exc.code
702 oversized.attempts = 1
703 raise oversized from exc
704 raise TransportError(
705 f"POST failed with HTTP {exc.code}",
706 http_status=exc.code,
707 response_body=body,
708 attempts=1,
709 ) from exc
710 except ResponseError:
711 raise
712 except (urllib.error.URLError, TimeoutError, socket.timeout) as exc:
713 raise TransportError(f"POST transport failure: {_sanitize_text(exc)}", attempts=1) from exc
716def _decode_json(body: bytes, label: str) -> Any:
717 def reject_constant(value: str) -> None:
718 raise ValueError(f"non-finite JSON number: {value}")
720 try:
721 return json.loads(body.decode("utf-8"), parse_constant=reject_constant)
722 except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as exc:
723 raise ResponseError(f"{label} was not valid strict UTF-8 JSON") from exc
726def _version_series(version: str | None) -> str | None:
727 if not version:
728 return None
729 parts = version.split(".")
730 if len(parts) < 2 or not all(part.isdigit() for part in parts[:2]):
731 return None
732 return ".".join(parts[:2])
735def _check_trapi_version(
736 trapi_version: str | None,
737 allow_untested: bool,
738 warnings: list[dict[str, Any]],
739) -> None:
740 series = _version_series(trapi_version)
741 if series not in SUPPORTED_TRAPI_SERIES and not allow_untested:
742 raise PreflightError(
743 "service reports an unknown or missing TRAPI series; use --allow-untested-version to continue"
744 )
745 if trapi_version != TESTED_TRAPI_VERSION:
746 _append_warning(
747 warnings,
748 "UNTESTED_SERVICE_VERSION",
749 "The service TRAPI version differs from the tested production version.",
750 {"reported": trapi_version, "tested": TESTED_TRAPI_VERSION},
751 )
754def parse_openapi_service_info(
755 payload: Any,
756 *,
757 base_url: str,
758 openapi_url: str,
759 allow_untested_version: bool,
760 initial_warnings: Sequence[dict[str, Any]] = (),
761) -> ServiceInfo:
762 if not isinstance(payload, dict):
763 raise PreflightError("OpenAPI document is not an object")
764 info = payload.get("info")
765 paths = payload.get("paths")
766 if not isinstance(info, dict) or not isinstance(paths, dict):
767 raise PreflightError("OpenAPI document is missing info or paths")
768 title = info.get("title")
769 if not isinstance(title, str) or "arax" not in title.lower():
770 raise PreflightError("OpenAPI title does not identify ARAX")
771 if not isinstance(paths.get("/query"), dict) or not isinstance(paths["/query"].get("post"), dict):
772 raise PreflightError("OpenAPI document does not expose POST /query")
773 if not isinstance(paths.get("/entity"), dict) or not isinstance(paths["/entity"].get("get"), dict):
774 raise PreflightError("OpenAPI document does not expose GET /entity")
775 version_match = TRAPI_VERSION_RE.search(title)
776 trapi_extension = info.get("x-trapi")
777 trapi_version = trapi_extension.get("version") if isinstance(trapi_extension, dict) else None
778 if not isinstance(trapi_version, str):
779 trapi_version = version_match.group(1) if version_match else None
780 arax_version = str(info["version"]) if info.get("version") is not None else None
781 warnings = list(initial_warnings)
782 _check_trapi_version(trapi_version, allow_untested_version, warnings)
783 if arax_version != TESTED_ARAX_VERSION:
784 _append_warning(
785 warnings,
786 "UNTESTED_SERVICE_VERSION",
787 "The ARAX version differs from the tested production version.",
788 {"reported": arax_version, "tested": TESTED_ARAX_VERSION},
789 )
790 return ServiceInfo(
791 base_url=base_url,
792 openapi_url=openapi_url,
793 arax_version=arax_version,
794 trapi_version=trapi_version,
795 warnings=tuple(warnings),
796 )
799def fetch_openapi(
800 base_url: str,
801 *,
802 allow_untested_version: bool,
803 initial_warnings: Sequence[dict[str, Any]] = (),
804) -> tuple[ServiceInfo, HttpResult]:
805 openapi_url = base_url + "/openapi.json"
806 try:
807 result = request_get_with_retry(openapi_url)
808 except TransportError as exc:
809 raise PreflightError(
810 str(exc),
811 http_status=exc.http_status,
812 response_body=exc.response_body,
813 attempts=exc.attempts,
814 ) from exc
815 try:
816 payload = _decode_json(result.body, "OpenAPI response")
817 service = parse_openapi_service_info(
818 payload,
819 base_url=base_url,
820 openapi_url=openapi_url,
821 allow_untested_version=allow_untested_version,
822 initial_warnings=initial_warnings,
823 )
824 except (PreflightError, ResponseError) as exc:
825 if isinstance(exc, ResponseError):
826 wrapped = PreflightError(str(exc))
827 else:
828 wrapped = exc
829 wrapped.http_status = result.status
830 wrapped.response_body = result.body
831 wrapped.attempts = result.attempts
832 raise wrapped
833 return service, result
836def normalize_entity_request(base_url: str, term: str) -> HttpResult:
837 query = urllib.parse.urlencode({"q": term})
838 return request_get_with_retry(base_url + "/entity?" + query)
841def parse_normalization_response(
842 payload: Any,
843 *,
844 term: str,
845 expected_category: str | None,
846 max_synonyms: int,
847 service: ServiceInfo,
848) -> tuple[dict[str, Any], bool]:
849 warnings = [dict(item) for item in service.warnings]
850 warnings.append(make_warning("PUBLIC_QUERY", "The normalization request was sent to a public service."))
851 requires_confirmation = not _looks_like_curie(term)
852 if requires_confirmation:
853 _append_warning(
854 warnings,
855 "NORMALIZATION_REQUIRES_CONFIRMATION",
856 "Review the normalized identifier and category before constructing a graph query.",
857 )
858 record: Any = payload.get(term) if isinstance(payload, dict) else None
859 if record is None and isinstance(payload, dict) and len(payload) == 1:
860 record = next(iter(payload.values()))
861 usable = isinstance(record, dict) and isinstance(record.get("id"), dict)
862 identifier_record = record.get("id", {}) if usable else {}
863 identifier = identifier_record.get("identifier") if isinstance(identifier_record, dict) else None
864 if not isinstance(identifier, str) or not _looks_like_curie(identifier):
865 usable = False
866 identifier = None
867 category = identifier_record.get("category") if isinstance(identifier_record, dict) else None
868 if expected_category is not None and category != expected_category and usable:
869 _append_warning(
870 warnings,
871 "NORMALIZATION_CATEGORY_MISMATCH",
872 "The canonical category differs from the expected category.",
873 {"expected": expected_category, "returned": category},
874 )
875 nodes = record.get("nodes", []) if isinstance(record, dict) else []
876 preview: list[dict[str, Any]] = []
877 if isinstance(nodes, list):
878 for node in nodes[:max_synonyms]:
879 if not isinstance(node, dict):
880 continue
881 preview.append(
882 {
883 "identifier": node.get("identifier", node.get("id")),
884 "name": node.get("name") or node.get("label"),
885 "category": node.get("category"),
886 }
887 )
888 if not usable:
889 _append_warning(
890 warnings,
891 "NO_RESULTS",
892 "Normalization returned no usable canonical CURIE.",
893 )
894 summary = {
895 "schema_version": "1.0",
896 "kind": "normalization",
897 "query": {
898 "input": term,
899 "expected_category": expected_category,
900 "requires_user_confirmation": requires_confirmation,
901 },
902 "service": service.as_dict(),
903 "canonical": {
904 "identifier": identifier,
905 "name": identifier_record.get("name") if isinstance(identifier_record, dict) else None,
906 "category": category,
907 },
908 "category_counts": record.get("categories", {}) if isinstance(record, dict) else {},
909 "total_synonyms": record.get("total_synonyms", 0) if isinstance(record, dict) else 0,
910 "synonym_preview": preview,
911 "warnings": warnings,
912 }
913 return summary, usable
916def _require_exact_keys(value: Mapping[str, Any], allowed: set[str], label: str) -> None:
917 extras = set(value) - allowed
918 missing = allowed - set(value)
919 if extras or missing:
920 details: list[str] = []
921 if missing:
922 details.append("missing " + ", ".join(sorted(missing)))
923 if extras:
924 details.append("unsupported " + ", ".join(sorted(extras)))
925 raise UsageError(f"{label} has " + "; ".join(details))
928def _validate_edge_contract(edge: Any) -> tuple[str, str]:
929 if not isinstance(edge, dict):
930 raise UsageError("saved request contains an invalid query edge")
931 allowed = {"subject", "object", "predicates"}
932 if "qualifier_constraints" in edge:
933 allowed.add("qualifier_constraints")
934 _require_exact_keys(edge, allowed, "saved query edge")
935 subject = edge.get("subject")
936 object_ = edge.get("object")
937 if not isinstance(subject, str) or not isinstance(object_, str):
938 raise UsageError("saved query edge lacks subject or object")
939 predicates = edge.get("predicates")
940 if not isinstance(predicates, list):
941 raise UsageError("saved query edge lacks predicates")
942 validate_predicates(predicates)
943 constraints = edge.get("qualifier_constraints")
944 if constraints is not None:
945 if not isinstance(constraints, list) or len(constraints) != 1:
946 raise UsageError("saved query edge must have one AND qualifier set")
947 if not isinstance(constraints[0], dict):
948 raise UsageError("saved query edge has invalid qualifiers")
949 _require_exact_keys(constraints[0], {"qualifier_set"}, "qualifier constraint")
950 qualifier_set = constraints[0].get("qualifier_set")
951 if not isinstance(qualifier_set, list):
952 raise UsageError("saved query edge has invalid qualifiers")
953 raw = []
954 for qualifier in qualifier_set:
955 if not isinstance(qualifier, dict):
956 raise UsageError("saved query edge has invalid qualifiers")
957 _require_exact_keys(
958 qualifier,
959 {"qualifier_type_id", "qualifier_value"},
960 "saved qualifier",
961 )
962 raw.append(f"{qualifier.get('qualifier_type_id')}={qualifier.get('qualifier_value')}")
963 parse_qualifiers(raw)
964 return subject, object_
967def _parse_kp_parameter(value: str) -> tuple[str, tuple[str, ...]]:
968 if value == "infores:rtx-kg2":
969 return "lookup", (value,)
970 if not value.startswith("[") or not value.endswith("]"):
971 raise UsageError("saved request does not use lookup or selected-provider federation")
972 providers = tuple(validate_provider_id(item) for item in value[1:-1].split(",") if item)
973 if not MIN_FEDERATED_PROVIDERS <= len(providers) <= MAX_PROVIDERS:
974 raise UsageError("saved federated request must name two to five providers")
975 _unique(list(providers), "provider")
976 return "federated", providers
979def validate_saved_request_contract(payload: Any) -> QueryContract:
980 if not isinstance(payload, dict) or "workflow" in payload:
981 raise UsageError("saved request is not a supported fixed-operation request")
982 _require_exact_keys(
983 payload,
984 {"message", "operations", "stream_progress", "submitter"},
985 "saved request",
986 )
987 if payload["stream_progress"] is not False or payload["submitter"] != SUBMITTER:
988 raise UsageError("saved request changes fixed submission controls")
989 message = payload.get("message")
990 if not isinstance(message, dict):
991 raise UsageError("saved request lacks a query graph")
992 _require_exact_keys(message, {"query_graph"}, "saved message")
993 graph = message.get("query_graph") if isinstance(message, dict) else None
994 if not isinstance(graph, dict):
995 raise UsageError("saved request lacks a query graph")
996 _require_exact_keys(graph, {"nodes", "edges"}, "saved query graph")
997 nodes = graph.get("nodes") if isinstance(graph, dict) else None
998 edges = graph.get("edges") if isinstance(graph, dict) else None
999 if not isinstance(nodes, dict) or not isinstance(edges, dict):
1000 raise UsageError("saved request lacks a query graph")
1001 if set(nodes) == {"n0", "n1"} and set(edges) == {"e0"}:
1002 kind = "one-hop"
1003 elif set(nodes) == {"n0", "n1", "n2"} and set(edges) == {"e0", "e1"}:
1004 kind = "two-hop"
1005 else:
1006 raise UsageError("saved request must be exactly one hop or endpoint-pinned two hop")
1007 qnode_ids: dict[str, list[str]] = {}
1008 for key, node in nodes.items():
1009 if not isinstance(node, dict):
1010 raise UsageError("saved request contains an invalid qnode")
1011 allowed_node_keys = {"categories"}
1012 if "ids" in node:
1013 allowed_node_keys.add("ids")
1014 _require_exact_keys(node, allowed_node_keys, "saved qnode")
1015 categories = node.get("categories")
1016 if not isinstance(categories, list) or len(categories) != 1:
1017 raise UsageError("every saved qnode must have exactly one category")
1018 validate_biolink_term(categories[0], "Biolink category")
1019 ids = node.get("ids", [])
1020 if not isinstance(ids, list) or len(ids) > 1:
1021 raise UsageError("saved qnodes may have at most one CURIE")
1022 qnode_ids[key] = [validate_curie(value) for value in ids]
1023 if kind == "one-hop" and not (qnode_ids["n0"] or qnode_ids["n1"]):
1024 raise UsageError("saved one-hop request has no pinned endpoint")
1025 if kind == "two-hop" and not (
1026 len(qnode_ids["n0"]) == 1 and not qnode_ids["n1"] and len(qnode_ids["n2"]) == 1
1027 ):
1028 raise UsageError("saved two-hop request is not endpoint pinned")
1029 query_edges = {key: _validate_edge_contract(edge) for key, edge in edges.items()}
1030 expected_edges = {"e0": ("n0", "n1")}
1031 if kind == "two-hop":
1032 expected_edges["e1"] = ("n1", "n2")
1033 if query_edges != expected_edges:
1034 raise UsageError("saved query edge topology is unsupported")
1035 operations = payload.get("operations")
1036 if not isinstance(operations, dict):
1037 raise UsageError("saved request lacks operations")
1038 _require_exact_keys(operations, {"actions"}, "saved operations")
1039 actions = operations.get("actions") if isinstance(operations, dict) else None
1040 expected_count = 4 if kind == "one-hop" else 5
1041 if not isinstance(actions, list) or len(actions) != expected_count:
1042 raise UsageError("saved request does not contain the fixed action sequence")
1043 if actions[-3] != "scoreless_resultify(ignore_edge_direction=true)":
1044 raise UsageError("saved request does not use scoreless resultification")
1045 filter_match = FILTER_RE.fullmatch(actions[-2]) if isinstance(actions[-2], str) else None
1046 if filter_match is None or actions[-1] != "return(response=true,store=false)":
1047 raise UsageError("saved request does not use the fixed filter and return actions")
1048 result_limit = validate_result_limit(int(filter_match.group(1)))
1049 expand_actions = actions[:-3]
1050 expand_keys: list[str] = []
1051 mode: str | None = None
1052 provider_ids: tuple[str, ...] | None = None
1053 for action in expand_actions:
1054 match = EXPAND_RE.fullmatch(action) if isinstance(action, str) else None
1055 if match is None:
1056 raise UsageError("saved request contains an unsupported expansion action")
1057 current_mode, current_providers = _parse_kp_parameter(match.group(2))
1058 if mode is not None and (current_mode, current_providers) != (mode, provider_ids):
1059 raise UsageError("saved request changes providers between edges")
1060 mode, provider_ids = current_mode, current_providers
1061 expand_keys.append(match.group(1))
1062 if kind == "one-hop":
1063 if expand_keys != ["e0"]:
1064 raise UsageError("saved one-hop request has invalid expansion order")
1065 expand_order = None
1066 elif expand_keys == ["e1", "e0"]:
1067 expand_order = "right-first"
1068 elif expand_keys == ["e0", "e1"]:
1069 expand_order = "left-first"
1070 else:
1071 raise UsageError("saved two-hop request has invalid expansion order")
1072 return QueryContract(
1073 kind=kind,
1074 mode=mode or "lookup",
1075 provider_ids=provider_ids or ("infores:rtx-kg2",),
1076 expand_order=expand_order,
1077 qnode_ids=qnode_ids,
1078 result_limit=result_limit,
1079 query_edges=query_edges,
1083def _normalize_node_bindings(value: Any) -> tuple[dict[str, list[dict[str, Any]]], dict[str, set[str]]]:
1084 if not isinstance(value, dict):
1085 raise ResponseError("result node_bindings is not an object")
1086 normalized: dict[str, list[dict[str, Any]]] = {}
1087 identifiers: dict[str, set[str]] = {}
1088 for qnode_key, bindings in value.items():
1089 if not isinstance(bindings, list):
1090 raise ResponseError("node binding list is malformed")
1091 normalized[qnode_key] = []
1092 identifiers[qnode_key] = set()
1093 for binding in bindings:
1094 if not isinstance(binding, dict) or not isinstance(binding.get("id"), str):
1095 raise ResponseError("node binding is missing an identifier")
1096 identifiers[qnode_key].add(binding["id"])
1097 normalized[qnode_key].append(
1099 "id": binding["id"],
1100 "query_id": binding.get("query_id"),
1101 "attributes": binding.get("attributes", []),
1104 return normalized, identifiers
1107def _resource_ids_for_role(sources: Sequence[dict[str, Any]], role: str) -> list[str]:
1108 found: list[str] = []
1109 for source in sources:
1110 resource_role = source.get("resource_role")
1111 if isinstance(resource_role, str) and resource_role.removeprefix("biolink:") == role:
1112 resource_id = source.get("resource_id")
1113 if isinstance(resource_id, str) and resource_id not in found:
1114 found.append(resource_id)
1115 return found
1118def extract_publications(edge: Mapping[str, Any]) -> list[str]:
1119 publications: list[str] = []
1120 attributes = edge.get("attributes", [])
1121 if not isinstance(attributes, list):
1122 return publications
1123 for attribute in attributes:
1124 if not isinstance(attribute, dict) or attribute.get("attribute_type_id") != "biolink:publications":
1125 continue
1126 value = attribute.get("value")
1127 candidates: Iterable[Any] = value if isinstance(value, list) else [value]
1128 for candidate in candidates:
1129 if isinstance(candidate, str) and candidate not in publications:
1130 publications.append(candidate)
1131 return publications
1134def _classify_logs(
1135 payload: Mapping[str, Any],
1136 contract: QueryContract,
1137 warnings: list[dict[str, Any]],
1138) -> bool:
1139 partial = False
1140 logs = payload.get("logs", [])
1141 if not isinstance(logs, list):
1142 return partial
1143 for log in logs:
1144 if not isinstance(log, dict):
1145 continue
1146 message = _sanitize_text(log.get("message", ""), MAX_WARNING_MESSAGE)
1147 code = _sanitize_text(log.get("code", ""), 100)
1148 level = _sanitize_text(log.get("level", ""), 30).lower()
1149 combined = f"{code} {message}".lower()
1150 provider = next((item for item in contract.provider_ids if item.lower() in combined), None)
1151 context = {"provider_id": provider} if provider else {}
1152 failure_level = level in {"warning", "error", "critical", "fatal"}
1153 timeout_failure = failure_level and (
1154 "timed out" in combined
1155 or "timeout error" in combined
1156 or "timeout failure" in combined
1157 or "timeout" in code.lower()
1159 malformed_failure = failure_level and any(
1160 token in combined for token in ("malformed", "deserial", "invalid trapi", "not a dict")
1162 if timeout_failure:
1163 _append_warning(warnings, "KP_TIMEOUT", message or "A provider timed out.", context)
1164 partial = True
1165 elif malformed_failure:
1166 _append_warning(
1167 warnings,
1168 "MALFORMED_KP_RESPONSE",
1169 message or "A provider returned a malformed response.",
1170 context,
1172 partial = True
1173 elif failure_level and (
1174 "kp" in combined or provider is not None
1175 ) and (
1176 level in {"error", "critical", "fatal"}
1177 or any(token in combined for token in (
1178 "error", "exception", "failed", "couldn't", "cannot", "no 'message'", "no knowledge graph",
1179 "skipping", "response of 4", "response of 5",
1180 ))
1181 ):
1182 _append_warning(warnings, "KP_ERROR", message or "A provider reported an error.", context)
1183 partial = True
1184 result_pruning = (
1185 "result_limit" in code.lower()
1186 or "result_prun" in code.lower()
1187 or "pruned result" in combined
1188 or ("removed" in combined and "result" in combined)
1190 if result_pruning:
1191 _append_warning(
1192 warnings,
1193 "INTERNAL_PRUNING_DETECTED",
1194 message or "ARAX reported internal pruning or result removal.",
1196 return partial
1199def parse_trapi_response(
1200 payload: Any,
1201 contract: QueryContract,
1202 *,
1203 service: ServiceInfo | None = None,
1204 initial_warnings: Sequence[dict[str, Any]] = (),
1205) -> dict[str, Any]:
1206 if not isinstance(payload, dict):
1207 raise ResponseError("TRAPI response is not an object")
1208 # ARAX can put an application error code in an otherwise valid HTTP 200
1209 # envelope. A populated graph or an empty results list does not make it a
1210 # successful query. Status is optional/nullable in TRAPI, so absent values
1211 # remain compatible with saved responses.
1212 status = payload.get("status")
1213 if status is not None and status not in ("Success", "OK"):
1214 raise ResponseError(
1215 f"ARAX reported non-success status {_sanitize_text(status, 100)}: "
1216 f"{_sanitize_text(payload.get('description') or 'inspect response.json')}"
1218 message = payload.get("message")
1219 if not isinstance(message, dict):
1220 raise ResponseError("TRAPI response is missing message")
1221 knowledge_graph = message.get("knowledge_graph")
1222 if not isinstance(knowledge_graph, dict):
1223 raise ResponseError("TRAPI response is missing knowledge_graph")
1224 kg_nodes = knowledge_graph.get("nodes")
1225 kg_edges = knowledge_graph.get("edges")
1226 results = message.get("results")
1227 if not isinstance(kg_nodes, dict) or not isinstance(kg_edges, dict) or not isinstance(results, list):
1228 raise ResponseError("TRAPI response has malformed nodes, edges, or results")
1229 auxiliary_graphs = message.get("auxiliary_graphs", {})
1230 if not isinstance(auxiliary_graphs, dict):
1231 auxiliary_graphs = {}
1232 warnings = [dict(item) for item in initial_warnings]
1233 _append_warning(
1234 warnings,
1235 "UNSCORED_RESPONSE_ORDER",
1236 "Result positions are preserved from the unscored ARAX response.",
1238 partial = _classify_logs(payload, contract, warnings)
1239 summarized_results: list[dict[str, Any]] = []
1240 analyses_count = 0
1241 bound_edges_count = 0
1242 for position, result in enumerate(results[: contract.result_limit], start=1):
1243 if not isinstance(result, dict):
1244 raise ResponseError("TRAPI result is not an object")
1245 node_bindings, binding_ids = _normalize_node_bindings(result.get("node_bindings"))
1246 if set(node_bindings) != set(contract.qnode_ids):
1247 raise ResponseError("result node_bindings do not match the supported query graph")
1248 analyses = result.get("analyses")
1249 if not isinstance(analyses, list):
1250 raise ResponseError("TRAPI result analyses is not a list")
1251 normalized_analyses: list[dict[str, Any]] = []
1252 for analysis in analyses:
1253 if not isinstance(analysis, dict):
1254 raise ResponseError("TRAPI analysis is not an object")
1255 edge_bindings = analysis.get("edge_bindings")
1256 if not isinstance(edge_bindings, dict):
1257 raise ResponseError("TRAPI analysis edge_bindings is not an object")
1258 if set(edge_bindings) != set(contract.query_edges):
1259 raise ResponseError("TRAPI analysis edge_bindings do not match the supported query graph")
1260 support_graph_ids = analysis.get("support_graphs", [])
1261 if not isinstance(support_graph_ids, list):
1262 support_graph_ids = []
1263 missing_support = [item for item in support_graph_ids if item not in auxiliary_graphs]
1264 if support_graph_ids and missing_support:
1265 support_status = "missing"
1266 for graph_id in missing_support:
1267 _append_warning(
1268 warnings,
1269 "MISSING_AUXILIARY_GRAPH",
1270 "A referenced auxiliary graph was not returned.",
1271 {"support_graph_id": graph_id},
1273 elif support_graph_ids:
1274 support_status = "available"
1275 else:
1276 support_status = "not_returned"
1277 normalized_edges: dict[str, list[dict[str, Any]]] = {}
1278 for qedge_key in contract.query_edges:
1279 bindings = edge_bindings.get(qedge_key, [])
1280 if not isinstance(bindings, list):
1281 raise ResponseError("query-edge binding list is malformed")
1282 normalized_edges[qedge_key] = []
1283 qsubject, qobject = contract.query_edges[qedge_key]
1284 expected_subjects = binding_ids.get(qsubject, set())
1285 expected_objects = binding_ids.get(qobject, set())
1286 for binding in bindings:
1287 if not isinstance(binding, dict) or not isinstance(binding.get("id"), str):
1288 raise ResponseError("query-edge binding is missing an edge identifier")
1289 edge_id = binding["id"]
1290 edge = kg_edges.get(edge_id)
1291 if not isinstance(edge, dict):
1292 raise ResponseError(f"bound knowledge-graph edge is missing: {edge_id}")
1293 subject = edge.get("subject")
1294 object_ = edge.get("object")
1295 predicate = edge.get("predicate")
1296 if not all(isinstance(item, str) for item in (subject, object_, predicate)):
1297 raise ResponseError(f"bound edge is malformed: {edge_id}")
1298 matches_direction = subject in expected_subjects and object_ in expected_objects
1299 matches_reverse = subject in expected_objects and object_ in expected_subjects
1300 if not matches_direction and not matches_reverse:
1301 raise ResponseError(
1302 f"bound edge endpoints do not match result node bindings: {edge_id}"
1304 if matches_reverse:
1305 _append_warning(
1306 warnings,
1307 "REVERSED_EDGE_BINDING",
1308 "A bound edge does not match the query's physical direction; it was preserved.",
1309 {"edge_id": edge_id, "query_edge": qedge_key},
1311 sources_value = edge.get("sources", [])
1312 sources = [dict(item) for item in sources_value if isinstance(item, dict)] if isinstance(sources_value, list) else []
1313 primary = _resource_ids_for_role(sources, "primary_knowledge_source")
1314 aggregators = _resource_ids_for_role(sources, "aggregator_knowledge_source")
1315 supporting = _resource_ids_for_role(sources, "supporting_data_source")
1316 if not primary:
1317 _append_warning(
1318 warnings,
1319 "NO_PRIMARY_SOURCE_RETURNED",
1320 "No primary knowledge source was returned for a bound edge.",
1321 {"edge_id": edge_id},
1323 publications = extract_publications(edge)
1324 if not publications:
1325 _append_warning(
1326 warnings,
1327 "NO_PUBLICATIONS_RETURNED",
1328 "No recognized publication metadata was returned for a bound edge.",
1329 {"edge_id": edge_id},
1331 qualifiers_value = edge.get("qualifiers", [])
1332 qualifiers = []
1333 if isinstance(qualifiers_value, list):
1334 qualifiers = [
1336 "qualifier_type_id": item.get("qualifier_type_id"),
1337 "qualifier_value": item.get("qualifier_value"),
1339 for item in qualifiers_value
1340 if isinstance(item, dict)
1342 subject_node = kg_nodes.get(subject, {})
1343 object_node = kg_nodes.get(object_, {})
1344 normalized_edges[qedge_key].append(
1346 "edge_id": edge_id,
1347 "subject": subject,
1348 "subject_name": subject_node.get("name") if isinstance(subject_node, dict) else None,
1349 "predicate": predicate,
1350 "object": object_,
1351 "object_name": object_node.get("name") if isinstance(object_node, dict) else None,
1352 "matches_query_direction": matches_direction,
1353 "qualifiers": qualifiers,
1354 "sources": sources,
1355 "primary_knowledge_sources": primary,
1356 "aggregator_knowledge_sources": aggregators,
1357 "supporting_data_sources": supporting,
1358 "publication_ids": publications,
1359 "publication_availability": "available" if publications else "not_returned",
1360 "support_graph_ids": list(support_graph_ids),
1361 "support_graph_status": support_status,
1364 bound_edges_count += 1
1365 normalized_analyses.append(
1367 "resource_id": analysis.get("resource_id"),
1368 "score": analysis.get("score"),
1369 "support_graphs": list(support_graph_ids),
1370 "edge_bindings": normalized_edges,
1373 analyses_count += 1
1374 summarized_results.append(
1376 "position": position,
1377 "description": _sanitize_text(result.get("description"), MAX_DISPLAY_TEXT)
1378 if result.get("description") is not None
1379 else None,
1380 "node_bindings": node_bindings,
1381 "analyses": normalized_analyses,
1384 raw_count = len(results)
1385 total_results_count = payload.get("total_results_count")
1386 if not isinstance(total_results_count, int) or isinstance(total_results_count, bool):
1387 total_results_count = None
1388 pruning = any(item.get("code") == "INTERNAL_PRUNING_DETECTED" for item in warnings)
1389 if raw_count > contract.result_limit or pruning or (
1390 total_results_count is not None and total_results_count > raw_count
1391 ):
1392 truncation_status = "confirmed"
1393 elif raw_count == contract.result_limit:
1394 truncation_status = "possible"
1395 else:
1396 truncation_status = "no"
1397 if raw_count >= contract.result_limit:
1398 _append_warning(
1399 warnings,
1400 "RESULT_LIMIT_REACHED",
1401 "The response reached or exceeded the requested result limit.",
1402 {"result_limit": contract.result_limit},
1404 if raw_count == 0:
1405 _append_warning(
1406 warnings,
1407 "NO_RESULTS",
1408 "ARAX returned no results under these constraints.",
1410 if service is None:
1411 service_data = {
1412 "base_url": None,
1413 "openapi_url": None,
1414 "arax_version": payload.get("tool_version"),
1415 "trapi_version": payload.get("schema_version"),
1416 "biolink_version": payload.get("biolink_version"),
1418 else:
1419 service_data = service.as_dict(biolink_version=payload.get("biolink_version"))
1420 return {
1421 "schema_version": "1.0",
1422 "query": {
1423 "kind": contract.kind,
1424 "mode": contract.mode,
1425 "provider_ids": list(contract.provider_ids),
1426 "expand_order": contract.expand_order,
1427 "qnode_ids": contract.qnode_ids,
1428 "result_limit": contract.result_limit,
1429 },
1430 "service": service_data,
1431 "counts": {
1432 "results_returned": raw_count,
1433 "results_summarized": len(summarized_results),
1434 "analyses_summarized": analyses_count,
1435 "bound_edges_summarized": bound_edges_count,
1436 "knowledge_graph_nodes": len(kg_nodes),
1437 "knowledge_graph_edges": len(kg_edges),
1438 "server_total_results_count": total_results_count,
1439 },
1440 "truncation_status": truncation_status,
1441 "completeness": "partial" if partial else "complete",
1442 "results": summarized_results,
1443 "warnings": warnings,
1447def render_text_summary(
1448 summary: Mapping[str, Any],
1449 *,
1450 summary_path: str = "summary.json",
1451 response_path: str = "response.json",
1452) -> str:
1453 def rendered(value: Any, limit: int = 200) -> str:
1454 return _sanitize_text(value, limit)
1456 def rendered_list(values: Any, *, limit: int = 200) -> str:
1457 if not isinstance(values, list) or not values:
1458 return "not returned"
1459 return ", ".join(rendered(value, limit) for value in values)
1461 lines: list[str] = []
1462 query = summary.get("query", {})
1463 counts = summary.get("counts", {})
1464 lines.append(
1465 f"ARAX returned {counts.get('results_returned', 0)} result(s); "
1466 f"showing {counts.get('results_summarized', 0)} in unscored response order."
1468 lines.append(
1469 f"Query: {rendered(query.get('kind'))} / {rendered(query.get('mode'))}; "
1470 f"limit {rendered(query.get('result_limit'))}; "
1471 f"truncation {rendered(summary.get('truncation_status'))}; "
1472 f"completeness {rendered(summary.get('completeness'))}."
1474 warnings = summary.get("warnings", [])
1475 if isinstance(warnings, list) and warnings:
1476 lines.append("Warnings:")
1477 for warning in warnings:
1478 if isinstance(warning, dict):
1479 lines.append(
1480 f" - {rendered(warning.get('code'), 100)}: "
1481 f"{rendered(warning.get('message', ''), MAX_WARNING_MESSAGE)}"
1483 results = summary.get("results", [])
1484 if not results:
1485 lines.append("Not returned under these constraints.")
1486 for result in results if isinstance(results, list) else []:
1487 if not isinstance(result, dict):
1488 continue
1489 lines.append(
1490 f"Unscored position {rendered(result.get('position'))}: "
1491 f"{rendered(result.get('description') or 'candidate path', MAX_DISPLAY_TEXT)}"
1493 node_bindings = result.get("node_bindings", {})
1494 if isinstance(node_bindings, dict):
1495 for qnode, bindings in node_bindings.items():
1496 ids = [item.get("id") for item in bindings if isinstance(item, dict)] if isinstance(bindings, list) else []
1497 lines.append(f" {rendered(qnode, 100)}: {rendered_list(ids)}")
1498 analyses = result.get("analyses", [])
1499 for analysis_index, analysis in enumerate(analyses if isinstance(analyses, list) else [], start=1):
1500 if not isinstance(analysis, dict):
1501 continue
1502 lines.append(
1503 f" Analysis {analysis_index}: resource {rendered(analysis.get('resource_id'))}; "
1504 f"score {rendered(analysis.get('score'))}"
1506 edge_bindings = analysis.get("edge_bindings", {})
1507 if not isinstance(edge_bindings, dict):
1508 continue
1509 for qedge, edges in edge_bindings.items():
1510 for edge in edges if isinstance(edges, list) else []:
1511 if not isinstance(edge, dict):
1512 continue
1513 subject = edge.get("subject")
1514 subject_name = edge.get("subject_name")
1515 object_ = edge.get("object")
1516 object_name = edge.get("object_name")
1517 lines.append(
1518 f" {rendered(qedge, 100)}/{rendered(edge.get('edge_id'))}: "
1519 f"{rendered(subject)} ({rendered(subject_name)}) "
1520 f"--{rendered(edge.get('predicate'))}--> "
1521 f"{rendered(object_)} ({rendered(object_name)})"
1523 if not edge.get("matches_query_direction"):
1524 lines.append(" Direction warning: physical edge orientation differs from the query.")
1525 qualifiers = edge.get("qualifiers", [])
1526 if qualifiers:
1527 rendered_qualifiers = ", ".join(
1528 f"{_sanitize_text(item.get('qualifier_type_id'), 200)}="
1529 f"{_sanitize_text(item.get('qualifier_value'), 200)}"
1530 for item in qualifiers
1531 if isinstance(item, dict)
1533 lines.append(f" Qualifiers: {rendered_qualifiers}")
1534 for label, field in (
1535 ("Primary sources", "primary_knowledge_sources"),
1536 ("Aggregator sources", "aggregator_knowledge_sources"),
1537 ("Supporting sources", "supporting_data_sources"),
1538 ):
1539 values = edge.get(field, [])
1540 lines.append(f" {label}: {rendered_list(values)}")
1541 publications = edge.get("publication_ids", [])
1542 preview = publications[:PUBLICATION_PREVIEW] if isinstance(publications, list) else []
1543 lines.append(
1544 f" Publications: {len(publications) if isinstance(publications, list) else 0}; "
1545 f"preview {rendered_list(preview)}"
1547 lines.append(f"Complete bounded publication and source lists: {rendered(summary_path, 500)}")
1548 lines.append(f"Exact TRAPI payload: {rendered(response_path, 500)}")
1549 lines.append("Candidate paths require subsequent scientific verification.")
1550 return "\n".join(lines)
1553def sha256_bytes(value: bytes) -> str:
1554 return hashlib.sha256(value).hexdigest()
1557def prepare_output_directory(value: str | Path) -> Path:
1558 path = Path(value)
1559 try:
1560 if path.exists():
1561 if not path.is_dir():
1562 raise UsageError("output path exists and is not a directory")
1563 if any(path.iterdir()):
1564 raise UsageError("output directory must be new or empty")
1565 else:
1566 path.mkdir(parents=True)
1567 except UsageError:
1568 raise
1569 except OSError as exc:
1570 raise ResponseError(f"could not prepare output directory: {_sanitize_text(exc)}") from exc
1571 return path
1574def atomic_write_bytes(path: Path, value: bytes) -> None:
1575 try:
1576 if path.exists():
1577 raise ResponseError(f"refusing to overwrite existing artifact: {path.name}")
1578 except ResponseError:
1579 raise
1580 except OSError as exc:
1581 raise ResponseError(f"could not inspect artifact path {path.name}: {_sanitize_text(exc)}") from exc
1582 descriptor: int | None = None
1583 temporary: Path | None = None
1584 try:
1585 descriptor, temporary_name = tempfile.mkstemp(
1586 prefix=".ncats-arax-", dir=str(path.parent)
1588 temporary = Path(temporary_name)
1589 try:
1590 os.chmod(temporary, 0o600)
1591 except OSError:
1592 pass
1593 handle = os.fdopen(descriptor, "wb")
1594 descriptor = None
1595 with handle:
1596 handle.write(value)
1597 handle.flush()
1598 os.fsync(handle.fileno())
1599 if path.exists():
1600 raise ResponseError(f"refusing to overwrite concurrently created artifact: {path.name}")
1601 os.replace(temporary, path)
1602 except ResponseError:
1603 raise
1604 except OSError as exc:
1605 raise ResponseError(f"could not write artifact {path.name}: {_sanitize_text(exc)}") from exc
1606 finally:
1607 if descriptor is not None:
1608 try:
1609 os.close(descriptor)
1610 except OSError:
1611 pass
1612 try:
1613 if temporary is not None:
1614 temporary.unlink(missing_ok=True)
1615 except OSError:
1616 pass
1619def atomic_write_json(path: Path, value: Mapping[str, Any]) -> bytes:
1620 try:
1621 encoded = (
1622 json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False, allow_nan=False) + "\n"
1623 ).encode("utf-8")
1624 except (TypeError, ValueError) as exc:
1625 raise ResponseError("artifact data cannot be represented as strict JSON") from exc
1626 atomic_write_bytes(path, encoded)
1627 return encoded
1630def _utc_now() -> str:
1631 return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
1634def new_manifest(
1635 command: str,
1636 *,
1637 base_url: str,
1638 privacy_acknowledged: bool,
1639 http_timeout: int,
1640 result_limit: int | None = None,
1641) -> dict[str, Any]:
1642 return {
1643 "manifest_version": "1.0",
1644 "run_id": str(uuid.uuid4()),
1645 "command": command,
1646 "started_at": _utc_now(),
1647 "finished_at": None,
1648 "execution_status": "client_error",
1649 "result_status": "not_available",
1650 "privacy_acknowledged": privacy_acknowledged,
1651 "client": {"name": CLIENT_NAME, "version": CLIENT_VERSION},
1652 "service": {
1653 "base_url": base_url,
1654 "openapi_url": base_url + "/openapi.json",
1655 "arax_version": None,
1656 "trapi_version": None,
1657 "biolink_version": None,
1658 },
1659 "request": {
1660 "method": "GET",
1661 "url": base_url + "/openapi.json",
1662 "file": None,
1663 "bytes": None,
1664 "sha256": None,
1665 },
1666 "response": {
1667 "http_status": None,
1668 "file": None,
1669 "bytes": None,
1670 "sha256": None,
1671 "elapsed_ms": None,
1672 },
1673 "limits": {
1674 "result_limit": result_limit,
1675 "response_byte_limit": MAX_RESPONSE_BYTES,
1676 "kp_timeout_seconds": KP_TIMEOUT_SECONDS if command in {"one-hop", "two-hop"} else None,
1677 "http_timeout_seconds": http_timeout,
1678 },
1679 "attempts": {"openapi_get": 0, "entity_get": 0, "query_post": 0},
1680 "artifacts": {"request": None, "response": None, "summary": None},
1681 "error": None,
1682 "warnings": [],
1686def _apply_service_to_manifest(manifest: dict[str, Any], service: ServiceInfo) -> None:
1687 manifest["service"].update(service.as_dict())
1688 manifest["warnings"] = [dict(item) for item in service.warnings]
1691def _record_response(manifest: dict[str, Any], result: HttpResult, filename: str = "response.json") -> None:
1692 manifest["response"].update(
1694 "http_status": result.status,
1695 "file": filename,
1696 "bytes": len(result.body),
1697 "sha256": sha256_bytes(result.body),
1698 "elapsed_ms": result.elapsed_ms,
1701 manifest["artifacts"]["response"] = filename
1704def _write_manifest(output_dir: Path, manifest: dict[str, Any]) -> None:
1705 manifest["finished_at"] = _utc_now()
1706 atomic_write_json(output_dir / "manifest.json", manifest)
1709def _retain_failure(
1710 output_dir: Path | None,
1711 manifest: dict[str, Any] | None,
1712 error: AraxClientError,
1713 stage: str,
1714) -> None:
1715 if output_dir is None or manifest is None:
1716 return
1717 try:
1718 manifest["execution_status"] = "http_error" if error.exit_code == 5 else "client_error"
1719 manifest["result_status"] = "not_available"
1720 manifest["error"] = {"kind": error.kind, "message": _sanitize_text(error)}
1721 if error.http_status is not None:
1722 manifest["response"]["http_status"] = error.http_status
1723 attempt_key = {
1724 "openapi": "openapi_get",
1725 "entity": "entity_get",
1726 "query": "query_post",
1727 }.get(stage)
1728 if error.attempts and attempt_key is not None:
1729 manifest["attempts"][attempt_key] = error.attempts
1730 if error.response_body is not None and not (output_dir / "response.json").exists():
1731 atomic_write_bytes(output_dir / "response.json", error.response_body)
1732 failed = HttpResult(
1733 status=error.http_status or 0,
1734 body=error.response_body,
1735 headers={},
1736 elapsed_ms=0,
1737 attempts=error.attempts,
1739 _record_response(manifest, failed)
1740 if not (output_dir / "manifest.json").exists():
1741 _write_manifest(output_dir, manifest)
1742 except Exception as manifest_error:
1743 print(
1744 f"warning: could not retain failure artifacts: {_sanitize_text(manifest_error)}",
1745 file=sys.stderr,
1749def _json_from_http(result: HttpResult, label: str) -> Any:
1750 return _decode_json(result.body, label)
1753def _preflight_summary(service: ServiceInfo) -> dict[str, Any]:
1754 return {
1755 "schema_version": "1.0",
1756 "kind": "preflight",
1757 "service": service.as_dict(),
1758 "checks": {
1759 "title_identifies_arax": True,
1760 "query_path_present": True,
1761 "version_supported": _version_series(service.trapi_version) in SUPPORTED_TRAPI_SERIES,
1762 },
1763 "warnings": [dict(item) for item in service.warnings],
1767def handle_preflight(args: argparse.Namespace) -> int:
1768 base_url, endpoint_warnings = validate_base_url(args.base_url, args.allow_nonproduction_endpoint)
1769 output_dir = prepare_output_directory(args.output_dir) if args.output_dir else None
1770 manifest = new_manifest(
1771 "preflight", base_url=base_url, privacy_acknowledged=False, http_timeout=GET_TIMEOUT_SECONDS
1772 ) if output_dir else None
1773 stage = "openapi"
1774 try:
1775 service, result = fetch_openapi(
1776 base_url,
1777 allow_untested_version=args.allow_untested_version,
1778 initial_warnings=endpoint_warnings,
1780 summary = _preflight_summary(service)
1781 if output_dir is not None and manifest is not None:
1782 _apply_service_to_manifest(manifest, service)
1783 manifest["attempts"]["openapi_get"] = result.attempts
1784 atomic_write_bytes(output_dir / "response.json", result.body)
1785 _record_response(manifest, result)
1786 atomic_write_json(output_dir / "summary.json", summary)
1787 manifest["artifacts"]["summary"] = "summary.json"
1788 manifest["execution_status"] = "success"
1789 manifest["result_status"] = "not_available"
1790 _write_manifest(output_dir, manifest)
1791 print(
1792 f"ARAX {service.arax_version or 'unknown'}; TRAPI {service.trapi_version or 'unknown'}; "
1793 f"/query available at {service.base_url}"
1795 for warning in service.warnings:
1796 print(f"warning {warning['code']}: {warning['message']}", file=sys.stderr)
1797 return 0
1798 except AraxClientError as exc:
1799 _retain_failure(output_dir, manifest, exc, stage)
1800 raise
1803def handle_normalize(args: argparse.Namespace) -> int:
1804 if not 0 <= args.max_synonyms <= 50:
1805 raise UsageError("--max-synonyms must be between 0 and 50")
1806 expected_category = (
1807 validate_biolink_term(args.expected_category, "Biolink category")
1808 if args.expected_category
1809 else None
1811 if not args.term or len(args.term) > MAX_IDENTIFIER_LENGTH or _has_control(args.term):
1812 raise UsageError("normalization term must be control-free and at most 200 characters")
1813 base_url, endpoint_warnings = validate_base_url(args.base_url, args.allow_nonproduction_endpoint)
1814 output_dir = prepare_output_directory(args.output_dir)
1815 manifest = new_manifest(
1816 "normalize", base_url=base_url, privacy_acknowledged=True, http_timeout=GET_TIMEOUT_SECONDS
1818 stage = "openapi"
1819 try:
1820 service, preflight_result = fetch_openapi(
1821 base_url,
1822 allow_untested_version=args.allow_untested_version,
1823 initial_warnings=endpoint_warnings,
1825 _apply_service_to_manifest(manifest, service)
1826 manifest["attempts"]["openapi_get"] = preflight_result.attempts
1827 stage = "entity"
1828 manifest["request"].update(
1830 "method": "GET",
1831 "url": base_url + "/entity?" + urllib.parse.urlencode({"q": args.term}),
1834 entity_result = normalize_entity_request(base_url, args.term)
1835 manifest["attempts"]["entity_get"] = entity_result.attempts
1836 atomic_write_bytes(output_dir / "response.json", entity_result.body)
1837 _record_response(manifest, entity_result)
1838 payload = _json_from_http(entity_result, "entity response")
1839 summary, usable = parse_normalization_response(
1840 payload,
1841 term=args.term,
1842 expected_category=expected_category,
1843 max_synonyms=args.max_synonyms,
1844 service=service,
1846 atomic_write_json(output_dir / "summary.json", summary)
1847 manifest["artifacts"]["summary"] = "summary.json"
1848 manifest["warnings"] = summary["warnings"]
1849 manifest["execution_status"] = "success"
1850 manifest["result_status"] = "results" if usable else "no_results"
1851 _write_manifest(output_dir, manifest)
1852 if not usable:
1853 raise NormalizationError("normalization returned no usable canonical CURIE")
1854 canonical = summary["canonical"]
1855 print(
1856 f"Canonical: {canonical['identifier']} ({canonical.get('name') or 'name not returned'}; "
1857 f"{canonical.get('category') or 'category not returned'})"
1859 if summary["query"]["requires_user_confirmation"]:
1860 print("Review and confirm this normalization before constructing a graph query.")
1861 return 0
1862 except NormalizationError:
1863 raise
1864 except AraxClientError as exc:
1865 _retain_failure(output_dir, manifest, exc, stage)
1866 raise
1869def _build_from_args(args: argparse.Namespace) -> tuple[dict[str, Any], QueryContract, int]:
1870 providers, limit = resolve_mode(args.mode, args.kp or [], args.result_limit)
1871 if args.command == "one-hop":
1872 body = build_one_hop_query(
1873 subject_id=args.subject_id,
1874 subject_category=args.subject_category,
1875 predicates=args.predicate,
1876 object_id=args.object_id,
1877 object_category=args.object_category,
1878 qualifiers=parse_qualifiers(args.qualifier or []),
1879 mode=args.mode,
1880 provider_ids=providers,
1881 result_limit=limit,
1883 timeout = LOOKUP_TIMEOUT_SECONDS if args.mode == "lookup" else FEDERATED_TIMEOUT_SECONDS
1884 else:
1885 body = build_two_hop_query(
1886 subject_id=args.subject_id,
1887 subject_category=args.subject_category,
1888 predicates_1=args.predicate_1,
1889 intermediate_category=args.intermediate_category,
1890 predicates_2=args.predicate_2,
1891 object_id=args.object_id,
1892 object_category=args.object_category,
1893 qualifiers_1=parse_qualifiers(args.qualifier_1 or []),
1894 qualifiers_2=parse_qualifiers(args.qualifier_2 or []),
1895 mode=args.mode,
1896 provider_ids=providers,
1897 expand_order=args.expand_order,
1898 result_limit=limit,
1900 timeout = LOOKUP_TIMEOUT_SECONDS if args.mode == "lookup" else FEDERATED_TIMEOUT_SECONDS
1901 return body, validate_saved_request_contract(body), timeout
1904def _validate_query_response_version(
1905 payload: Mapping[str, Any],
1906 allow_untested_version: bool,
1907 warnings: list[dict[str, Any]],
1908) -> None:
1909 try:
1910 _check_trapi_version(
1911 payload.get("schema_version") if isinstance(payload.get("schema_version"), str) else None,
1912 allow_untested_version,
1913 warnings,
1915 except PreflightError:
1916 raise
1919def handle_graph_query(args: argparse.Namespace) -> int:
1920 body, contract, timeout = _build_from_args(args)
1921 request_bytes = serialize_request(body)
1922 base_url, endpoint_warnings = validate_base_url(args.base_url, args.allow_nonproduction_endpoint)
1923 output_dir = prepare_output_directory(args.output_dir)
1924 manifest = new_manifest(
1925 args.command,
1926 base_url=base_url,
1927 privacy_acknowledged=True,
1928 http_timeout=timeout,
1929 result_limit=contract.result_limit,
1931 stage = "openapi"
1932 try:
1933 service, preflight_result = fetch_openapi(
1934 base_url,
1935 allow_untested_version=args.allow_untested_version,
1936 initial_warnings=endpoint_warnings,
1938 _apply_service_to_manifest(manifest, service)
1939 manifest["attempts"]["openapi_get"] = preflight_result.attempts
1940 stage = "query"
1941 atomic_write_bytes(output_dir / "request.json", request_bytes)
1942 manifest["request"].update(
1944 "method": "POST",
1945 "url": base_url + "/query",
1946 "file": "request.json",
1947 "bytes": len(request_bytes),
1948 "sha256": sha256_bytes(request_bytes),
1951 manifest["artifacts"]["request"] = "request.json"
1952 result = post_query(base_url + "/query", request_bytes, timeout=timeout)
1953 manifest["attempts"]["query_post"] = result.attempts
1954 atomic_write_bytes(output_dir / "response.json", result.body)
1955 _record_response(manifest, result)
1956 payload = _json_from_http(result, "TRAPI response")
1957 if not isinstance(payload, dict):
1958 raise ResponseError("TRAPI response is not an object")
1959 warnings = [dict(item) for item in service.warnings]
1960 _append_warning(warnings, "PUBLIC_QUERY", "The graph query was sent to a public service.")
1961 _validate_query_response_version(payload, args.allow_untested_version, warnings)
1962 summary = parse_trapi_response(payload, contract, service=service, initial_warnings=warnings)
1963 atomic_write_json(output_dir / "summary.json", summary)
1964 manifest["artifacts"]["summary"] = "summary.json"
1965 manifest["service"]["biolink_version"] = payload.get("biolink_version")
1966 manifest["warnings"] = summary["warnings"]
1967 manifest["execution_status"] = "success"
1968 if summary["completeness"] == "partial":
1969 manifest["result_status"] = "partial"
1970 exit_code = 7
1971 elif summary["counts"]["results_returned"] == 0:
1972 manifest["result_status"] = "no_results"
1973 exit_code = 0
1974 else:
1975 manifest["result_status"] = "results"
1976 exit_code = 0
1977 _write_manifest(output_dir, manifest)
1978 print(
1979 render_text_summary(
1980 summary,
1981 summary_path=str(output_dir / "summary.json"),
1982 response_path=str(output_dir / "response.json"),
1985 return exit_code
1986 except AraxClientError as exc:
1987 _retain_failure(output_dir, manifest, exc, stage)
1988 raise
1991def handle_summarize(args: argparse.Namespace) -> int:
1992 request_path = Path(args.request)
1993 response_path = Path(args.response)
1994 try:
1995 request_bytes = request_path.read_bytes()
1996 response_bytes = response_path.read_bytes()
1997 except OSError as exc:
1998 raise UsageError(f"could not read saved artifact: {_sanitize_text(exc)}") from exc
1999 if len(response_bytes) > MAX_RESPONSE_BYTES:
2000 raise ResponseError("saved response exceeded the 25 MiB byte limit")
2001 request_payload = _decode_json(request_bytes, "saved request")
2002 response_payload = _decode_json(response_bytes, "saved response")
2003 contract = validate_saved_request_contract(request_payload)
2004 warnings: list[dict[str, Any]] = []
2005 if isinstance(response_payload, dict):
2006 _validate_query_response_version(response_payload, False, warnings)
2007 summary = parse_trapi_response(
2008 response_payload,
2009 contract,
2010 service=None,
2011 initial_warnings=warnings,
2013 if args.format == "json":
2014 print(json.dumps(summary, indent=2, sort_keys=True, ensure_ascii=False))
2015 else:
2016 print(
2017 render_text_summary(
2018 summary,
2019 summary_path="standard output with --format json",
2020 response_path=str(response_path),
2023 return 7 if summary["completeness"] == "partial" else 0
2026def _add_network_policy_options(parser: argparse.ArgumentParser) -> None:
2027 parser.add_argument("--base-url", default=PRODUCTION_BASE_URL)
2028 parser.add_argument("--allow-nonproduction-endpoint", action="store_true")
2029 parser.add_argument("--allow-untested-version", action="store_true")
2032def _add_query_common(parser: argparse.ArgumentParser) -> None:
2033 parser.add_argument("--mode", choices=("lookup", "federated"), default="lookup")
2034 parser.add_argument("--kp", action="append", default=[])
2035 parser.add_argument("--result-limit", type=int)
2036 parser.add_argument("--acknowledge-public-query", action="store_true", required=True)
2037 parser.add_argument("--output-dir", required=True)
2038 _add_network_policy_options(parser)
2041def build_parser() -> argparse.ArgumentParser:
2042 parser = argparse.ArgumentParser(
2043 description="Bounded NCATS Translator ARAX lookup and provenance inspection"
2045 subparsers = parser.add_subparsers(dest="command", required=True)
2047 preflight = subparsers.add_parser("preflight", help="verify the ARAX OpenAPI and versions")
2048 preflight.add_argument("--output-dir")
2049 _add_network_policy_options(preflight)
2051 normalize = subparsers.add_parser("normalize", help="normalize one term or CURIE for review")
2052 normalize.add_argument("term")
2053 normalize.add_argument("--expected-category")
2054 normalize.add_argument("--max-synonyms", type=int, default=10)
2055 normalize.add_argument("--acknowledge-public-query", action="store_true", required=True)
2056 normalize.add_argument("--output-dir", required=True)
2057 _add_network_policy_options(normalize)
2059 one_hop = subparsers.add_parser("one-hop", help="run a typed one-hop lookup")
2060 one_hop.add_argument("--subject-id", action=RejectRepeatedValueAction)
2061 one_hop.add_argument("--subject-category", required=True)
2062 one_hop.add_argument("--predicate", action="append", required=True)
2063 one_hop.add_argument("--object-id", action=RejectRepeatedValueAction)
2064 one_hop.add_argument("--object-category", required=True)
2065 one_hop.add_argument("--qualifier", action="append", default=[])
2066 _add_query_common(one_hop)
2068 two_hop = subparsers.add_parser("two-hop", help="run an endpoint-pinned two-hop lookup")
2069 two_hop.add_argument("--subject-id", required=True, action=RejectRepeatedValueAction)
2070 two_hop.add_argument("--subject-category", required=True)
2071 two_hop.add_argument("--predicate-1", action="append", required=True)
2072 two_hop.add_argument("--intermediate-category", required=True)
2073 two_hop.add_argument("--predicate-2", action="append", required=True)
2074 two_hop.add_argument("--object-id", required=True, action=RejectRepeatedValueAction)
2075 two_hop.add_argument("--object-category", required=True)
2076 two_hop.add_argument("--qualifier-1", action="append", default=[])
2077 two_hop.add_argument("--qualifier-2", action="append", default=[])
2078 two_hop.add_argument(
2079 "--expand-order", choices=("right-first", "left-first"), default="right-first"
2081 _add_query_common(two_hop)
2083 summarize = subparsers.add_parser("summarize", help="inspect a saved supported ARAX response")
2084 summarize.add_argument("--request", required=True)
2085 summarize.add_argument("--response", required=True)
2086 summarize.add_argument("--format", choices=("text", "json"), default="text")
2087 return parser
2090def main(argv: Sequence[str] | None = None) -> int:
2091 parser = build_parser()
2092 args = parser.parse_args(argv)
2093 try:
2094 if args.command == "preflight":
2095 return handle_preflight(args)
2096 if args.command == "normalize":
2097 return handle_normalize(args)
2098 if args.command in {"one-hop", "two-hop"}:
2099 return handle_graph_query(args)
2100 if args.command == "summarize":
2101 return handle_summarize(args)
2102 raise UsageError(f"unsupported command: {args.command}")
2103 except AraxClientError as exc:
2104 print(f"error: {exc}", file=sys.stderr)
2105 return exc.exit_code
2108if __name__ == "__main__":
2109 raise SystemExit(main())