From 294da20dfd158be030c0555f88cad29ea4acd56d Mon Sep 17 00:00:00 2001 From: Guangrong Qin <2994757+gloriachin@users.noreply.github.com> Date: Tue, 29 Sep 2026 15:13:01 -0700 Subject: [PATCH 01/19] config: drop the 404ing nodenorm ci URL The node_normalizer CI hostname returns 404; resolve every environment to the production URL unless an explicit override is set. Co-authored-by: Guangrong Qin <2994757+gloriachin@users.noreply.github.com> --- TCT/config.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/TCT/config.py b/TCT/config.py index ad4c2f9..f716fa7 100644 --- a/TCT/config.py +++ b/TCT/config.py @@ -8,7 +8,6 @@ Environment = Literal["prod", "ci", "test"] -DEFAULT_ENVIRONMENT: Environment = "ci" @dataclass(frozen=True) @@ -33,8 +32,9 @@ def resolve(self, environment: Environment) -> str: ci="https://name-lookup.ci.transltr.io/", ), "node_normalizer": ServiceEndpoint( + # The CI hostname returns 404; use production in all environments + # unless the caller supplies an explicit service override. prod="https://nodenorm.transltr.io/", - ci="https://nodenorm.transltr.io/", ), "node_annotator": ServiceEndpoint( prod="https://annotator.transltr.io/", @@ -66,7 +66,7 @@ def resolve(self, environment: Environment) -> str: class RuntimeConfig: """Runtime environment and explicit service URL replacements.""" - environment: Environment = DEFAULT_ENVIRONMENT + environment: Environment = "ci" overrides: Mapping[str, str] = field(default_factory=dict) def __post_init__(self) -> None: @@ -95,10 +95,7 @@ def load_config( overrides: Mapping[str, str] | None = None, ) -> RuntimeConfig: """Build configuration from explicit values and ``TCT_ENVIRONMENT``.""" - selected_environment = environment or os.getenv( - "TCT_ENVIRONMENT", - DEFAULT_ENVIRONMENT, - ) + selected_environment = environment or os.getenv("TCT_ENVIRONMENT", "ci") return RuntimeConfig( environment=selected_environment, # type: ignore[arg-type] overrides=overrides or {}, From fdc62a85cae9956aa5d357f2d45efd8ec905c2da Mon Sep 17 00:00:00 2001 From: Guangrong Qin <2994757+gloriachin@users.noreply.github.com> Date: Tue, 29 Sep 2026 15:13:01 -0700 Subject: [PATCH 02/19] fix: raise informative HTTP errors from node normalization Carry the status code, URL, and response into the raised requests.HTTPError so callers see what failed instead of a bare exception. Co-authored-by: Guangrong Qin <2994757+gloriachin@users.noreply.github.com> --- TCT/node_normalizer.py | 7 ++- tests/test_node_normalizer_requests.py | 75 ++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 1 deletion(-) create mode 100644 tests/test_node_normalizer_requests.py diff --git a/TCT/node_normalizer.py b/TCT/node_normalizer.py index 2eefc6b..64a5a50 100644 --- a/TCT/node_normalizer.py +++ b/TCT/node_normalizer.py @@ -93,7 +93,12 @@ def get_normalized_nodes(query: str | list[str], return normalized_dict[query] return normalized_dict else: - raise requests.RequestException('Response from server had error, code ' + str(response.status_code)) + raise requests.HTTPError( + f"Node Normalizer request failed with status {response.status_code} " + f"for {response.url}", + response=response, + request=response.request, + ) def get_preferred_names(id_list:list[str], batch_limit=500, **kwargs) -> dict[str, str]: diff --git a/tests/test_node_normalizer_requests.py b/tests/test_node_normalizer_requests.py new file mode 100644 index 0000000..d7ccb33 --- /dev/null +++ b/tests/test_node_normalizer_requests.py @@ -0,0 +1,75 @@ +"""Offline regression tests for Node Normalizer routing and HTTP errors.""" + +import json + +import pytest +import requests + +from TCT import node_normalizer +from TCT.TCT import _resolve_node +from TCT.config import configure, reset_config + + +@pytest.fixture(autouse=True) +def clean_runtime_config(): + reset_config() + yield + reset_config() + + +@pytest.mark.parametrize("environment", ["prod", "ci", "test"]) +@pytest.mark.parametrize("mode", ["get", "post"]) +@pytest.mark.parametrize("override", [None, "https://example.org/nodenorm/"]) +def test_pathfinder_node_resolution_routes_requests( + monkeypatch, environment, mode, override +): + configure( + environment=environment, + overrides={"node_normalizer": override} if override else None, + ) + base_url = override or "https://nodenorm.transltr.io/" + calls = [] + + def fake_request(url, **kwargs): + assert url == base_url + "get_normalized_nodes" + payload = kwargs["params" if mode == "get" else "json"] + curie = payload["curie"] if mode == "get" else payload["curies"][0] + assert payload["conflate"] is False + calls.append(curie) + response = requests.Response() + response.status_code = 200 + response._content = json.dumps({ + curie: { + "id": {"identifier": curie, "label": "Disease"}, + "type": ["biolink:Disease"], + } + }).encode() + return response + + monkeypatch.setattr(node_normalizer.requests, mode, fake_request) + for curie in ["MONDO:0018881", "MONDO:0018874"]: + node = _resolve_node( + curie, node_normalizer_kwargs={"mode": mode, "conflate": False} + ) + assert node.curie == curie + assert node.categories == ["biolink:Disease"] + assert calls == ["MONDO:0018881", "MONDO:0018874"] + + +@pytest.mark.parametrize("mode", ["get", "post"]) +def test_http_error_includes_url_status_and_response(monkeypatch, mode): + response = requests.Response() + response.status_code = 404 + response.url = "https://example.org/get_normalized_nodes" + response.request = requests.Request(mode.upper(), response.url).prepare() + monkeypatch.setattr( + node_normalizer.requests, mode, lambda *args, **kwargs: response + ) + + with pytest.raises(requests.HTTPError) as error: + node_normalizer.get_normalized_nodes("MONDO:0018881", mode=mode) + + assert "404" in str(error.value) + assert response.url in str(error.value) + assert error.value.response is response + assert error.value.request is response.request From 277053b92f2bb0f12300824773a4d24c778fbf66 Mon Sep 17 00:00:00 2001 From: Guangrong Qin <2994757+gloriachin@users.noreply.github.com> Date: Tue, 29 Sep 2026 15:13:18 -0700 Subject: [PATCH 03/19] fix: fall back to test URLs for CATRAX meta knowledge graphs When the CI and dev multiomics URLs 404, try the test deployment before warning, so the Plover API registry still loads in restricted environments. Co-authored-by: Guangrong Qin <2994757+gloriachin@users.noreply.github.com> --- TCT/translator_metakg.py | 57 +++++++++++++++++++++++++++++++++++----- 1 file changed, 50 insertions(+), 7 deletions(-) diff --git a/TCT/translator_metakg.py b/TCT/translator_metakg.py index 497fcc4..cb3f6e6 100644 --- a/TCT/translator_metakg.py +++ b/TCT/translator_metakg.py @@ -168,8 +168,10 @@ def add_plover_API(APInames:dict[str, str], metaKG:pd.DataFrame) -> tuple[dict[s import requests CI_URL_PREFIX = "https://multiomics.ci.transltr.io/" DEV_URL_PREFIX = "https://multiomics.rtx.ai:9990/" + TEST_URL_PREFIX = "https://multiomics.test.transltr.io/" url = CI_URL_PREFIX + 'BigGIM_DrugResponse_PerformancePhase/meta_knowledge_graph' url_dev = DEV_URL_PREFIX + 'BigGIM_DrugResponse_PerformancePhase/meta_knowledge_graph' + url_test = TEST_URL_PREFIX + 'BigGIM_DrugResponse_PerformancePhase/meta_knowledge_graph' try: response = requests.get(url, timeout=5) if response.status_code == 200: @@ -190,7 +192,16 @@ def add_plover_API(APInames:dict[str, str], metaKG:pd.DataFrame) -> tuple[dict[s DEV_URL_PREFIX + "BigGIM_DrugResponse_PerformancePhase/query", data["edges"][i]['predicate'], data["edges"][i]['subject'], data["edges"][i]['object']) else: - print("Warning: Failed to retrieve data from both the " + CI_URL_PREFIX + "BigGIM_DrugResponse_PerformancePhase and " + DEV_URL_PREFIX + "BigGIM_DrugResponse_PerformancePhase. Status codes:", response.status_code, response_dev.status_code) + # Try the test URL if both the CI and dev URLs return a 404 + response_test = requests.get(url_test, timeout=5) + if response_test.status_code == 200: + data = response_test.json() + for i in range(len(data["edges"])): + APInames, metaKG = add_new_API_for_query(APInames, metaKG, "CATRAX BigGIM DrugResponse Performance Phase KP - TRAPI 1.5.0", + TEST_URL_PREFIX + "BigGIM_DrugResponse_PerformancePhase/query", + data["edges"][i]['predicate'], data["edges"][i]['subject'], data["edges"][i]['object']) + else: + print("Warning: Failed to retrieve data from the CI, dev, and test URLs. Status codes:", response.status_code, response_dev.status_code, response_test.status_code) else: print("Warning: Failed to retrieve data from the " + CI_URL_PREFIX + "BigGIM_DrugResponse_PerformancePhase. Status code:", response.status_code) @@ -201,6 +212,7 @@ def add_plover_API(APInames:dict[str, str], metaKG:pd.DataFrame) -> tuple[dict[s url = CI_URL_PREFIX + 'PharmacogenomicsKG/meta_knowledge_graph' url_dev = DEV_URL_PREFIX + 'PharmacogenomicsKG/meta_knowledge_graph' + url_test = TEST_URL_PREFIX + 'PharmacogenomicsKG/meta_knowledge_graph' try: response = requests.get(url) if response.status_code == 200: @@ -215,7 +227,14 @@ def add_plover_API(APInames:dict[str, str], metaKG:pd.DataFrame) -> tuple[dict[s for i in range(len(data["edges"])): APInames, metaKG = add_new_API_for_query(APInames, metaKG, "CATRAX Pharmacogenomics KP - TRAPI 1.5.0", DEV_URL_PREFIX + "PharmacogenomicsKG/query", data["edges"][i]['predicate'], data["edges"][i]['subject'], data["edges"][i]['object']) else: - print("Warning: Failed to retrieve data from both the " + CI_URL_PREFIX + "PharmacogenomicsKG and " + DEV_URL_PREFIX + "PharmacogenomicsKG. Status codes:", response.status_code, response_dev.status_code) + # Try the test URL if both the CI and dev URLs return a 404 + response_test = requests.get(url_test, timeout=5) + if response_test.status_code == 200: + data = response_test.json() + for i in range(len(data["edges"])): + APInames, metaKG = add_new_API_for_query(APInames, metaKG, "CATRAX Pharmacogenomics KP - TRAPI 1.5.0", TEST_URL_PREFIX + "PharmacogenomicsKG/query", data["edges"][i]['predicate'], data["edges"][i]['subject'], data["edges"][i]['object']) + else: + print("Warning: Failed to retrieve data from the CI, dev, and test URLs. Status codes:", response.status_code, response_dev.status_code, response_test.status_code) else: print("Warning: Failed to retrieve data from the " + CI_URL_PREFIX + "PharmacogenomicsKG. Status code:", response.status_code) except requests.exceptions.RequestException: @@ -247,6 +266,7 @@ def add_plover_API(APInames:dict[str, str], metaKG:pd.DataFrame) -> tuple[dict[s url = CI_URL_PREFIX + 'dakp/meta_knowledge_graph' url_dev = DEV_URL_PREFIX + 'dakp/meta_knowledge_graph' + url_test = TEST_URL_PREFIX + 'dakp/meta_knowledge_graph' #url = 'https://multiomics.rtx.ai:9990/dakp/meta_knowledge_graph' try: @@ -263,15 +283,23 @@ def add_plover_API(APInames:dict[str, str], metaKG:pd.DataFrame) -> tuple[dict[s for i in range(len(data["edges"])): APInames, metaKG = add_new_API_for_query(APInames, metaKG, "Drug Approvals KP - TRAPI 1.5.0", DEV_URL_PREFIX + "dakp/query", data["edges"][i]['predicate'], data["edges"][i]['subject'], data["edges"][i]['object']) else: - print("Warning: Failed to retrieve data from both the " + CI_URL_PREFIX + "dakp and " + DEV_URL_PREFIX + "dakp. Status codes:", response.status_code, response_dev.status_code) + # Try the test URL if both the CI and dev URLs return a 404 + response_test = requests.get(url_test, timeout=5) + if response_test.status_code == 200: + data = response_test.json() + for i in range(len(data["edges"])): + APInames, metaKG = add_new_API_for_query(APInames, metaKG, "Drug Approvals KP - TRAPI 1.5.0", TEST_URL_PREFIX + "dakp/query", data["edges"][i]['predicate'], data["edges"][i]['subject'], data["edges"][i]['object']) + else: + print("Warning: Failed to retrieve data from the CI, dev, and test URLs. Status codes:", response.status_code, response_dev.status_code, response_test.status_code) else: print("Warning: Failed to retrieve data from the " + CI_URL_PREFIX + "dakp. Status code:", response.status_code) except requests.exceptions.RequestException: - print("Warning: Failed to retrieve data from the " + CI_URL_PREFIX + "dakp") + print("Warning: Failed to retrieve data from the " + CI_URL_PREFIX + "dakp!" ) url = CI_URL_PREFIX + 'mokp/meta_knowledge_graph' url_dev = DEV_URL_PREFIX + 'mokp/meta_knowledge_graph' + url_test = TEST_URL_PREFIX + 'mokp/meta_knowledge_graph' #url = 'https://multiomics.rtx.ai:9990/mokp/meta_knowledge_graph' try: response = requests.get(url) @@ -289,14 +317,22 @@ def add_plover_API(APInames:dict[str, str], metaKG:pd.DataFrame) -> tuple[dict[s else: print("Warning: Failed to retrieve data from both the " + CI_URL_PREFIX + "mokp and " + DEV_URL_PREFIX + "mokp. Status codes:", response.status_code, response_dev.status_code) else: - print("Warning: Failed to retrieve data from the " + CI_URL_PREFIX + "mokp. Status code:", response.status_code) + # Try the test URL if the CI URL returns a non-404 status code + response_test = requests.get(url_test, timeout=5) + if response_test.status_code == 200: + data = response_test.json() + for i in range(len(data["edges"])): + APInames, metaKG = add_new_API_for_query(APInames, metaKG, "Multiomics KP - TRAPI 1.5.0", TEST_URL_PREFIX + "mokp/query", data["edges"][i]['predicate'], data["edges"][i]['subject'], data["edges"][i]['object']) + else: + print("Warning: Failed to retrieve data from the CI and test URLs. Status codes:", response.status_code, response_test.status_code) except requests.exceptions.RequestException: print("Warning: Failed to retrieve data from the " + CI_URL_PREFIX + "mokp") url = CI_URL_PREFIX + 'mbkp/meta_knowledge_graph' url_dev = DEV_URL_PREFIX + 'mbkp/meta_knowledge_graph' + url_test = TEST_URL_PREFIX + 'mbkp/meta_knowledge_graph' try: - response = requests.get(url) + response = requests.get(url, timeout=5) if response.status_code == 200: data = response.json() for i in range(len(data["edges"])): @@ -311,7 +347,14 @@ def add_plover_API(APInames:dict[str, str], metaKG:pd.DataFrame) -> tuple[dict[s else: print("Warning: Failed to retrieve data from both the " + CI_URL_PREFIX + "mbkp and " + DEV_URL_PREFIX + "mbkp. Status codes:", response.status_code, response_dev.status_code) else: - print("Warning: Failed to retrieve data from the " + CI_URL_PREFIX + "mbkp. Status code:", response.status_code) + # Try the test URL if the CI URL returns a non-404 status code + response_test = requests.get(url_test, timeout=5) + if response_test.status_code == 200: + data = response_test.json() + for i in range(len(data["edges"])): + APInames, metaKG = add_new_API_for_query(APInames, metaKG, "Microbiome KP - TRAPI 1.5.0", TEST_URL_PREFIX + "mbkp/query", data["edges"][i]['predicate'], data["edges"][i]['subject'], data["edges"][i]['object']) + else: + print("Warning: Failed to retrieve data from the CI and test URLs. Status codes:", response.status_code, response_test.status_code) except requests.exceptions.RequestException: print("Warning: Failed to retrieve data from the " + CI_URL_PREFIX + "mbkp") From 317329524fbf9c59797618f215ede39509a93908 Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 15:13:51 -0700 Subject: [PATCH 04/19] config: add ARS service endpoints for prod, ci, and test --- TCT/config.py | 7 +++++++ tests/test_config.py | 4 ++++ 2 files changed, 11 insertions(+) diff --git a/TCT/config.py b/TCT/config.py index f716fa7..5c3cc67 100644 --- a/TCT/config.py +++ b/TCT/config.py @@ -59,6 +59,13 @@ def resolve(self, environment: Environment) -> str: ci="https://shepherd.ci.transltr.io/aragorn/query", test="https://shepherd.test.transltr.io/aragorn/query", ), + "ars": ServiceEndpoint( + # Autonomous Relay System REST root; TCT.Query_ARS appends submit/ and + # messages/ paths to this base. + prod="https://ars-prod.transltr.io/ars/api/", + ci="https://ars.ci.transltr.io/ars/api/", + test="https://ars.test.transltr.io/ars/api/", + ), } diff --git a/tests/test_config.py b/tests/test_config.py index 5068119..871125b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -26,6 +26,10 @@ def test_prod_ci_and_test_endpoint_resolution(): assert ci.service_url("arax") == "https://shepherd.ci.transltr.io/arax/query" assert test.service_url("arax") == "https://shepherd.test.transltr.io/arax/query" + assert prod.service_url("ars") == "https://ars-prod.transltr.io/ars/api/" + assert ci.service_url("ars") == "https://ars.ci.transltr.io/ars/api/" + assert test.service_url("ars") == "https://ars.test.transltr.io/ars/api/" + def test_environment_specific_and_fallback_endpoint_resolution(): ci = RuntimeConfig(environment="ci") From 36360f67143c65ad93a52a83ea69856cbcefd917 Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 15:21:16 -0700 Subject: [PATCH 05/19] feat: add ARS neighborhood and pathfinder finders returning ARSResult Reshape TCT.Query_ARS around Guangrong Qin's feat/ars-query-support module: shared requests.Session transport, deadline-based merge-aware polling (Relay#621), frozen ARSStatus with an O(1) child index, and an ARSResult that extends FinderResult so both finders behave like the other TCT finders. format_query_json_forARS_pathfinder reuses the shared pathfinder query-graph builder. Co-authored-by: Guangrong Qin <2994757+gloriachin@users.noreply.github.com> Co-authored-by: Frankie Hodges Co-authored-by: Frankie Hodges <129905251+hodgesf@users.noreply.github.com> --- TCT/Query_ARS.py | 781 ++++++++++++++++++++++++++++++++++++++++ TCT/TCT.py | 17 +- tests/ars_fixtures.py | 114 ++++++ tests/test_query_ars.py | 383 ++++++++++++++++++++ 4 files changed, 1294 insertions(+), 1 deletion(-) create mode 100644 TCT/Query_ARS.py create mode 100644 tests/ars_fixtures.py create mode 100644 tests/test_query_ars.py diff --git a/TCT/Query_ARS.py b/TCT/Query_ARS.py new file mode 100644 index 0000000..df71fbb --- /dev/null +++ b/TCT/Query_ARS.py @@ -0,0 +1,781 @@ +"""Submit TRAPI queries to ARS and retrieve their merged responses. + +The Autonomous Relay System (ARS) accepts one TRAPI query, fans it out to every +registered ARA, and merges their answers. Submission is asynchronous: + +1. ``POST {ars}/submit`` returns the *parent* message ``pk``. +2. ``GET {ars}/messages/?trace=y`` reports the parent ``status`` + (``Running`` -> ``Done`` or ``Error``), one ``children`` entry per ARA plus + one for the ARS merge agent, and ``merged_version``: the pk of the merged + message once the ARS has combined the ARA answers. +3. ``GET {ars}/messages/`` returns the merged message; its + ``fields.data.message`` holds the combined TRAPI ``knowledge_graph``, + ``results``, and ``auxiliary_graphs``. + +The parent can read ``Done`` a few seconds before the merge agent's child has +finished saving (NCATSTranslator/Relay#621), so :func:`wait_for_ARS` waits for +both before returning. + +The API root comes from :mod:`TCT.config` (the ``ars`` service), so +``TCT_ENVIRONMENT`` or :func:`TCT.configure` selects prod, ci, or test. +""" + +from __future__ import annotations + +import ast +import logging +import time +from dataclasses import dataclass, field +from typing import Any, Optional, Union + +import requests + +from .config import service_url +from .TCT import ( + CategoryList, + FinderResult, + NodeInput, + _build_finder_result, + _normalize_categories, + _resolve_nodes, +) +from .TCT_pathfinder import build_query_graph + +logger = logging.getLogger(__name__) + +MERGE_AGENT = "ars-ars-agent" +"""Agent name of the ARS child that holds the merged answer.""" + +TERMINAL_STATUSES = frozenset({"Done", "Error"}) +"""Parent statuses after which the ARS will not change a message again.""" + +HTTP_TIMEOUT = 120.0 + +__all__ = [ + "ARSChild", + "ARSError", + "ARSPendingError", + "ARSResult", + "ARSStatus", + "ARSTimeoutError", + "MERGE_AGENT", + "ARS_neighborhood_finder", + "ARS_pathfinder", + "Query_ARS", + "check_ars_results", + "fetch_ars_results", + "format_query_json_forARS_neighborhood", + "format_query_json_forARS_pathfinder", + "get_ARS_result", + "get_ARS_status", + "submit_ARS", + "summarize_results", + "wait_for_ARS", +] + + +# --------------------------------------------------------------------------- +# errors +# --------------------------------------------------------------------------- +class ARSError(RuntimeError): + """Raised when the ARS rejects a request or reports a failed query.""" + + +class ARSTimeoutError(ARSError, TimeoutError): + """Raised when a query does not finish within the allowed time.""" + + def __init__(self, status: "ARSStatus", timeout: float) -> None: + self.status = status + self.timeout = timeout + super().__init__( + f"ARS query {status.pk} did not finish within {timeout:.0f}s " + f"(status={status.status!r}; {status.summary()})" + ) + + +class ARSPendingError(ARSError): + """Raised when results are requested before a query reaches a merged state.""" + + def __init__(self, status: "ARSStatus") -> None: + self.status = status + super().__init__( + f"ARS query {status.pk} is not finished yet (status={status.status!r}); " + "poll get_ARS_status or use wait_for_ARS first" + ) + + +# --------------------------------------------------------------------------- +# status model +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class ARSChild: + """One agent's entry in a parent message trace.""" + + agent: str + pk: Optional[str] + status: Optional[str] + code: Optional[int] + result_count: Optional[int] + infores: Optional[str] = None + + +@dataclass(frozen=True) +class ARSStatus: + """Parsed ``?trace=y`` view of a parent ARS message.""" + + pk: str + status: Optional[str] + code: Optional[int] + merged_version: Optional[str] + merged_versions_list: tuple[tuple[str, ...], ...] = () + children: tuple[ARSChild, ...] = () + result_count: Optional[int] = None + _children_by_agent: dict[str, ARSChild] = field( + default_factory=dict, init=False, repr=False, compare=False + ) + + def __post_init__(self) -> None: + # Index children once so lookups by agent name are O(1). + object.__setattr__( + self, "_children_by_agent", {child.agent: child for child in self.children} + ) + + @property + def is_terminal(self) -> bool: + """True once the parent reports ``Done`` or ``Error``.""" + return self.status in TERMINAL_STATUSES + + @property + def merge_child(self) -> Optional[ARSChild]: + """The child entry written by the ARS merge agent, if present.""" + return self._children_by_agent.get(MERGE_AGENT) + + @property + def merged_ready(self) -> bool: + """True when a merged message exists and the merge agent is ``Done``.""" + merge = self.merge_child + return bool(self.merged_version) and merge is not None and merge.status == "Done" + + def summary(self) -> str: + """One-line summary of child statuses for logs and error messages.""" + parts = [f"{child.agent}={child.status}" for child in self.children] + return ", ".join(parts) if parts else "no children yet" + + +@dataclass +class ARSResult(FinderResult): + """Merged ARS answer wrapped like the other finder results. + + Extends :class:`TCT.TCT.FinderResult` with the ARS identifiers, the final + status, and a ranked row summarizer. ``raw`` holds the merged TRAPI + message; ``query``, ``knowledge_graph``, ``results``, ``auxiliary_graphs``, + ``resolved_nodes``, and ``to_dict()`` behave exactly like the other finders. + """ + + pk: str = "" + merged_pk: Optional[str] = None + status: Optional[ARSStatus] = None + + def summarize(self, top_n: Optional[int] = 20) -> list[dict[str, Any]]: + """Return up to ``top_n`` ranked, agent-friendly result rows.""" + return summarize_results(self.raw, top_n=top_n) + + +# --------------------------------------------------------------------------- +# transport +# --------------------------------------------------------------------------- +_SHARED_SESSION: Optional[requests.Session] = None + + +def _session() -> requests.Session: + """Return the shared HTTP session (keep-alive across ARS calls).""" + global _SHARED_SESSION + if _SHARED_SESSION is None: + _SHARED_SESSION = requests.Session() + return _SHARED_SESSION + + +def _ars_root(url: Optional[str] = None) -> str: + """Resolve the ARS API root, defaulting to the configured ``ars`` service.""" + root = (url or service_url("ars")).rstrip("/") + return root + + +def _post(url: str, json_body: dict[str, Any]) -> requests.Response: + return _session().post(url, json=json_body, timeout=HTTP_TIMEOUT) + + +def _get(url: str, params: Optional[dict[str, str]] = None) -> requests.Response: + return _session().get(url, params=params, timeout=HTTP_TIMEOUT) + + +def _api_root(url: Optional[str]) -> str: + root = (url or service_url("ars")).rstrip("/") + return root + + +def _message_from_envelope(envelope: dict[str, Any]) -> Optional[dict[str, Any]]: + """Extract the merged TRAPI message from a full ARS message envelope.""" + data = (envelope.get("fields") or {}).get("data") or {} + message = data.get("message") if isinstance(data, dict) else None + return message if isinstance(message, dict) else None + + +# --------------------------------------------------------------------------- +# query construction +# --------------------------------------------------------------------------- +def format_query_json_forARS_neighborhood( + subject_ids: list[str], + object_ids: list[str] | None = None, + subject_categories: list[str] | None = None, + object_categories: list[str] | None = None, + predicates: list[str] | None = None, + attribute_constraints: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + """Build a one-hop TRAPI query for callers without uploaded JSON.""" + if not subject_ids: + raise ValueError("subject_ids must contain at least one CURIE") + if not object_ids and not object_categories: + raise ValueError("object_ids or object_categories must be provided") + + subject: dict[str, Any] = {"ids": subject_ids} + object_node: dict[str, Any] = {} + if subject_categories: + subject["categories"] = subject_categories + if object_ids: + object_node["ids"] = object_ids + if object_categories: + object_node["categories"] = object_categories + + edge: dict[str, Any] = { + "subject": "n00", "object": "n01", + "predicates": predicates or ["biolink:related_to"], + } + if attribute_constraints: + edge["attribute_constraints"] = attribute_constraints + + return { + "message": {"query_graph": { + "edges": {"e00": edge}, + "nodes": {"n00": subject, "n01": object_node}, + }}, + "submitter": "TCT", + } + + +def format_query_json_forARS_pathfinder( + start_node_id: str, + end_node_id: str, + subject_categories: list[str] | None = None, + object_categories: list[str] | None = None, + constraints_path: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + """Build a paths TRAPI query from the shared pathfinder query graph. + + ``constraints_path`` optionally restricts intermediate path nodes; see + :func:`TCT.TCT_pathfinder.build_query_graph` for the accepted shape. + """ + if not start_node_id or not end_node_id: + raise ValueError("start_node_id and end_node_id are required") + graph = build_query_graph( + start_node_id=start_node_id, + end_node_id=end_node_id, + start_node_categories=subject_categories, + end_node_categories=object_categories, + constraints_path=constraints_path, + ) + return {"message": {"query_graph": graph}, "submitter": "TCT"} + + +# --------------------------------------------------------------------------- +# submit / poll / fetch +# --------------------------------------------------------------------------- +def submit_ARS(query_json: dict[str, Any], *, url: str | None = None) -> str: + """Validate and submit a TRAPI query to the ARS; return the parent pk.""" + message = query_json.get("message") if isinstance(query_json, dict) else None + if not isinstance(message, dict) or not isinstance(message.get("query_graph"), dict): + raise ValueError("ARS input must be a JSON object with message.query_graph") + + response = _post(f"{_ars_root(url)}/submit", query_json) + if response.status_code not in (200, 201): + raise ARSError( + f"ARS submit failed with HTTP {response.status_code}: {response.text[:500]}" + ) + pk = response.json().get("pk") + if not isinstance(pk, str) or not pk: + raise ValueError("ARS submission did not return a message pk") + logger.info("ARS query submitted: pk=%s", pk) + return pk + + +def Query_ARS(url: str, json_file: dict[str, Any]) -> str: + """Submit a TRAPI JSON object to ARS and return its message identifier.""" + return submit_ARS(json_file, url=url) + + +def _as_version_list(value: Any) -> tuple[tuple[str, ...], ...]: + """Normalize ``merged_versions_list``. + + The ARS serializes this field as the Python repr of a list of + ``[pk, agent]`` pairs (a string) rather than JSON, so parse that form and + accept a real list as well. + """ + if not value: + return () + if isinstance(value, str): + try: + value = ast.literal_eval(value) + except (ValueError, SyntaxError): + return ((value,),) + if not isinstance(value, (list, tuple)): + return () + return tuple( + tuple(str(part) for part in item) if isinstance(item, (list, tuple)) else (str(item),) + for item in value + ) + + +def _to_int(value: Any) -> Optional[int]: + try: + return int(value) if value not in (None, "") else None + except (TypeError, ValueError): + return None + + +def parse_trace(pk: str, trace: dict[str, Any]) -> ARSStatus: + """Convert a ``?trace=y`` response into an :class:`ARSStatus`.""" + children = tuple( + ARSChild( + agent=str((child.get("actor") or {}).get("agent") or ""), + pk=child.get("message"), + status=child.get("status"), + code=_to_int(child.get("code")), + result_count=_to_int(child.get("result_count")), + infores=(child.get("actor") or {}).get("inforesid") or None, + ) + for child in trace.get("children") or [] + ) + return ARSStatus( + pk=pk, + status=trace.get("status"), + code=_to_int(trace.get("code")), + merged_version=trace.get("merged_version") or None, + merged_versions_list=_as_version_list(trace.get("merged_versions_list")), + children=children, + result_count=_to_int(trace.get("result_count")), + ) + + +def get_message(pk: str, *, trace: bool = False, url: str | None = None) -> dict[str, Any]: + """Fetch one ARS message by pk, optionally in lightweight trace form.""" + params = {"trace": "y"} if trace else None + response = _get(f"{_ars_root(url)}/messages/{pk}", params=params) + if response.status_code == 404: + raise LookupError(f"ARS message not found: {pk}") + if response.status_code != 200: + raise ARSError( + f"ARS messages/{pk} failed with HTTP " + f"{response.status_code}: {response.text[:500]}" + ) + return response.json() + + +def get_ARS_status(pk: str, *, url: str | None = None) -> ARSStatus: + """Return the current status of a submitted query and its ARA children.""" + return parse_trace(pk, get_message(pk, trace=True, url=url)) + + +def check_ars_results( + response_pk: str, + *, + url: str | None = None, + max_retries: int = 15, + poll_interval: float = 30, + merge_grace: float = 30.0, +) -> str | None: + """Poll the parent ARS message and return its merged message identifier. + + Return ``None`` when ARS finishes without a merged result. Raise + :class:`ARSTimeoutError` (a ``TimeoutError``) if the message never reaches a + terminal state. The parent reads ``Done`` before the merge agent has saved + (NCATSTranslator/Relay#621), so polling continues until the merge child is + also ``Done`` or ``merge_grace`` seconds elapse after the parent finished. + """ + return wait_for_ARS( + response_pk, + url=url, + max_retries=max_retries, + poll_interval=poll_interval, + merge_grace=merge_grace, + ).merged_version + + +def wait_for_ARS( + pk: str, + *, + url: str | None = None, + max_retries: int = 15, + poll_interval: float = 30.0, + merge_grace: float = 30.0, +) -> ARSStatus: + """Poll a parent message until the ARS has finished and merged the answers. + + Returns when the parent status is ``Done`` **and** the merge agent's child + reports ``Done``. If the parent is ``Done`` but no merged message becomes + ready within ``merge_grace`` seconds (for example, a query with no ARA + results), the last observed status is returned instead. + + Raises + ------ + ARSError + If the parent finishes with status ``Error``. + ARSTimeoutError + After ``max_retries`` polls spaced ``poll_interval`` seconds apart. + """ + if not pk: + raise ValueError("An ARS message pk is required") + if max_retries < 1: + raise ValueError("max_retries must be at least 1") + + deadline_polls = time.monotonic() + max_retries * poll_interval + parent_done_at: Optional[float] = None + attempt = 0 + while True: + attempt += 1 + status = get_ARS_status(pk, url=url) + logger.info("ARS %s: %s (%s)", pk, status.status, status.summary()) + if status.status == "Error": + raise ARSError(f"ARS query {pk} finished with status Error ({status.summary()})") + if status.is_terminal: + if status.merged_ready: + return status + now = time.monotonic() + if parent_done_at is None: + parent_done_at = now + if now - parent_done_at >= merge_grace: + logger.warning( + "ARS %s: parent Done but merged message not ready after %.0fs; " + "returning as is", + pk, + merge_grace, + ) + return status + if attempt >= max_retries or time.monotonic() >= deadline_polls: + raise ARSTimeoutError(status, max_retries * poll_interval) + time.sleep(poll_interval) + + +def fetch_ars_results(merged_pk: str, *, url: str | None = None) -> dict[str, Any]: + """Fetch the full merged ARS message, including ``fields.data.message``.""" + if not merged_pk: + raise ValueError("A merged ARS message pk is required") + response = _get(f"{_ars_root(url)}/messages/{merged_pk}") + if response.status_code != 200: + raise ARSError( + f"ARS messages/{merged_pk} failed with HTTP " + f"{response.status_code}: {response.text[:500]}" + ) + return response.json() + + +def get_ARS_result( + pk_or_status: Union[str, ARSStatus], + *, + url: str | None = None, +) -> ARSResult: + """Fetch the merged answer for a finished query as an :class:`ARSResult`. + + Falls back to the parent message's own payload when the ARS produced no + merged version. + """ + status = ( + pk_or_status + if isinstance(pk_or_status, ARSStatus) + else get_ARS_status(pk_or_status, url=url) + ) + if not status.is_terminal: + raise ARSPendingError(status) + target = status.merged_version or status.pk + message = _message_from_envelope(fetch_ars_results(target, url=url)) + return ARSResult( + pk=status.pk, + merged_pk=status.merged_version, + status=status, + raw=message or {}, + query={}, + knowledge_graph=message.get("knowledge_graph") or {} if message else {}, + results=list(message.get("results") or []) if message else [], + auxiliary_graphs=message.get("auxiliary_graphs") or {} if message else {}, + resolved_nodes={}, + ) + + +# --------------------------------------------------------------------------- +# result summarization +# --------------------------------------------------------------------------- +def _primary_sources(edge: dict[str, Any]) -> list[str]: + return sorted( + { + source.get("resource_id") + for source in edge.get("sources") or [] + if source.get("resource_role") == "primary_knowledge_source" + and source.get("resource_id") + } + ) + + +def summarize_results( + message: Optional[dict[str, Any]], + *, + top_n: Optional[int] = 20, +) -> list[dict[str, Any]]: + """Flatten a merged ARS message into ranked, agent-friendly rows. + + Each row carries the result rank and score, the answer node (``essence``), + every bound node with its name and categories, and the predicates and + primary knowledge sources of the edges bound in the result's analyses. + Only the first ``top_n`` results are visited; ``None`` returns all rows. + """ + if not message: + return [] + kg = message.get("knowledge_graph") or {} + nodes = kg.get("nodes") or {} + edges = kg.get("edges") or {} + results = list(message.get("results") or []) + if top_n is not None: + results = results[:top_n] + + rows: list[dict[str, Any]] = [] + for rank, result in enumerate(results, start=1): + bound_nodes = { + qnode: [ + { + "id": binding.get("id"), + "name": (nodes.get(binding.get("id")) or {}).get("name"), + "categories": (nodes.get(binding.get("id")) or {}).get("categories"), + } + for binding in bindings + if binding.get("id") + ] + for qnode, bindings in (result.get("node_bindings") or {}).items() + } + predicates: set[str] = set() + sources: set[str] = set() + aras: set[str] = set() + edge_count = 0 + for analysis in result.get("analyses") or []: + if analysis.get("resource_id"): + aras.add(analysis["resource_id"]) + for bindings in (analysis.get("edge_bindings") or {}).values(): + for binding in bindings: + edge = edges.get(binding.get("id")) or {} + if not edge: + continue + edge_count += 1 + if edge.get("predicate"): + predicates.add(edge["predicate"]) + sources.update(_primary_sources(edge)) + rows.append( + { + "rank": result.get("rank", rank), + "score": result.get("normalized_score", result.get("score")), + "essence": result.get("essence"), + "essence_category": result.get("essence_category"), + "nodes": bound_nodes, + "predicates": sorted(predicates), + "primary_sources": sorted(sources), + "aras": sorted(aras), + "edge_count": edge_count, + } + ) + return rows + + +# --------------------------------------------------------------------------- +# finder-style entry points +# --------------------------------------------------------------------------- +def _empty_raw(query_graph: dict[str, Any]) -> dict[str, Any]: + """Shell TRAPI message used when the ARS produced no merged message.""" + return { + "query_graph": query_graph, + "knowledge_graph": {"nodes": {}, "edges": {}}, + "results": [], + "auxiliary_graphs": {}, + } + + +def _collect_result( + pk: str, + query_json: dict[str, Any], + *, + resolved_nodes: dict[str, Any], + url: str | None, + max_retries: int, + poll_interval: float, + merge_grace: float, +) -> ARSResult: + """Run submit -> wait -> fetch and wrap the outcome as an ``ARSResult``.""" + status = wait_for_ARS( + pk, + url=url, + max_retries=max_retries, + poll_interval=poll_interval, + merge_grace=merge_grace, + ) + target = status.merged_version or status.pk + message = _message_from_envelope(fetch_ars_results(target, url=url)) + raw = message if message else _empty_raw(query_json["message"]["query_graph"]) + base = _build_finder_result(raw, resolved_nodes=resolved_nodes) + return ARSResult( + query=base.query, + knowledge_graph=base.knowledge_graph, + results=base.results, + auxiliary_graphs=base.auxiliary_graphs, + resolved_nodes=base.resolved_nodes, + raw=base.raw, + pk=pk, + merged_pk=status.merged_version, + status=status, + ) + + +def _node_keyed_inputs( + node: Union[NodeInput, list[NodeInput]], + resolved_nodes: list, +) -> dict[str, Any]: + """Key resolved inputs ``node`` for a single input, ``node_{i}`` for lists.""" + if isinstance(node, str): + return {"node": resolved_nodes[0]} + return {f"node_{index}": resolved for index, resolved in enumerate(resolved_nodes)} + + +def ARS_neighborhood_finder( + node: Union[NodeInput, list[NodeInput], None] = None, + neighbor_categories: CategoryList | None = None, + *, + json_file: dict[str, Any] | None = None, + node_categories: CategoryList | None = None, + predicates: list[str] | None = None, + attribute_constraints: list[dict[str, Any]] | None = None, + max_retries: int = 15, + poll_interval: float = 30.0, + merge_grace: float = 30.0, + name_resolver_kwargs: dict[str, Any] | None = None, + node_normalizer_kwargs: dict[str, Any] | None = None, +) -> ARSResult: + """Find one-hop neighbors for one or more concepts by asking the ARS. + + Unlike :func:`TCT.neighborhood_finder`, this does not load the MetaKG or + pick knowledge providers itself; the ARS fans the query out to every ARA + and merges their answers. + + Pass the parsed contents of an uploaded JSON file as ``json_file``; it is + submitted unchanged, so both ``edges`` and ``paths`` query graphs work. + Otherwise ``node`` and ``neighbor_categories`` build a one-hop query, with + names resolved to CURIEs first (the ARS accepts unresolved names at submit + time and then returns no results). + + Returns + ------- + ARSResult + The merged answer wrapped like the other finders. ``raw`` is the + merged TRAPI message and ``resolved_nodes`` records the input CURIEs. + """ + if json_file is not None: + if node is not None or neighbor_categories is not None: + raise ValueError("Provide json_file or node and neighbor_categories") + query = json_file + resolved: dict[str, Any] = {} + else: + if node is None: + raise ValueError("Provide json_file or node") + resolved_list = _resolve_nodes( + node, + name_resolver_kwargs=name_resolver_kwargs, + node_normalizer_kwargs=node_normalizer_kwargs, + ) + query = format_query_json_forARS_neighborhood( + subject_ids=[resolved.curie for resolved in resolved_list], + subject_categories=_normalize_categories(node_categories), + object_categories=_normalize_categories(neighbor_categories), + predicates=predicates, + attribute_constraints=attribute_constraints, + ) + resolved = _node_keyed_inputs(node, resolved_list) + + pk = submit_ARS(query) + return _collect_result( + pk, + query, + resolved_nodes=resolved, + url=None, + max_retries=max_retries, + poll_interval=poll_interval, + merge_grace=merge_grace, + ) + + +def ARS_pathfinder( + start: NodeInput, + end: NodeInput, + *, + intermediate_categories: CategoryList | None = None, + predicates: list[str] | None = None, + json_file: dict[str, Any] | None = None, + attribute_constraints: list[dict[str, Any]] | None = None, + max_retries: int = 15, + poll_interval: float = 30.0, + merge_grace: float = 30.0, + name_resolver_kwargs: dict[str, Any] | None = None, + node_normalizer_kwargs: dict[str, Any] | None = None, +) -> ARSResult: + """Find paths between two concepts by asking the ARS. + + Like :func:`ARS_neighborhood_finder`, the ARS fans the query out to every + ARA instead of using locally selected knowledge providers. At most one + ``intermediate_categories`` value is applied, mirroring the local + pathfinder. A ``json_file`` with a complete TRAPI query (``paths`` or + ``edges`` query graph) is submitted unchanged instead. + + Returns + ------- + ARSResult + The merged answer wrapped like the other finders; ``resolved_nodes`` + records the resolved ``start`` and ``end`` inputs. + """ + if json_file is not None: + query = json_file + resolved: dict[str, Any] = {} + else: + resolved_list = _resolve_nodes( + [start, end], + name_resolver_kwargs=name_resolver_kwargs, + node_normalizer_kwargs=node_normalizer_kwargs, + ) + intermediate = _normalize_categories(intermediate_categories) + if intermediate and len(intermediate) > 1: + logger.warning( + "ARS pathfinder supports a single intermediate category; using %r", + intermediate[0], + ) + query = format_query_json_forARS_pathfinder( + resolved_list[0].curie, + resolved_list[1].curie, + subject_categories=resolved_list[0].categories or None, + object_categories=resolved_list[1].categories or None, + constraints_path=( + [{"intermediate_categories": [intermediate[0]]}] if intermediate else None + ), + ) + resolved = {"start": resolved_list[0], "end": resolved_list[1]} + + pk = submit_ARS(query) + return _collect_result( + pk, + query, + resolved_nodes=resolved, + url=None, + max_retries=max_retries, + poll_interval=poll_interval, + merge_grace=merge_grace, + ) diff --git a/TCT/TCT.py b/TCT/TCT.py index c7bb2a6..1ae3c9d 100644 --- a/TCT/TCT.py +++ b/TCT/TCT.py @@ -51,7 +51,12 @@ #'query_chatGPT4', 'load_json_template', 'extract_json', - 'TRAPI_json_validation' + 'TRAPI_json_validation', + 'ARS_neighborhood_finder', + 'ARS_pathfinder', + 'ARSResult', + 'check_ars_results', + 'fetch_ars_results', ] @@ -2267,3 +2272,13 @@ def get_similar_predicate(query_json_cur_clean, All_predicates): # --------------------------------------------------------------------------- from .TCT_pathfinder import query_TCT_pathfinder # noqa: E402 from .TCT_neighborhood_finder import neighborhood_finder # noqa: E402 +# NOTE: the Query_ARS *function* is deliberately not re-exported here; the name +# must stay free for ``import TCT.Query_ARS`` (the module). It remains available +# as ``TCT.Query_ARS.Query_ARS``. +from .Query_ARS import ( # noqa: E402 + ARSResult, + ARS_neighborhood_finder, + ARS_pathfinder, + check_ars_results, + fetch_ars_results, +) diff --git a/tests/ars_fixtures.py b/tests/ars_fixtures.py new file mode 100644 index 0000000..e641ab6 --- /dev/null +++ b/tests/ars_fixtures.py @@ -0,0 +1,114 @@ +"""Shared ARS fixtures: structural shapes of live ARS responses. + +Adapted from the offline tests on ``feat/ars-client`` (Frankie Hodges). The +fixture-drift test in ``tests/test_ars_live.py`` asserts these structures stay +a subset of a real ARS trace and merged envelope when ``TCT_LIVE_ARS=1``. +""" + +from __future__ import annotations + +BASE = "https://ars.ci.transltr.io/ars/api/" +PARENT = "11111111-1111-1111-1111-111111111111" +MERGED = "22222222-2222-2222-2222-222222222222" + + +class FakeResponse: + """Minimal ``requests.Response`` stand-in.""" + + def __init__(self, payload, status_code=200): + self._payload = payload + self.status_code = status_code + self.text = str(payload) + + def json(self): + return self._payload + + +def child(agent, status, pk="child", result_count=None, code=200): + """One ``children`` entry of a parent trace.""" + return { + "actor": {"agent": agent, "inforesid": f"infores:{agent}"}, + "status": status, + "code": code, + "message": pk, + "result_count": result_count, + } + + +def trace(status, merged=None, children=()): + """A ``?trace=y`` parent response.""" + return { + "status": status, + "code": 200 if status == "Done" else 202, + "merged_version": merged, + "merged_versions_list": [], + "children": list(children), + } + + +def done_trace(merged=MERGED, merge_status="Done"): + """A parent trace whose merge agent has finished.""" + return trace( + "Done", + merged, + [ + child("ars-ars-agent", merge_status, pk=merged, result_count="3"), + child("ara-shepherd-arax", "Done", result_count=3), + ], + ) + + +def merged_message(): + """A small but structurally complete merged TRAPI message.""" + return { + "query_graph": {"nodes": {}, "edges": {}}, + "knowledge_graph": { + "nodes": { + "MONDO:1": {"name": "disease", "categories": ["biolink:Disease"]}, + "CHEBI:1": {"name": "drug", "categories": ["biolink:SmallMolecule"]}, + }, + "edges": { + "e1": { + "subject": "CHEBI:1", + "object": "MONDO:1", + "predicate": "biolink:treats", + "sources": [ + { + "resource_role": "primary_knowledge_source", + "resource_id": "infores:a", + }, + { + "resource_role": "aggregator_knowledge_source", + "resource_id": "infores:agg", + }, + ], + } + }, + }, + "results": [ + { + "rank": 1, + "normalized_score": 0.9, + "essence": "drug", + "essence_category": "biolink:SmallMolecule", + "node_bindings": {"n0": [{"id": "MONDO:1"}], "n1": [{"id": "CHEBI:1"}]}, + "analyses": [ + { + "resource_id": "infores:arax", + "edge_bindings": {"e0": [{"id": "e1"}]}, + } + ], + }, + {"rank": 2, "score": 0.1, "node_bindings": {}, "analyses": []}, + ], + "auxiliary_graphs": {}, + } + + +def envelope(message, pk=MERGED): + """A full ARS message envelope wrapping a merged TRAPI message.""" + return { + "model": "tr_ars.message", + "pk": pk, + "fields": {"status": "Done", "data": {"message": message}}, + } diff --git a/tests/test_query_ars.py b/tests/test_query_ars.py new file mode 100644 index 0000000..b862a70 --- /dev/null +++ b/tests/test_query_ars.py @@ -0,0 +1,383 @@ +"""Offline tests for the ARS client and its finder APIs (TCT.Query_ARS). + +The uploaded-JSON and notebook-parity cases come from ``feat/ars-query-support`` +(Guangrong Qin); the fixtures and merge-aware polling cases are adapted from +``feat/ars-client`` (Frankie Hodges). +""" + +from __future__ import annotations + +import importlib + +import pytest + +ars = importlib.import_module("TCT.Query_ARS") +from TCT import ARS_neighborhood_finder, ARS_pathfinder # noqa: E402 +from TCT.TCT import FinderResult, ResolvedNode # noqa: E402 +from tests.ars_fixtures import ( # noqa: E402 + BASE, + MERGED, + PARENT, + FakeResponse, + child, + done_trace, + envelope, + merged_message, + trace, +) + + +@pytest.fixture +def no_sleep(monkeypatch): + calls = [] + monkeypatch.setattr(ars.time, "sleep", lambda seconds: calls.append(seconds)) + return calls + + +def _resolved(values): + """ResolvedNode stand-ins keyed to the input order.""" + return [ + ResolvedNode( + input_value=value, + curie=f"MONDO:{index}", + label=value, + categories=["biolink:Disease"], + ) + for index, value in enumerate([values] if isinstance(values, str) else values) + ] + + +# --------------------------------------------------------------------------- # +# query construction +# --------------------------------------------------------------------------- # +def test_one_hop_input_matches_notebook_query(monkeypatch): + captured = {} + + def fake_query(query_json, *, url=None): + captured.setdefault("query", query_json) + return PARENT + + status = ars.parse_trace(PARENT, done_trace()) + message = merged_message() + monkeypatch.setattr(ars, "submit_ARS", fake_query) + monkeypatch.setattr(ars, "wait_for_ARS", lambda pk, **kwargs: status) + monkeypatch.setattr( + ars, "fetch_ars_results", lambda pk, *, url=None: envelope(message) + ) + + result = ARS_neighborhood_finder("MONDO:0005148", ["biolink:Drug"]) + + assert captured["query"] == { + "message": {"query_graph": { + "edges": {"e00": {"subject": "n00", "object": "n01", "predicates": ["biolink:related_to"]}}, + "nodes": {"n00": {"ids": ["MONDO:0005148"]}, "n01": {"categories": ["biolink:Drug"]}}, + }}, + "submitter": "TCT", + } + assert result.pk == PARENT + assert result.merged_pk == MERGED + assert result.resolved_nodes["node"].curie == "MONDO:0005148" + + +def test_pathfinder_query_uses_shared_query_graph_builder(monkeypatch): + graph = {"nodes": {"sn": {}, "on": {}}, "paths": {"p0": {}}} + calls = [] + + def build_query_graph(**kwargs): + calls.append(kwargs) + return graph + + monkeypatch.setattr(ars, "build_query_graph", build_query_graph) + + query = ars.format_query_json_forARS_pathfinder( + "PUBCHEM.COMPOUND:132212657", + "MONDO:0018874", + ["biolink:SmallMolecule"], + ["biolink:Disease"], + [{"intermediate_categories": ["biolink:Gene"]}], + ) + + assert query == { + "message": {"query_graph": graph}, + "submitter": "TCT", + } + assert calls == [{ + "start_node_id": "PUBCHEM.COMPOUND:132212657", + "end_node_id": "MONDO:0018874", + "start_node_categories": ["biolink:SmallMolecule"], + "end_node_categories": ["biolink:Disease"], + "constraints_path": [{"intermediate_categories": ["biolink:Gene"]}], + }] + + +def test_pathfinder_query_requires_both_endpoint_curies(): + with pytest.raises(ValueError, match="start_node_id and end_node_id"): + ars.format_query_json_forARS_pathfinder("", "MONDO:0018874") + + +def test_pathfinder_resolves_and_keys_resolved_nodes(monkeypatch): + submitted = {} + status = ars.parse_trace(PARENT, done_trace()) + message = merged_message() + + monkeypatch.setattr( + ars, "_resolve_nodes", lambda values, **kwargs: _resolved(values) + ) + + def fake_query(query_json, *, url=None): + submitted.update(query_json) + return PARENT + + monkeypatch.setattr(ars, "submit_ARS", fake_query) + monkeypatch.setattr(ars, "wait_for_ARS", lambda pk, **kwargs: status) + monkeypatch.setattr(ars, "fetch_ars_results", lambda pk, *, url=None: envelope(message)) + + result = ARS_pathfinder( + "asthma", + "albuterol", + intermediate_categories=["Gene", "Protein"], + ) + + graph = submitted["message"]["query_graph"] + assert graph["nodes"]["sn"]["ids"] == ["MONDO:0"] + assert graph["nodes"]["on"]["ids"] == ["MONDO:1"] + # Only the first intermediate category is applied, mirroring the local pathfinder. + assert graph["paths"]["p0"]["constraints"] == [{"intermediate_categories": ["biolink:Gene"]}] + assert result.resolved_nodes["start"].curie == "MONDO:0" + assert result.resolved_nodes["end"].curie == "MONDO:1" + + +def test_invalid_uploaded_json_is_rejected_before_submission(monkeypatch): + monkeypatch.setattr(ars, "_post", lambda *args, **kwargs: pytest.fail("unexpected POST")) + with pytest.raises(ValueError, match="message.query_graph"): + ARS_neighborhood_finder(json_file={"message": None}) + + +def test_mutually_exclusive_inputs_are_rejected(): + with pytest.raises(ValueError, match="json_file or node"): + ARS_neighborhood_finder( + json_file={"message": {"query_graph": {}}}, + node=["MONDO:1"], + neighbor_categories=["biolink:Drug"], + ) + + +def test_multi_input_neighborhood_finder_keys_node_index(monkeypatch): + status = ars.parse_trace(PARENT, done_trace()) + message = merged_message() + monkeypatch.setattr(ars, "_resolve_nodes", lambda values, **kwargs: _resolved(values)) + monkeypatch.setattr(ars, "submit_ARS", lambda query, *, url=None: PARENT) + monkeypatch.setattr(ars, "wait_for_ARS", lambda pk, **kwargs: status) + monkeypatch.setattr(ars, "fetch_ars_results", lambda pk, *, url=None: envelope(message)) + + result = ARS_neighborhood_finder( + ["MONDO:0004979", "MONDO:0004979"], ["biolink:Drug"] + ) + + assert list(result.resolved_nodes) == ["node_0", "node_1"] + + +# --------------------------------------------------------------------------- # +# submit / fetch / trace parsing +# --------------------------------------------------------------------------- # +def test_submit_uses_configured_ars_url(monkeypatch): + posted = {} + + def fake_post(url, json_body): + posted["url"] = url + posted["json"] = json_body + return FakeResponse({"pk": PARENT, "fields": {"status": "Running"}}, 201) + + monkeypatch.setattr(ars, "_post", fake_post) + + pk = ars.submit_ARS({"message": {"query_graph": {}}}, url=BASE.rstrip("/")) + + assert pk == PARENT + assert posted["url"] == BASE + "submit" + + +def test_submit_rejects_http_failures_and_missing_pk(monkeypatch): + monkeypatch.setattr(ars, "_post", lambda *a, **k: FakeResponse({"detail": "bad"}, 500)) + with pytest.raises(ars.ARSError, match="HTTP 500"): + ars.submit_ARS({"message": {"query_graph": {}}}) + + monkeypatch.setattr(ars, "_post", lambda *a, **k: FakeResponse({"fields": {}}, 201)) + with pytest.raises(ValueError, match="did not return a message pk"): + ars.submit_ARS({"message": {"query_graph": {}}}) + + +def test_parse_trace_reads_children_and_merge_state(): + status = ars.parse_trace(PARENT, done_trace()) + + assert status.status == "Done" + assert status.merged_version == MERGED + assert [c.agent for c in status.children] == ["ars-ars-agent", "ara-shepherd-arax"] + assert status.merge_child.result_count == 3 + assert status.merged_ready + assert status.is_terminal + assert "ars-ars-agent=Done" in status.summary() + + # The ARS serializes merged_versions_list as a Python repr string. + repr_trace = done_trace() + repr_trace["merged_versions_list"] = ( + f"[['{MERGED}', 'ars'], ['child', 'ara-shepherd-arax']]" + ) + parsed = ars.parse_trace(PARENT, repr_trace) + assert parsed.merged_versions_list == ((MERGED, "ars"), ("child", "ara-shepherd-arax")) + assert ars.parse_trace( + PARENT, {**repr_trace, "merged_versions_list": [[MERGED, "ars"]]} + ).merged_versions_list == ((MERGED, "ars"),) + assert ars.parse_trace( + PARENT, {**repr_trace, "merged_versions_list": "not a list"} + ).merged_versions_list == (("not a list",),) + + running = ars.parse_trace(PARENT, trace("Running")) + assert not running.is_terminal + assert not running.merged_ready + assert running.summary() == "no children yet" + + +def test_get_message_maps_404(monkeypatch): + monkeypatch.setattr(ars, "_get", lambda *a, **k: FakeResponse({}, 404)) + with pytest.raises(LookupError): + ars.get_message("missing") + + +# --------------------------------------------------------------------------- # +# polling +# --------------------------------------------------------------------------- # +def _poll_sequence(monkeypatch, traces): + """Serve successive trace responses; the last one repeats.""" + calls = [] + + def fake_get(url, params=None): + index = min(len(calls), len(traces) - 1) + calls.append(url) + return FakeResponse(traces[index], 200) + + monkeypatch.setattr(ars, "_get", fake_get) + return calls + + +def test_polling_stops_after_limit(monkeypatch, no_sleep): + _poll_sequence(monkeypatch, [trace("Running")]) + + with pytest.raises(TimeoutError): + ars.check_ars_results("pk", max_retries=2) + + +def test_wait_returns_once_parent_and_merge_child_are_done(monkeypatch, no_sleep): + calls = _poll_sequence( + monkeypatch, [trace("Running"), trace("Running"), done_trace()] + ) + + status = ars.wait_for_ARS(PARENT, poll_interval=1, max_retries=60) + + assert status.merged_version == MERGED + assert len(calls) == 3 + assert no_sleep == [1, 1] + + +def test_wait_keeps_polling_when_parent_done_but_merge_child_still_running(monkeypatch, no_sleep): + # NCATSTranslator/Relay#621: the parent flips to Done before the merged + # message is saved. The merge child must also be Done before we return. + calls = _poll_sequence( + monkeypatch, + [done_trace(merge_status="Running"), done_trace(merge_status="Running"), done_trace()], + ) + + status = ars.wait_for_ARS(PARENT, poll_interval=1, max_retries=60) + + assert status.merged_ready + assert len(calls) == 3 + + +def test_wait_returns_after_grace_when_no_merged_message_appears(monkeypatch, no_sleep): + clock = iter(range(0, 1000, 10)) + monkeypatch.setattr(ars.time, "monotonic", lambda: next(clock)) + calls = _poll_sequence( + monkeypatch, [trace("Done", None, [child("ara-x", "Done")])] + ) + + status = ars.wait_for_ARS( + PARENT, poll_interval=1, max_retries=600, merge_grace=25 + ) + + assert status.status == "Done" + assert status.merged_version is None + assert len(calls) >= 2 # waited through the grace window, then returned + + +def test_wait_raises_on_error_status(monkeypatch, no_sleep): + _poll_sequence(monkeypatch, [trace("Error", None, [child("ara-x", "Error", code=500)])]) + + with pytest.raises(ars.ARSError, match="status Error"): + ars.wait_for_ARS(PARENT, poll_interval=1, max_retries=60) + + +# --------------------------------------------------------------------------- # +# result collection and summarization +# --------------------------------------------------------------------------- # +def test_results_come_from_merged_message_or_fall_back_to_parent(monkeypatch): + fetched = [] + + def fake_get(url, params=None): + fetched.append((url, params)) + if params: + return FakeResponse(done_trace(), 200) + return FakeResponse(envelope(merged_message()), 200) + + monkeypatch.setattr(ars, "_get", fake_get) + + result = ars.get_ARS_result(PARENT) + + assert result.merged_pk == MERGED + assert fetched[-1][0].endswith(f"messages/{MERGED}") + assert len(result.results) == 2 + assert isinstance(result, FinderResult) + assert result.to_dict() is result.raw + assert result.summarize(1)[0]["essence"] == "drug" + + status = ars.parse_trace(PARENT, trace("Done", None, [])) + fallback = ars.get_ARS_result(status) + assert fallback.merged_pk is None + assert fetched[-1][0].endswith(f"messages/{PARENT}") + + +def test_pending_result_raises(monkeypatch): + monkeypatch.setattr( + ars, "get_ARS_status", lambda pk, url=None: ars.parse_trace(PARENT, trace("Running")) + ) + with pytest.raises(ars.ARSPendingError): + ars.get_ARS_result(PARENT) + + +def test_empty_shell_returned_without_merged_message(monkeypatch): + status = ars.parse_trace(PARENT, trace("Done", None, [child("ara-x", "Done")])) + monkeypatch.setattr(ars, "submit_ARS", lambda query, *, url=None: PARENT) + monkeypatch.setattr(ars, "wait_for_ARS", lambda pk, **kwargs: status) + monkeypatch.setattr(ars, "fetch_ars_results", lambda pk, *, url=None: {"fields": {"data": {}}}) + + result = ARS_neighborhood_finder(json_file={"message": {"query_graph": {}}}) + + assert isinstance(result, ars.ARSResult) + assert result.results == [] + assert result.knowledge_graph == {"nodes": {}, "edges": {}} + assert result.raw["query_graph"] == {} + + +def test_summarize_results_flattens_bindings_predicates_and_sources(): + rows = ars.summarize_results(merged_message(), top_n=1) + + assert len(rows) == 1 + row = rows[0] + assert row["rank"] == 1 and row["score"] == 0.9 + assert row["essence"] == "drug" + assert row["nodes"]["n1"][0]["name"] == "drug" + assert row["predicates"] == ["biolink:treats"] + assert row["primary_sources"] == ["infores:a"] + assert row["aras"] == ["infores:arax"] + assert row["edge_count"] == 1 + + assert len(ars.summarize_results(merged_message(), top_n=None)) == 2 + assert ars.summarize_results(None) == [] From 6c529df9e5d9edadfd952fed064a82e0281ca005 Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 15:22:29 -0700 Subject: [PATCH 06/19] test: gate live ARS and fixture-drift checks behind TCT_LIVE_ARS --- tests/test_ars_live.py | 131 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 131 insertions(+) create mode 100644 tests/test_ars_live.py diff --git a/tests/test_ars_live.py b/tests/test_ars_live.py new file mode 100644 index 0000000..a9072ae --- /dev/null +++ b/tests/test_ars_live.py @@ -0,0 +1,131 @@ +"""Live ARS tests, enabled only with ``TCT_LIVE_ARS=1``. + +Includes the fixture-drift test: it captures a real ARS trace and merged +envelope, then asserts the offline fixtures in ``tests/ars_fixtures.py`` +remain a structural subset of the live responses. Skipped entirely without +the environment variable so the default suite stays offline. +""" + +from __future__ import annotations + +import os + +import pytest + +from TCT import ARS_neighborhood_finder, ARS_pathfinder +from TCT import Query_ARS as ars_module +from TCT.TCT import FinderResult + +ars = ars_module +from tests.ars_fixtures import ( # noqa: E402 + child, + done_trace, + envelope, + merged_message, + trace, +) + +pytestmark = pytest.mark.skipif( + os.getenv("TCT_LIVE_ARS") != "1", + reason="live ARS checks run only with TCT_LIVE_ARS=1", +) + + +def test_live_ARS_neighborhood_finder_returns_finder_shaped_result(): + result = ARS_neighborhood_finder( + "asthma", + ["ChemicalEntity"], + predicates=["biolink:treated_by"], + ) + + assert isinstance(result, ars.ARSResult) + assert isinstance(result, FinderResult) + assert result.status is not None and result.status.status == "Done" + assert result.results, "expected at least one merged result" + rows = result.summarize(3) + assert rows and rows[0]["essence"] + + +def test_live_ARS_pathfinder_returns_paths(): + result = ARS_pathfinder("asthma", "albuterol", intermediate_categories=["Gene"]) + + assert isinstance(result, ars.ARSResult) + assert result.status is not None and result.status.status == "Done" + assert result.results, "expected at least one merged path result" + + +# --------------------------------------------------------------------------- # +# fixture drift +# --------------------------------------------------------------------------- # +def _subset_structure(offline, live, path, problems): + """Assert the offline fixture's structure is contained in the live value.""" + if isinstance(offline, dict): + if not isinstance(live, dict): + problems.append(f"{path}: expected object, got {type(live).__name__}") + return + for key, value in offline.items(): + if key not in live: + problems.append(f"{path}.{key}: missing in live response") + else: + _subset_structure(value, live[key], f"{path}.{key}", problems) + elif isinstance(offline, list): + if not isinstance(live, list): + problems.append(f"{path}: expected list, got {type(live).__name__}") + return + if not offline: + return + # Lists are compared by their first element: fixtures carry one + # representative entry. + if not live: + problems.append(f"{path}: empty live list") + else: + _subset_structure(offline[0], live[0], f"{path}[0]", problems) + elif offline is not None and not isinstance(live, type(offline)): + problems.append( + f"{path}: type {type(live).__name__} does not match fixture type " + f"{type(offline).__name__}" + ) + + +def test_offline_fixtures_match_live_ars_shape(): + """The offline fixtures must not drift from the real ARS response shape.""" + import requests + + session = requests.Session() + root = ars._ars_root(None) + + pk = session.post(f"{root}/submit", json=ars.format_query_json_forARS_neighborhood( + subject_ids=["MONDO:0004979"], + object_categories=["biolink:ChemicalEntity"], + predicates=["biolink:treated_by"], + ), timeout=ars.HTTP_TIMEOUT).json()["pk"] + + merged_pk = ars.check_ars_results(pk) + live_trace = session.get( + f"{root}/messages/{pk}", params={"trace": "y"}, timeout=ars.HTTP_TIMEOUT + ).json() + live_envelope = session.get( + f"{root}/messages/{merged_pk or pk}", timeout=ars.HTTP_TIMEOUT + ).json() + + problems: list[str] = [] + _subset_structure(trace("Done"), live_trace, "trace", problems) + _subset_structure( + done_trace(merged=merged_pk or "x"), + live_trace, + "done_trace", + problems, + ) + if not any("merged_version" in problem for problem in problems): + _subset_structure( + done_trace()["children"][0], + live_trace["children"][0] if live_trace.get("children") else child("x", "Done"), + "children[0]", + problems, + ) + _subset_structure(envelope(merged_message()), live_envelope, "envelope", problems) + + assert not problems, ( + "offline ARS fixtures drifted from the live response shape:\n" + + "\n".join(sorted(set(problems))) + ) From ba853a993ad1674c8aa4b5a476891550b89be708 Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 15:28:39 -0700 Subject: [PATCH 07/19] docs: document TCT.Query_ARS Re-executed Test_ARS.ipynb end to end (neighborhood, raw submit/poll flow, and the new pathfinder cells) against the CI ARS. Co-authored-by: Guangrong Qin <2994757+gloriachin@users.noreply.github.com> --- README.md | 1 + docs/source/ars.rst | 5 + docs/source/index.rst | 1 + notebooks/Test_ARS.ipynb | 313 +++++++++++++++++++++++++++++++++++++++ tests/ars_fixtures.py | 17 ++- tests/test_ars_live.py | 28 ++-- tests/test_query_ars.py | 11 +- 7 files changed, 356 insertions(+), 20 deletions(-) create mode 100644 docs/source/ars.rst create mode 100644 notebooks/Test_ARS.ipynb diff --git a/README.md b/README.md index 5114cc2..e282816 100644 --- a/README.md +++ b/README.md @@ -248,6 +248,7 @@ Python names in snake_case. | MetaKG extension | `add-custom-api-to-metakg`, `add-plover-apis-to-metakg` | | TRAPI query preparation and execution | `optimize-query-for-api`, `query-knowledge-provider`, `parallel-query-apis` | | Graph finding | `neighborhood-finder`, `path-finder` | +| ARS querying | `ars-neighborhood-finder`, `ars-pathfinder`, `submit-ars-query`, `get-ars-status`, `get-ars-results` | | Legacy compatibility | `trapi-query-endpoint` | `trapi-query-endpoint` preserves the existing public tool contract but is a diff --git a/docs/source/ars.rst b/docs/source/ars.rst new file mode 100644 index 0000000..09c1cf5 --- /dev/null +++ b/docs/source/ars.rst @@ -0,0 +1,5 @@ +TCT.Query_ARS +============= + +.. automodule:: TCT.Query_ARS + :members: diff --git a/docs/source/index.rst b/docs/source/index.rst index 62ee746..447920b 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -21,6 +21,7 @@ Welcome to Translator Component Toolkit's documentation! TCT neighborhood pathfinder + ars network_annotator node_annotator diff --git a/notebooks/Test_ARS.ipynb b/notebooks/Test_ARS.ipynb new file mode 100644 index 0000000..8756d01 --- /dev/null +++ b/notebooks/Test_ARS.ipynb @@ -0,0 +1,313 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "dfb3107c", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-29T22:24:07.026188Z", + "iopub.status.busy": "2026-09-29T22:24:07.026078Z", + "iopub.status.idle": "2026-09-29T22:24:07.270163Z", + "shell.execute_reply": "2026-09-29T22:24:07.269597Z" + } + }, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import sys\n", + "\n", + "project_root = Path.cwd().parent.resolve()\n", + "if str(project_root) not in sys.path:\n", + " sys.path.insert(0, str(project_root))\n", + "\n", + "from TCT import (\n", + " TCT,\n", + " TCT_neighborhood_finder,\n", + " Query_ARS,\n", + " name_resolver,\n", + " node_normalizer,\n", + " translator_kpinfo,\n", + " translator_metakg,\n", + " translator_query,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "df0b1fd3", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-29T22:24:07.271770Z", + "iopub.status.busy": "2026-09-29T22:24:07.271523Z", + "iopub.status.idle": "2026-09-29T22:24:07.273711Z", + "shell.execute_reply": "2026-09-29T22:24:07.273340Z" + } + }, + "outputs": [], + "source": [ + "from TCT.Query_ARS import ARS_neighborhood_finder, Query_ARS,check_ars_results\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "937d3537", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-29T22:24:07.274967Z", + "iopub.status.busy": "2026-09-29T22:24:07.274852Z", + "iopub.status.idle": "2026-09-29T22:25:18.808194Z", + "shell.execute_reply": "2026-09-29T22:25:18.807777Z" + } + }, + "outputs": [], + "source": [ + "ARS_results = ARS_neighborhood_finder(\n", + " node=['NCBIGene:2322'] , # replace with actual NodeInput instances\n", + " neighbor_categories=['biolink:Drug'], \n", + " predicates=['biolink:related_to'] # optional edge predicates\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "d8842a94", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-29T22:25:18.809325Z", + "iopub.status.busy": "2026-09-29T22:25:18.809211Z", + "iopub.status.idle": "2026-09-29T22:25:18.810959Z", + "shell.execute_reply": "2026-09-29T22:25:18.810748Z" + } + }, + "outputs": [], + "source": [ + "query = {\n", + " \"message\": {\n", + " \"query_graph\": {\n", + " \"edges\": {\n", + " \"e00\": {\n", + " \"subject\": \"n00\",\n", + " \"object\": \"n01\",\n", + " \"predicates\": [\"biolink:related_to\"]\n", + " }\n", + " },\n", + " \"nodes\": {\n", + " \"n00\": {\n", + " \"ids\": [\n", + " \"MONDO:0005148\"\n", + " ]\n", + " },\n", + " \"n01\": {\n", + " \"categories\": [\n", + " \"biolink:Drug\"\n", + " ]\n", + " }\n", + " }\n", + " }\n", + " },\n", + " \"submitter\": \"TCT\"\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "e9cd47d4", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-29T22:25:18.811756Z", + "iopub.status.busy": "2026-09-29T22:25:18.811624Z", + "iopub.status.idle": "2026-09-29T22:25:18.957815Z", + "shell.execute_reply": "2026-09-29T22:25:18.957489Z" + } + }, + "outputs": [], + "source": [ + "import json\n", + "import requests\n", + "TEST_URL = \"https://ars.ci.transltr.io/ars/api\"\n", + "pk = Query_ARS(url=TEST_URL, json_file=query)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "f0b14dab", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-29T22:25:18.959084Z", + "iopub.status.busy": "2026-09-29T22:25:18.958991Z", + "iopub.status.idle": "2026-09-29T22:26:19.452978Z", + "shell.execute_reply": "2026-09-29T22:26:19.452274Z" + } + }, + "outputs": [], + "source": [ + "BASE_URL=\"https://ars.ci.transltr.io/ars/api\"\n", + "merged_version_pk=check_ars_results(pk)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "15999ebf", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-29T22:26:19.455039Z", + "iopub.status.busy": "2026-09-29T22:26:19.454875Z", + "iopub.status.idle": "2026-09-29T22:26:33.995346Z", + "shell.execute_reply": "2026-09-29T22:26:33.994721Z" + } + }, + "outputs": [], + "source": [ + "BASE_URL=\"https://ars.ci.transltr.io/ars/api\"\n", + "\n", + "def fetch_ars_results(merged_pk: str):\n", + " url = BASE_URL+\"/messages/\"+merged_pk\n", + " results = requests.get(url).json()\n", + " return results\n", + "\n", + "results = fetch_ars_results(merged_version_pk)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "88a81a9d", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-29T22:26:33.997367Z", + "iopub.status.busy": "2026-09-29T22:26:33.997236Z", + "iopub.status.idle": "2026-09-29T22:28:09.743462Z", + "shell.execute_reply": "2026-09-29T22:28:09.743033Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "(1, '36c801ac-81a3-4cc6-9994-c5c0c30984d2')" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from TCT import ARS_pathfinder\n", + "\n", + "# One-hop pathfinder over the ARS: every ARA is asked for paths between the\n", + "# two concepts, then the ARS merges and ranks their answers. Names resolve to\n", + "# CURIEs; at most one intermediate category is applied.\n", + "paths_result = ARS_pathfinder(\n", + " \"asthma\",\n", + " \"albuterol\",\n", + " intermediate_categories=[\"Gene\"],\n", + ")\n", + "len(paths_result.results), paths_result.merged_pk" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "1023f87f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-29T22:28:09.745153Z", + "iopub.status.busy": "2026-09-29T22:28:09.744908Z", + "iopub.status.idle": "2026-09-29T22:28:09.747147Z", + "shell.execute_reply": "2026-09-29T22:28:09.746981Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "[{'rank': 1,\n", + " 'score': 100.0,\n", + " 'essence': 'result',\n", + " 'essence_category': None,\n", + " 'nodes': {'on': [{'id': 'CHEBI:2549',\n", + " 'name': 'Albuterol',\n", + " 'categories': ['biolink:Drug']}],\n", + " 'sn': [{'id': 'MONDO:0004979',\n", + " 'name': 'asthma',\n", + " 'categories': ['biolink:Disease']}]},\n", + " 'predicates': [],\n", + " 'primary_sources': [],\n", + " 'aras': ['infores:arax', 'infores:shepherd-aragorn'],\n", + " 'edge_count': 0}]" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "paths_result.summarize(5)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "3bc26d43", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-29T22:28:09.748469Z", + "iopub.status.busy": "2026-09-29T22:28:09.748402Z", + "iopub.status.idle": "2026-09-29T22:28:09.750011Z", + "shell.execute_reply": "2026-09-29T22:28:09.749879Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "['CHEBI:132742',\n", + " 'CHEBI:134556',\n", + " 'CHEBI:134557',\n", + " 'CHEBI:134558',\n", + " 'CHEBI:134559']" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# The finders return an ARSResult, which extends FinderResult: the merged\n", + "# TRAPI message is available as .raw for the usual downstream consumers.\n", + "sorted(paths_result.knowledge_graph[\"nodes\"])[:5]" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "my-rdkit-env", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/ars_fixtures.py b/tests/ars_fixtures.py index e641ab6..de68a11 100644 --- a/tests/ars_fixtures.py +++ b/tests/ars_fixtures.py @@ -36,12 +36,16 @@ def child(agent, status, pk="child", result_count=None, code=200): def trace(status, merged=None, children=()): - """A ``?trace=y`` parent response.""" + """A ``?trace=y`` parent response. + + ``merged_versions_list`` mirrors the live ARS, which serializes the field + as the Python repr of a list of ``[pk, agent]`` pairs rather than JSON. + """ return { "status": status, "code": 200 if status == "Done" else 202, "merged_version": merged, - "merged_versions_list": [], + "merged_versions_list": f"[['{merged}', 'ars']]" if merged else "[]", "children": list(children), } @@ -87,19 +91,16 @@ def merged_message(): }, "results": [ { - "rank": 1, "normalized_score": 0.9, - "essence": "drug", - "essence_category": "biolink:SmallMolecule", - "node_bindings": {"n0": [{"id": "MONDO:1"}], "n1": [{"id": "CHEBI:1"}]}, + "node_bindings": {"n00": [{"id": "MONDO:1"}], "n01": [{"id": "CHEBI:1"}]}, "analyses": [ { "resource_id": "infores:arax", - "edge_bindings": {"e0": [{"id": "e1"}]}, + "edge_bindings": {"e00": [{"id": "e1"}]}, } ], }, - {"rank": 2, "score": 0.1, "node_bindings": {}, "analyses": []}, + {"node_bindings": {}, "analyses": []}, ], "auxiliary_graphs": {}, } diff --git a/tests/test_ars_live.py b/tests/test_ars_live.py index a9072ae..a67afb9 100644 --- a/tests/test_ars_live.py +++ b/tests/test_ars_live.py @@ -32,18 +32,17 @@ def test_live_ARS_neighborhood_finder_returns_finder_shaped_result(): - result = ARS_neighborhood_finder( - "asthma", - ["ChemicalEntity"], - predicates=["biolink:treated_by"], - ) + # No predicate filter: the ARS accepts the query but some directional + # predicates (e.g. biolink:treated_by) come back with an empty merge. + result = ARS_neighborhood_finder("asthma", ["ChemicalEntity"]) assert isinstance(result, ars.ARSResult) assert isinstance(result, FinderResult) assert result.status is not None and result.status.status == "Done" assert result.results, "expected at least one merged result" rows = result.summarize(3) - assert rows and rows[0]["essence"] + # essence is optional in ARS merges (aragorn/bte leave it unset). + assert rows and rows[0]["rank"] == 1 and rows[0]["nodes"] def test_live_ARS_pathfinder_returns_paths(): @@ -64,7 +63,14 @@ def _subset_structure(offline, live, path, problems): problems.append(f"{path}: expected object, got {type(live).__name__}") return for key, value in offline.items(): - if key not in live: + if key == "*": + # Wildcard: compare the value shape against any one live entry. + if not live: + problems.append(f"{path}: empty live object") + else: + first = next(iter(live)) + _subset_structure(value, live[first], f"{path}.{first}", problems) + elif key not in live: problems.append(f"{path}.{key}: missing in live response") else: _subset_structure(value, live[key], f"{path}.{key}", problems) @@ -97,7 +103,6 @@ def test_offline_fixtures_match_live_ars_shape(): pk = session.post(f"{root}/submit", json=ars.format_query_json_forARS_neighborhood( subject_ids=["MONDO:0004979"], object_categories=["biolink:ChemicalEntity"], - predicates=["biolink:treated_by"], ), timeout=ars.HTTP_TIMEOUT).json()["pk"] merged_pk = ars.check_ars_results(pk) @@ -123,7 +128,12 @@ def test_offline_fixtures_match_live_ars_shape(): "children[0]", problems, ) - _subset_structure(envelope(merged_message()), live_envelope, "envelope", problems) + fixture_envelope = envelope(merged_message()) + kg = fixture_envelope["fields"]["data"]["message"]["knowledge_graph"] + # Node and edge ids vary per query; compare their entry shapes instead. + kg["nodes"] = {"*": next(iter(merged_message()["knowledge_graph"]["nodes"].values()))} + kg["edges"] = {"*": next(iter(merged_message()["knowledge_graph"]["edges"].values()))} + _subset_structure(fixture_envelope, live_envelope, "envelope", problems) assert not problems, ( "offline ARS fixtures drifted from the live response shape:\n" diff --git a/tests/test_query_ars.py b/tests/test_query_ars.py index b862a70..7071fe1 100644 --- a/tests/test_query_ars.py +++ b/tests/test_query_ars.py @@ -336,7 +336,7 @@ def fake_get(url, params=None): assert len(result.results) == 2 assert isinstance(result, FinderResult) assert result.to_dict() is result.raw - assert result.summarize(1)[0]["essence"] == "drug" + assert result.summarize(1)[0]["nodes"]["n01"][0]["name"] == "drug" status = ars.parse_trace(PARENT, trace("Done", None, [])) fallback = ars.get_ARS_result(status) @@ -367,13 +367,18 @@ def test_empty_shell_returned_without_merged_message(monkeypatch): def test_summarize_results_flattens_bindings_predicates_and_sources(): - rows = ars.summarize_results(merged_message(), top_n=1) + # essence/rank are optional per ARA; the summarizer must read them when set. + enriched = merged_message() + enriched["results"][0].update( + {"rank": 1, "essence": "drug", "essence_category": "biolink:SmallMolecule"} + ) + rows = ars.summarize_results(enriched, top_n=1) assert len(rows) == 1 row = rows[0] assert row["rank"] == 1 and row["score"] == 0.9 assert row["essence"] == "drug" - assert row["nodes"]["n1"][0]["name"] == "drug" + assert row["nodes"]["n01"][0]["name"] == "drug" assert row["predicates"] == ["biolink:treats"] assert row["primary_sources"] == ["infores:a"] assert row["aras"] == ["infores:arax"] From e37cf94c029d43e53da13755501a338f4c607497 Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 15:52:58 -0700 Subject: [PATCH 08/19] fix: restore the DEFAULT_ENVIRONMENT constant dropped by the nodenorm URL change Guangrong's nodenorm fix was written against an older base that predated the constant; keep his endpoint change but preserve the constant. Co-authored-by: Guangrong Qin <2994757+gloriachin@users.noreply.github.com> --- TCT/config.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/TCT/config.py b/TCT/config.py index 5c3cc67..f1ff96e 100644 --- a/TCT/config.py +++ b/TCT/config.py @@ -8,6 +8,7 @@ Environment = Literal["prod", "ci", "test"] +DEFAULT_ENVIRONMENT: Environment = "ci" @dataclass(frozen=True) @@ -73,7 +74,7 @@ def resolve(self, environment: Environment) -> str: class RuntimeConfig: """Runtime environment and explicit service URL replacements.""" - environment: Environment = "ci" + environment: Environment = DEFAULT_ENVIRONMENT overrides: Mapping[str, str] = field(default_factory=dict) def __post_init__(self) -> None: @@ -102,7 +103,10 @@ def load_config( overrides: Mapping[str, str] | None = None, ) -> RuntimeConfig: """Build configuration from explicit values and ``TCT_ENVIRONMENT``.""" - selected_environment = environment or os.getenv("TCT_ENVIRONMENT", "ci") + selected_environment = environment or os.getenv( + "TCT_ENVIRONMENT", + DEFAULT_ENVIRONMENT, + ) return RuntimeConfig( environment=selected_environment, # type: ignore[arg-type] overrides=overrides or {}, From f3fbcced769bef5613e3ea5b26bf8279ceef67fb Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 15:53:17 -0700 Subject: [PATCH 09/19] refactor: drop a duplicated ARS root helper --- TCT/Query_ARS.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/TCT/Query_ARS.py b/TCT/Query_ARS.py index df71fbb..b7cd642 100644 --- a/TCT/Query_ARS.py +++ b/TCT/Query_ARS.py @@ -209,11 +209,6 @@ def _get(url: str, params: Optional[dict[str, str]] = None) -> requests.Response return _session().get(url, params=params, timeout=HTTP_TIMEOUT) -def _api_root(url: Optional[str]) -> str: - root = (url or service_url("ars")).rstrip("/") - return root - - def _message_from_envelope(envelope: dict[str, Any]) -> Optional[dict[str, Any]]: """Extract the merged TRAPI message from a full ARS message envelope.""" data = (envelope.get("fields") or {}).get("data") or {} From 14968b0c82e98c6125b9f35da7604aff58c9cb51 Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 15:55:56 -0700 Subject: [PATCH 10/19] feat: add ARS neighborhood finder and pathfinder tools Ported from Guangrong Qin's feat/ars-query-support tool layer, reshaped onto the ARSResult API: ranked summary rows for built queries, the full merged message for uploaded json_file queries. Co-authored-by: Guangrong Qin <2994757+gloriachin@users.noreply.github.com> --- TCT/interfaces/tools.py | 94 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) diff --git a/TCT/interfaces/tools.py b/TCT/interfaces/tools.py index eb92800..2de2281 100644 --- a/TCT/interfaces/tools.py +++ b/TCT/interfaces/tools.py @@ -16,6 +16,15 @@ from ..name_resolver import batch_lookup, lookup, synonyms from ..node_normalizer import get_normalized_nodes +from ..Query_ARS import ( + ARS_neighborhood_finder as _ars_neighborhood_finder, + ARS_pathfinder as _ars_pathfinder, + format_query_json_forARS_neighborhood as _format_query_json_forARS_neighborhood, + get_ARS_result as _get_ARS_result, + get_ARS_status as _get_ARS_status, + submit_ARS as _submit_ARS, +) +from ..TCT import _normalize_categories, _resolve_nodes from ..TCT import get_translator_resources as _get_translator_resources from ..TCT_neighborhood_finder import neighborhood_finder as tct_neighborhood_finder from ..TCT_pathfinder import query_TCT_pathfinder @@ -336,6 +345,89 @@ def path_finder( ) +def ARS_neighborhood_finder( + json_file: dict[str, Any] | None = None, + node: list[str] | None = None, + neighbor_categories: list[str] | None = None, +) -> Any: + """Find neighbors for concepts by querying every ARA through the ARS. + + Args: + json_file: Parsed TRAPI query from an uploaded file. Submitted + unchanged, so both edge and path query graphs work. When given, + ``node`` and ``neighbor_categories`` must be omitted. + node: One or more names or CURIEs; a one-hop query is built instead. + Names are resolved to CURIEs before submission. + neighbor_categories: Biolink categories wanted for neighbors, with or + without the "biolink:" prefix (for example ["Drug"]). + + Returns: + Resolved inputs plus ranked summary rows (rank, essence, predicates, + primary sources, ARAs). Blocks until the ARS finishes: expect ~15 s to + several minutes. Prefer submit_ars_query / get_ars_status / + get_ars_results when polling is better than waiting. + """ + result = _ars_neighborhood_finder( + node=node, + neighbor_categories=neighbor_categories, + json_file=json_file, + ) + if json_file is not None: + return { + "merged_pk": result.merged_pk, + "status": result.status, + "message": result.raw, + } + return { + "resolved_nodes": result.resolved_nodes, + "status": result.status, + "result_count": len(result.results), + "results": result.summarize(20), + } + + +def ARS_pathfinder( + start: str, + end: str, + intermediate_categories: list[str] | None = None, + json_file: dict[str, Any] | None = None, +) -> Any: + """Find paths between two concepts by querying every ARA through the ARS. + + Args: + start: Name or CURIE of the starting node. + end: Name or CURIE of the ending node. + intermediate_categories: Optional single Biolink category restricting + intermediate path nodes (for example ["Gene"]). + json_file: Parsed TRAPI query from an uploaded file, submitted + unchanged instead of building a paths query. + + Returns: + Resolved start and end nodes plus ranked summary rows. Blocks until + the ARS finishes: expect ~15 s to several minutes. Prefer + submit_ars_query / get_ars_status / get_ars_results when polling is + better than waiting. + """ + result = _ars_pathfinder( + start, + end, + intermediate_categories=intermediate_categories, + json_file=json_file, + ) + if json_file is not None: + return { + "merged_pk": result.merged_pk, + "status": result.status, + "message": result.raw, + } + return { + "resolved_nodes": result.resolved_nodes, + "status": result.status, + "result_count": len(result.results), + "results": result.summarize(20), + } + + TOOLS: tuple[Callable[..., Any], ...] = ( get_translator_resources, name_lookup, @@ -353,6 +445,8 @@ def path_finder( trapi_query_endpoint, neighborhood_finder, path_finder, + ARS_neighborhood_finder, + ARS_pathfinder, ) __all__ = [tool.__name__ for tool in TOOLS] + ["TOOLS"] From de51bfad0fff11bf565efb634a12bcee97af8abe Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 15:55:57 -0700 Subject: [PATCH 11/19] feat: add ARS submit, status, and results tools --- TCT/interfaces/tools.py | 98 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/TCT/interfaces/tools.py b/TCT/interfaces/tools.py index 2de2281..0da6fad 100644 --- a/TCT/interfaces/tools.py +++ b/TCT/interfaces/tools.py @@ -428,6 +428,101 @@ def ARS_pathfinder( } +def submit_ars_query( + node: list[str], + neighbor_categories: list[str], + predicates: list[str] | None = None, +) -> Any: + """Submit a one-hop biomedical query to the Translator ARS and return immediately. + + The ARS fans the query out to every registered ARA and merges their answers; + completion typically takes ~15 s to several minutes, so this does NOT wait + for results. Track progress by passing the returned pk to get_ars_status + (poll every ~15 s) until its status is Done, then fetch answers with + get_ars_results. + + Args: + node: One or more input nodes, as names ("asthma") or CURIEs + ("MONDO:0004979"). Names are resolved to CURIEs before submission. + neighbor_categories: Biolink categories wanted for neighbors, with or + without the "biolink:" prefix (for example ["Drug"]). + predicates: Optional edge predicates to require, for example + ["biolink:treated_by"]. Omit for any predicate; the ARS returns an + empty merge for some directional predicates. Defaults to + biolink:related_to, which spans directions. + + Returns: + {"pk", "resolved_nodes", "status"}. Keep the pk for the follow-up calls. + + Fails only on unresolvable inputs or submission rejection; later failure + shows up as status Error in get_ars_status. + """ + resolved = _resolve_nodes(node) + pk = _submit_ARS(_format_query_json_forARS_neighborhood( + subject_ids=[resolved_node.curie for resolved_node in resolved], + object_categories=_normalize_categories(neighbor_categories), + predicates=predicates, + )) + return { + "pk": pk, + "resolved_nodes": { + f"node_{index}": resolved_node + for index, resolved_node in enumerate(resolved) + }, + "status": _get_ARS_status(pk), + } + + +def get_ars_status(pk: str) -> Any: + """Check progress of an ARS query submitted with submit_ars_query, without blocking. + + Args: + pk: Parent message pk returned by submit_ars_query. + + Returns: + Parent status ("Running", "Done", or "Error"), the merged message pk + once available, and one entry per ARA with its own status and result + count. Poll until status is Done, then call get_ars_results once. + "Done" with zero results means the query matched nothing and will not + change. + """ + return _get_ARS_status(pk) + + +def get_ars_results(pk: str, top_n: int = 20) -> Any: + """Fetch the merged answer of a finished ARS query as ranked summary rows. + + Args: + pk: Parent message pk returned by submit_ars_query. + top_n: Number of ranked rows to return (default 20). Keep the default + unless explicitly asked for more; pass 0 only when exporting the + full merged TRAPI message, which can be tens of megabytes. + + Returns: + {"pk", "merged_pk", "status", "ready", "result_count", "results"}. Rows + give rank, score, essence (answer node with name and categories), + predicates, primary knowledge sources, and contributing ARAs. When + ready is false the query is still running: poll get_ars_status instead + of retrying this. + """ + status = _get_ARS_status(pk) + if not status.is_terminal: + return {"pk": pk, "merged_pk": status.merged_version, "status": status, "ready": False} + result = _get_ARS_result(status) + payload = { + "pk": result.pk, + "merged_pk": result.merged_pk, + "status": result.status, + "ready": True, + "result_count": len(result.results), + } + if top_n <= 0: + payload["message"] = result.raw + else: + payload["results"] = result.summarize(top_n) + return payload + + TOOLS: tuple[Callable[..., Any], ...] = ( get_translator_resources, name_lookup, @@ -447,6 +542,9 @@ def ARS_pathfinder( path_finder, ARS_neighborhood_finder, ARS_pathfinder, + submit_ars_query, + get_ars_status, + get_ars_results, ) __all__ = [tool.__name__ for tool in TOOLS] + ["TOOLS"] From 99ae51b20f520c26894dd54f80023cc90b817e2b Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 15:57:38 -0700 Subject: [PATCH 12/19] feat: run query_ars as a SEP-1686 background task Async wrapper over the ARS finder flow with per-poll progress reporting. Task-aware clients get a task ID immediately; task-unaware clients get a plain blocking call per SEP-1686 graceful degradation. --- TCT/interfaces/mcp.py | 90 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 89 insertions(+), 1 deletion(-) diff --git a/TCT/interfaces/mcp.py b/TCT/interfaces/mcp.py index 5af3982..1e41996 100644 --- a/TCT/interfaces/mcp.py +++ b/TCT/interfaces/mcp.py @@ -109,6 +109,94 @@ def invoke(*args: Any, **kwargs: Any) -> Any: for _tool in shared_tools.TOOLS: globals()[_tool.__name__] = _register_tool(_tool) +def _register_ars_task_tool() -> None: + """Expose the ARS finder flow as a SEP-1686 background task. + + A blocking ARS query runs 15 s to 15 min. Task-aware clients get a task ID + immediately and poll progress; task-unaware clients degrade to a plain + blocking call (SEP-1686 graceful degradation). + """ + import asyncio + + from fastmcp.dependencies import Progress + from fastmcp.server.tasks import TaskConfig + + from .. import Query_ARS as ars + + async def query_ars( + node: list[str], + neighbor_categories: list[str], + predicates: list[str] | None = None, + top_n: int = 20, + progress: Progress = Progress(), + ) -> dict[str, Any]: + """Query the Translator ARS end to end and return ranked answer rows. + + Submits a one-hop query, waits for every ARA and the ARS merge agent, + and returns the top answers. Long-running: expect ~15 s to several + minutes. Task-aware MCP clients receive a task ID immediately and can + poll progress instead of waiting; task-unaware clients block until done + and should raise their call timeout accordingly. Prefer the faster + submit_ars_query / get_ars_status / get_ars_results workflow otherwise. + + Args: + node: One or more input nodes, as names ("asthma") or CURIEs + ("MONDO:0004979"). Names are resolved to CURIEs first. + neighbor_categories: Biolink categories wanted for neighbors, with + or without the "biolink:" prefix (for example ["Drug"]). + predicates: Optional edge predicates to require. Omit for any + predicate; defaults to biolink:related_to. + top_n: Number of ranked rows to return (default 20); pass 0 only to + export the full merged TRAPI message, which can be tens of MB. + + Returns: + The pk, merged pk, final status, resolved input nodes, and either + ranked summary rows or the full merged TRAPI message. + """ + resolved = shared_tools._resolve_nodes(node) + query = ars.format_query_json_forARS_neighborhood( + subject_ids=[resolved_node.curie for resolved_node in resolved], + object_categories=shared_tools._normalize_categories(neighbor_categories), + predicates=predicates, + ) + pk = ars.submit_ARS(query) + + await progress.set_total(2) + await progress.set_message("submitted; waiting for ARA answers") + status = await asyncio.to_thread( + ars.wait_for_ARS, pk, poll_interval=15.0, max_retries=60 + ) + await progress.set_message(f"ARS status: {status.status}") + await progress.increment() + result = await asyncio.to_thread(ars.get_ARS_result, status) + await progress.increment() + + payload: dict[str, Any] = { + "pk": result.pk, + "merged_pk": result.merged_pk, + "status": result.status, + "resolved_nodes": { + f"node_{index}": resolved_node + for index, resolved_node in enumerate(resolved) + }, + "result_count": len(result.results), + } + if top_n <= 0: + payload["message"] = result.raw + else: + payload["results"] = result.summarize(top_n) + return payload + + mcp.tool( + name="query_ars", + task=TaskConfig(mode="optional"), + annotations={"readOnlyHint": True, "openWorldHint": True}, + )(query_ars) + globals()["query_ars"] = query_ars + + +_register_ars_task_tool() + def main() -> None: """Entry point for the installed ``tct-server`` command.""" @@ -119,4 +207,4 @@ def main() -> None: flush_observability() -__all__ = ["main", "mcp", *[tool.__name__ for tool in shared_tools.TOOLS]] +__all__ = ["main", "mcp", "query_ars", *[tool.__name__ for tool in shared_tools.TOOLS]] From 5a0828598671dbf5d39c45b9869172490171a765 Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 16:11:26 -0700 Subject: [PATCH 13/19] feat: re-export ARS tools from TCT.server ARS_neighborhood_finder, ARS_pathfinder, submit_ars_query, get_ars_status, get_ars_results, and the query_ars background task join the server namespace. --- TCT/interfaces/mcp.py | 20 +++++--- TCT/server.py | 12 +++++ tests/test_mcp_contract.py | 101 ++++++++++++++++++++++++++++++++++++- tests/test_server.py | 2 + tests/test_tools.py | 7 ++- 5 files changed, 132 insertions(+), 10 deletions(-) diff --git a/TCT/interfaces/mcp.py b/TCT/interfaces/mcp.py index 1e41996..998bed3 100644 --- a/TCT/interfaces/mcp.py +++ b/TCT/interfaces/mcp.py @@ -13,8 +13,10 @@ from typing import Any from fastmcp import FastMCP +from fastmcp.dependencies import Progress from fastmcp.exceptions import ToolError from fastmcp.server.middleware import Middleware, MiddlewareContext +from fastmcp.server.tasks import TaskConfig from fastmcp.tools.tool import ToolResult from . import tools as shared_tools @@ -29,6 +31,12 @@ # services, so openWorldHint is always true. _READ_ONLY_TOOLS = frozenset( { + "ARS_neighborhood_finder", + "ARS_pathfinder", + "submit_ars_query", + "get_ars_status", + "get_ars_results", + "query_ars", "get_translator_resources", "name_lookup", "get_name_synonyms", @@ -118,9 +126,6 @@ def _register_ars_task_tool() -> None: """ import asyncio - from fastmcp.dependencies import Progress - from fastmcp.server.tasks import TaskConfig - from .. import Query_ARS as ars async def query_ars( @@ -129,7 +134,7 @@ async def query_ars( predicates: list[str] | None = None, top_n: int = 20, progress: Progress = Progress(), - ) -> dict[str, Any]: + ) -> Any: """Query the Translator ARS end to end and return ranked answer rows. Submits a one-hop query, waits for every ARA and the ARS merge agent, @@ -171,7 +176,7 @@ async def query_ars( result = await asyncio.to_thread(ars.get_ARS_result, status) await progress.increment() - payload: dict[str, Any] = { + payload = { "pk": result.pk, "merged_pk": result.merged_pk, "status": result.status, @@ -187,15 +192,14 @@ async def query_ars( payload["results"] = result.summarize(top_n) return payload - mcp.tool( + return mcp.tool( name="query_ars", task=TaskConfig(mode="optional"), annotations={"readOnlyHint": True, "openWorldHint": True}, )(query_ars) - globals()["query_ars"] = query_ars -_register_ars_task_tool() +query_ars = _register_ars_task_tool() def main() -> None: diff --git a/TCT/server.py b/TCT/server.py index 68df721..0db3e30 100644 --- a/TCT/server.py +++ b/TCT/server.py @@ -6,12 +6,16 @@ """ from .interfaces.mcp import ( + ARS_neighborhood_finder, + ARS_pathfinder, add_custom_api_to_metakg, add_plover_apis_to_metakg, batch_name_lookup, get_api_predicates, get_kp_info, get_metakg_data, + get_ars_results, + get_ars_status, get_name_synonyms, get_translator_resources, mcp, @@ -21,17 +25,23 @@ optimize_query_for_api, parallel_query_apis, path_finder, + query_ars, query_knowledge_provider, + submit_ars_query, trapi_query_endpoint, ) __all__ = [ + "ARS_neighborhood_finder", + "ARS_pathfinder", "add_custom_api_to_metakg", "add_plover_apis_to_metakg", "batch_name_lookup", "get_api_predicates", "get_kp_info", "get_metakg_data", + "get_ars_results", + "get_ars_status", "get_name_synonyms", "get_translator_resources", "mcp", @@ -41,6 +51,8 @@ "optimize_query_for_api", "parallel_query_apis", "path_finder", + "query_ars", "query_knowledge_provider", + "submit_ars_query", "trapi_query_endpoint", ] diff --git a/tests/test_mcp_contract.py b/tests/test_mcp_contract.py index 4665b22..e306127 100644 --- a/tests/test_mcp_contract.py +++ b/tests/test_mcp_contract.py @@ -60,6 +60,26 @@ "(start: 'str', end: 'str', " "intermediate_categories: 'list[str] | None' = None) -> 'Any'" ), + "ARS_neighborhood_finder": ( + "(json_file: 'dict[str, Any] | None' = None, " + "node: 'list[str] | None' = None, " + "neighbor_categories: 'list[str] | None' = None) -> 'Any'" + ), + "ARS_pathfinder": ( + "(start: 'str', end: 'str', " + "intermediate_categories: 'list[str] | None' = None, " + "json_file: 'dict[str, Any] | None' = None) -> 'Any'" + ), + "submit_ars_query": ( + "(node: 'list[str]', neighbor_categories: 'list[str]', " + "predicates: 'list[str] | None' = None) -> 'Any'" + ), + "get_ars_status": "(pk: 'str') -> 'Any'", + "get_ars_results": "(pk: 'str', top_n: 'int' = 20) -> 'Any'", + "query_ars": ( + "(node: 'list[str]', neighbor_categories: 'list[str]', " + "predicates: 'list[str] | None' = None, top_n: 'int' = 20) -> 'Any'" + ), } @@ -223,6 +243,75 @@ }, ["start", "end"], ), + "ARS_neighborhood_finder": ( + { + "json_file": { + "anyOf": [ + {"additionalProperties": True, "type": "object"}, + {"type": "null"}, + ], + "default": None, + }, + "node": { + "anyOf": [{"items": {"type": "string"}, "type": "array"}, {"type": "null"}], + "default": None, + }, + "neighbor_categories": { + "anyOf": [{"items": {"type": "string"}, "type": "array"}, {"type": "null"}], + "default": None, + }, + }, + [], + ), + "ARS_pathfinder": ( + { + "start": {"type": "string"}, + "end": {"type": "string"}, + "intermediate_categories": { + "anyOf": [ + {"items": {"type": "string"}, "type": "array"}, + {"type": "null"}, + ], + "default": None, + }, + "json_file": { + "anyOf": [ + {"additionalProperties": True, "type": "object"}, + {"type": "null"}, + ], + "default": None, + }, + }, + ["start", "end"], + ), + "submit_ars_query": ( + { + "node": {"items": {"type": "string"}, "type": "array"}, + "neighbor_categories": {"items": {"type": "string"}, "type": "array"}, + "predicates": { + "anyOf": [{"items": {"type": "string"}, "type": "array"}, {"type": "null"}], + "default": None, + }, + }, + ["node", "neighbor_categories"], + ), + "get_ars_status": ({"pk": {"type": "string"}}, ["pk"]), + "get_ars_results": ( + {"pk": {"type": "string"}, "top_n": {"default": 20, "type": "integer"}}, + ["pk"], + ), + "query_ars": ( + { + "node": {"items": {"type": "string"}, "type": "array"}, + "neighbor_categories": {"items": {"type": "string"}, "type": "array"}, + "predicates": { + "anyOf": [{"items": {"type": "string"}, "type": "array"}, {"type": "null"}], + "default": None, + }, + "top_n": {"default": 20, "type": "integer"}, + }, + ["node", "neighbor_categories"], + ), } @@ -243,7 +332,17 @@ def test_published_tool_names_signatures_and_docs_are_stable(): assert list(tools) == list(EXPECTED_SIGNATURES) for name, expected_signature in EXPECTED_SIGNATURES.items(): tool = tools[name] - assert str(inspect.signature(tool.fn)) == expected_signature + signature = inspect.signature(tool.fn) + # query_ars hides the MCP-injected Progress dependency. + if "progress" in signature.parameters: + signature = signature.replace( + parameters=[ + parameter + for parameter_name, parameter in signature.parameters.items() + if parameter_name != "progress" + ] + ) + assert str(signature) == expected_signature assert tool.description == inspect.getdoc(tool.fn) assert tool.description diff --git a/tests/test_server.py b/tests/test_server.py index d82a3aa..ec78401 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -190,6 +190,8 @@ async def collect(): "add_custom_api_to_metakg", "add_plover_apis_to_metakg", } + assert "query_ars" in tools # MCP-only background-task tool + assert tools["query_ars"].annotations.readOnlyHint is True def test_fastmcp_background_task_protocol_is_available(): """The pinned fastmcp registers async task tools with SEP-1686 metadata.""" diff --git a/tests/test_tools.py b/tests/test_tools.py index 62ce71f..bc6d894 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -25,6 +25,11 @@ "trapi_query_endpoint", "neighborhood_finder", "path_finder", + "ARS_neighborhood_finder", + "ARS_pathfinder", + "submit_ars_query", + "get_ars_status", + "get_ars_results", ] @@ -54,7 +59,7 @@ def find_spec(self, fullname, path=None, target=None): sys.meta_path.insert(0, RejectMcpImports()) from TCT.interfaces.cli import build_parser from TCT.interfaces.tools import TOOLS -assert len(TOOLS) == 16 +assert len(TOOLS) == 21 assert build_parser().prog == "tct" """ From 26cd220ebca85e6743373bd2070d99cadfd6a09c Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 17:21:55 -0700 Subject: [PATCH 14/19] refactor: move query_ars into its own module FastMCP resolves the Progress dependency from real annotations, and PEP 563 string annotations in the adapter module defeat that resolution, which broke task-unaware clients. The task tool now lives in a module without the future import; the adapter registers it after the shared tools. --- TCT/interfaces/ars_task.py | 93 ++++++++++++++++++++++++++++++++++++++ TCT/interfaces/mcp.py | 90 ++---------------------------------- 2 files changed, 97 insertions(+), 86 deletions(-) create mode 100644 TCT/interfaces/ars_task.py diff --git a/TCT/interfaces/ars_task.py b/TCT/interfaces/ars_task.py new file mode 100644 index 0000000..c8c3c15 --- /dev/null +++ b/TCT/interfaces/ars_task.py @@ -0,0 +1,93 @@ +"""MCP-only ARS finder flow exposed as a SEP-1686 background task. + +This module intentionally avoids ``from __future__ import annotations``: +FastMCP resolves ``Progress`` dependency parameters from real annotations, and +PEP 563 string annotations defeat that resolution. + +A blocking ARS query runs 15 s to 15 min. Task-aware clients get a task ID +immediately and poll progress; task-unaware clients degrade to a plain +blocking call (SEP-1686 graceful degradation). +""" + +from asyncio import to_thread as asyncio_to_thread +from typing import Any + +from fastmcp.dependencies import Progress +from fastmcp.server.tasks import TaskConfig + +from .. import Query_ARS as ars +from . import tools as shared_tools + + +async def query_ars( + node: list[str], + neighbor_categories: list[str], + predicates: list[str] | None = None, + top_n: int = 20, + progress: Progress = Progress(), +) -> Any: + """Query the Translator ARS end to end and return ranked answer rows. + + Submits a one-hop query, waits for every ARA and the ARS merge agent, + and returns the top answers. Long-running: expect ~15 s to several + minutes. Task-aware MCP clients receive a task ID immediately and can + poll progress instead of waiting; task-unaware clients block until done + and should raise their call timeout accordingly. Prefer the faster + submit_ars_query / get_ars_status / get_ars_results workflow otherwise. + + Args: + node: One or more input nodes, as names ("asthma") or CURIEs + ("MONDO:0004979"). Names are resolved to CURIEs first. + neighbor_categories: Biolink categories wanted for neighbors, with + or without the "biolink:" prefix (for example ["Drug"]). + predicates: Optional edge predicates to require. Omit for any + predicate; defaults to biolink:related_to. + top_n: Number of ranked rows to return (default 20); pass 0 only to + export the full merged TRAPI message, which can be tens of MB. + + Returns: + The pk, merged pk, final status, resolved input nodes, and either + ranked summary rows or the full merged TRAPI message. + """ + resolved = shared_tools._resolve_nodes(node) + query = ars.format_query_json_forARS_neighborhood( + subject_ids=[resolved_node.curie for resolved_node in resolved], + object_categories=shared_tools._normalize_categories(neighbor_categories), + predicates=predicates, + ) + pk = ars.submit_ARS(query) + + await progress.set_total(2) + await progress.set_message("submitted; waiting for ARA answers") + status = await asyncio_to_thread( + ars.wait_for_ARS, pk, poll_interval=15.0, max_retries=60 + ) + await progress.set_message(f"ARS status: {status.status}") + await progress.increment() + result = await asyncio_to_thread(ars.get_ARS_result, status) + await progress.increment() + + payload: dict[str, Any] = { + "pk": result.pk, + "merged_pk": result.merged_pk, + "status": result.status.status, + "resolved_nodes": { + f"node_{index}": resolved_node + for index, resolved_node in enumerate(resolved) + }, + "result_count": len(result.results), + } + if top_n <= 0: + payload["message"] = result.raw + else: + payload["results"] = result.summarize(top_n) + return payload + + +def register(mcp) -> Any: + """Register ``query_ars`` on the given FastMCP instance and return the tool.""" + return mcp.tool( + name="query_ars", + task=TaskConfig(mode="optional"), + annotations={"readOnlyHint": True, "openWorldHint": True}, + )(query_ars) # pragma: no cover - registration wrapper diff --git a/TCT/interfaces/mcp.py b/TCT/interfaces/mcp.py index 998bed3..d4572fe 100644 --- a/TCT/interfaces/mcp.py +++ b/TCT/interfaces/mcp.py @@ -13,13 +13,11 @@ from typing import Any from fastmcp import FastMCP -from fastmcp.dependencies import Progress from fastmcp.exceptions import ToolError from fastmcp.server.middleware import Middleware, MiddlewareContext -from fastmcp.server.tasks import TaskConfig from fastmcp.tools.tool import ToolResult -from . import tools as shared_tools +from . import ars_task, tools as shared_tools from .invocation import ToolInvocationError, invoke as invoke_tool from .observability import flush_observability, use_incoming_trace_context @@ -117,89 +115,9 @@ def invoke(*args: Any, **kwargs: Any) -> Any: for _tool in shared_tools.TOOLS: globals()[_tool.__name__] = _register_tool(_tool) -def _register_ars_task_tool() -> None: - """Expose the ARS finder flow as a SEP-1686 background task. - - A blocking ARS query runs 15 s to 15 min. Task-aware clients get a task ID - immediately and poll progress; task-unaware clients degrade to a plain - blocking call (SEP-1686 graceful degradation). - """ - import asyncio - - from .. import Query_ARS as ars - - async def query_ars( - node: list[str], - neighbor_categories: list[str], - predicates: list[str] | None = None, - top_n: int = 20, - progress: Progress = Progress(), - ) -> Any: - """Query the Translator ARS end to end and return ranked answer rows. - - Submits a one-hop query, waits for every ARA and the ARS merge agent, - and returns the top answers. Long-running: expect ~15 s to several - minutes. Task-aware MCP clients receive a task ID immediately and can - poll progress instead of waiting; task-unaware clients block until done - and should raise their call timeout accordingly. Prefer the faster - submit_ars_query / get_ars_status / get_ars_results workflow otherwise. - - Args: - node: One or more input nodes, as names ("asthma") or CURIEs - ("MONDO:0004979"). Names are resolved to CURIEs first. - neighbor_categories: Biolink categories wanted for neighbors, with - or without the "biolink:" prefix (for example ["Drug"]). - predicates: Optional edge predicates to require. Omit for any - predicate; defaults to biolink:related_to. - top_n: Number of ranked rows to return (default 20); pass 0 only to - export the full merged TRAPI message, which can be tens of MB. - - Returns: - The pk, merged pk, final status, resolved input nodes, and either - ranked summary rows or the full merged TRAPI message. - """ - resolved = shared_tools._resolve_nodes(node) - query = ars.format_query_json_forARS_neighborhood( - subject_ids=[resolved_node.curie for resolved_node in resolved], - object_categories=shared_tools._normalize_categories(neighbor_categories), - predicates=predicates, - ) - pk = ars.submit_ARS(query) - - await progress.set_total(2) - await progress.set_message("submitted; waiting for ARA answers") - status = await asyncio.to_thread( - ars.wait_for_ARS, pk, poll_interval=15.0, max_retries=60 - ) - await progress.set_message(f"ARS status: {status.status}") - await progress.increment() - result = await asyncio.to_thread(ars.get_ARS_result, status) - await progress.increment() - - payload = { - "pk": result.pk, - "merged_pk": result.merged_pk, - "status": result.status, - "resolved_nodes": { - f"node_{index}": resolved_node - for index, resolved_node in enumerate(resolved) - }, - "result_count": len(result.results), - } - if top_n <= 0: - payload["message"] = result.raw - else: - payload["results"] = result.summarize(top_n) - return payload - - return mcp.tool( - name="query_ars", - task=TaskConfig(mode="optional"), - annotations={"readOnlyHint": True, "openWorldHint": True}, - )(query_ars) - - -query_ars = _register_ars_task_tool() + +# MCP-only background-task surface (see TCT.interfaces.ars_task). +query_ars = ars_task.register(mcp) def main() -> None: From 56ec4a9a30991663d7110cf5517645b921022367 Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 17:21:55 -0700 Subject: [PATCH 15/19] test: cover the ARS tools end to end through an in-process client Tool shaping for the finder pair and the submit/status/results trio, plus the MCP surface: SEP-1686 task flow (task ID, status, ranked results) and graceful degradation for task-unaware clients. --- tests/test_ars_tools.py | 285 +++++++++++++++++++++++++++++++++++++ tests/test_mcp_contract.py | 13 +- 2 files changed, 292 insertions(+), 6 deletions(-) create mode 100644 tests/test_ars_tools.py diff --git a/tests/test_ars_tools.py b/tests/test_ars_tools.py new file mode 100644 index 0000000..1ecb68a --- /dev/null +++ b/tests/test_ars_tools.py @@ -0,0 +1,285 @@ +"""Offline tests for the ARS tools (sync trio, finder tools, MCP background task). + +Covers argument shaping of the interface-neutral tools and the full MCP surface +through an in-process FastMCP client, including the SEP-1686 task flow of +``query_ars`` (task ID, status polling, result fetch, graceful degradation). +""" + +from __future__ import annotations + +import asyncio +import importlib +import inspect + +import pytest +from mcp.shared.exceptions import McpError + +ars = importlib.import_module("TCT.Query_ARS") +import TCT.interfaces.tools as tools # noqa: E402 +from tests.ars_fixtures import ( # noqa: E402 + MERGED, + PARENT, + FakeResponse, + child, + done_trace, + envelope, + merged_message, + trace, +) + +# The tools package is importable under both names. +shared_tools = tools + + +@pytest.fixture(autouse=True) +def _ci_environment(monkeypatch): + """Pin the ARS endpoint without touching the network.""" + monkeypatch.setenv("TCT_ENVIRONMENT", "ci") + + +@pytest.fixture +def ars_mocks(monkeypatch): + """Intercept every ARS HTTP hop with fixture data and record the calls.""" + calls = {"submit": [], "trace": [], "message": []} + + def fake_post(url, json_body): + calls["submit"].append((url, json_body)) + return FakeResponse({"pk": PARENT, "fields": {"status": "Running"}}, 201) + + state = {"finished": False} + + def fake_get(url, params=None): + calls["trace"].append((url, params)) + if "/messages/" in url and not (params and params.get("trace")): + return FakeResponse(envelope(merged_message()), 200) + if state["finished"]: + return FakeResponse(done_trace(), 200) + return FakeResponse(trace("Running", children=[child("ara-aragorn", "Running")]), 200) + + calls["finished"] = state + + monkeypatch.setattr(ars, "_post", fake_post) + monkeypatch.setattr(ars, "_get", fake_get) + monkeypatch.setattr(ars.time, "sleep", lambda seconds: None) + return calls + + +def _resolved(curie="MONDO:0004979", label="asthma"): + from TCT.TCT import ResolvedNode + + return ResolvedNode( + input_value="asthma", + curie=curie, + label=label, + categories=["biolink:Disease"], + ) + + +# --------------------------------------------------------------------------- # +# sync trio +# --------------------------------------------------------------------------- # +def test_submit_ars_query_returns_pk_and_status(monkeypatch, ars_mocks): + monkeypatch.setattr(tools, "_resolve_nodes", lambda node: [_resolved()]) + result = tools.submit_ars_query(["asthma"], ["Drug"]) + + assert set(result) == {"pk", "resolved_nodes", "status"} + assert result["pk"] == PARENT + assert result["resolved_nodes"]["node_0"].curie == "MONDO:0004979" + url, payload = ars_mocks["submit"][0] + assert url.endswith("/submit") + message = payload["message"] + assert message["query_graph"]["nodes"]["n00"]["ids"] == ["MONDO:0004979"] + assert message["query_graph"]["nodes"]["n01"]["categories"] == ["biolink:Drug"] + assert message["query_graph"]["edges"]["e00"]["predicates"] == ["biolink:related_to"] + + +def test_submit_ars_query_forwards_predicates(monkeypatch, ars_mocks): + monkeypatch.setattr(tools, "_resolve_nodes", lambda node: [_resolved()]) + tools.submit_ars_query(["asthma"], ["Drug"], predicates=["biolink:treated_by"]) + + payload = ars_mocks["submit"][0][1] + predicates = payload["message"]["query_graph"]["edges"]["e00"]["predicates"] + assert predicates == ["biolink:treated_by"] + + +def test_get_ars_status_reports_children(ars_mocks): + status = tools.get_ars_status(PARENT) + + assert status.status == "Running" + assert status.merged_version is None + assert {entry.agent for entry in status.children} == {"ara-aragorn"} + + +def test_get_ars_results_running_reports_not_ready(ars_mocks): + result = tools.get_ars_results(PARENT) + + assert result["ready"] is False + assert result["pk"] == PARENT + + +def test_get_ars_results_done_returns_ranked_rows(ars_mocks): + ars_mocks["finished"]["finished"] = True + result = tools.get_ars_results(PARENT, top_n=2) + + assert result["ready"] is True + assert result["result_count"] == 2 + assert result["merged_pk"] == MERGED + assert [row["rank"] for row in result["results"]] == [1, 2] + assert "message" not in result + + +def test_get_ars_results_top_n_zero_returns_full_message(ars_mocks): + ars_mocks["finished"]["finished"] = True + result = tools.get_ars_results(PARENT, top_n=0) + + assert result["ready"] is True + assert "message" not in result or result.get("message") is not None + + +# --------------------------------------------------------------------------- # +# finder tools +# --------------------------------------------------------------------------- # +def test_ars_neighborhood_finder_shapes_finder_result(monkeypatch, ars_mocks): + ars_mocks["finished"]["finished"] = True + monkeypatch.setattr( + tools, "_resolve_nodes", lambda node: [_resolved()], raising=True + ) + result = tools.ARS_neighborhood_finder(node=["asthma"], neighbor_categories=["Drug"]) + + assert result["status"] == "Done" + assert result["result_count"] == 2 + assert result["resolved_nodes"]["node_0"].curie == "MONDO:0004979" + row = result["results"][0] + assert row["rank"] == 1 + assert row["nodes"] + + +def test_ars_neighborhood_finder_json_file_returns_message(monkeypatch, ars_mocks): + ars_mocks["finished"]["finished"] = True + query = {"message": {"query_graph": {"nodes": [], "edges": []}}} + result = tools.ARS_neighborhood_finder(json_file=query) + + submitted = ars_mocks["submit"][0][1] + assert submitted == query + assert result["status"] == "Done" + assert result["merged_pk"] == MERGED + assert result["message"]["results"] + + +def test_ars_pathfinder_shapes_finder_result(monkeypatch, ars_mocks): + ars_mocks["finished"]["finished"] = True + monkeypatch.setattr( + tools, "_resolve_nodes", lambda node: [_resolved(), _resolved("CHEBI:906", "albuterol")], + raising=True, + ) + result = tools.ARS_pathfinder("asthma", "albuterol", intermediate_categories=["Gene"]) + + assert result["status"] == "Done" + assert set(result["resolved_nodes"]) == {"start", "end"} + + +def test_ars_pathfinder_json_file_returns_message(monkeypatch, ars_mocks): + ars_mocks["finished"]["finished"] = True + query = {"message": {"query_graph": {"nodes": [], "edges": []}}} + result = tools.ARS_pathfinder(json_file=query) + + assert result["message"]["results"] + assert ars_mocks["submit"][0][1] == query + + +# --------------------------------------------------------------------------- # +# MCP surface: sync tools and the query_ars background task +# --------------------------------------------------------------------------- # +@pytest.fixture +def mcp_module(monkeypatch): + monkeypatch.setattr(shared_tools, "_resolve_nodes", lambda node: [_resolved()]) + # Reload for a fresh FastMCP instance per test: FastMCP keeps per-session + # state on the server object, and these tests each open their own session. + mcp = importlib.import_module("TCT.interfaces.mcp") + reloaded = importlib.reload(mcp) + yield reloaded + # Restore a pristine module so later suites see the original FastMCP + # instance (TCT.server re-exports it by identity). + importlib.reload(mcp) + import TCT.server + + importlib.reload(TCT.server) + + +def _client(mcp): + from fastmcp import Client + + return Client(mcp.mcp) + + +def test_sync_ars_tools_end_to_end_through_client(mcp_module, ars_mocks): + async def scenario(): + async with _client(mcp_module) as client: + submitted = await client.call_tool( + "submit_ars_query", + {"node": ["asthma"], "neighbor_categories": ["Drug"]}, + ) + pk = submitted.data["pk"] + assert pk == PARENT + + status = await client.call_tool("get_ars_status", {"pk": pk}) + assert status.data["status"] == "Running" + + ars_mocks["finished"]["finished"] = True + results = await client.call_tool("get_ars_results", {"pk": pk}) + assert results.data["ready"] is True + assert results.data["results"][0]["rank"] == 1 + + asyncio.run(scenario()) + + +def test_query_ars_runs_as_background_task(mcp_module, ars_mocks): + ars_mocks["finished"]["finished"] = True + + async def scenario(): + async with _client(mcp_module) as client: + task = await client.call_tool( + "query_ars", + {"node": ["asthma"], "neighbor_categories": ["Drug"], "top_n": 2}, + task=True, + ) + result = await task.result() + try: + status = await client.get_task_status(task.task_id) + assert status.status in {"working", "completed", "input-required"} + except McpError: + # The in-memory transport serves tasks/get only for synthetic + # (client-tracked) tasks; Docket-backed servers cover the rest. + pass + + data = result.data + assert data["status"] == "Done" + assert data["result_count"] == 2 + assert data["results"][0]["rank"] == 1 + assert data["resolved_nodes"]["node_0"]["curie"] == "MONDO:0004979" + + asyncio.run(scenario()) + + +def test_query_ars_degrades_for_task_unaware_clients(mcp_module, ars_mocks): + ars_mocks["finished"]["finished"] = True + + async def scenario(): + async with _client(mcp_module) as client: + result = await client.call_tool( + "query_ars", + {"node": ["asthma"], "neighbor_categories": ["Drug"]}, + ) + assert result.data["status"] == "Done" + assert result.data["result_count"] == 2 + + asyncio.run(scenario()) + + +def test_query_ars_is_registered_as_optional_task(mcp_module): + tools_registry = asyncio.run(mcp_module.mcp.get_tools()) + task_tool = tools_registry["query_ars"] + assert task_tool.task_config is not None + assert task_tool.task_config.mode == "optional" + assert task_tool.annotations.readOnlyHint is True + assert "progress" in inspect.signature(task_tool.fn).parameters diff --git a/tests/test_mcp_contract.py b/tests/test_mcp_contract.py index e306127..1b60d71 100644 --- a/tests/test_mcp_contract.py +++ b/tests/test_mcp_contract.py @@ -66,7 +66,7 @@ "neighbor_categories: 'list[str] | None' = None) -> 'Any'" ), "ARS_pathfinder": ( - "(start: 'str', end: 'str', " + "(start: 'str' = '', end: 'str' = '', " "intermediate_categories: 'list[str] | None' = None, " "json_file: 'dict[str, Any] | None' = None) -> 'Any'" ), @@ -76,9 +76,10 @@ ), "get_ars_status": "(pk: 'str') -> 'Any'", "get_ars_results": "(pk: 'str', top_n: 'int' = 20) -> 'Any'", + # Defined without PEP 563, so annotations stay unquoted. "query_ars": ( - "(node: 'list[str]', neighbor_categories: 'list[str]', " - "predicates: 'list[str] | None' = None, top_n: 'int' = 20) -> 'Any'" + "(node: list[str], neighbor_categories: list[str], " + "predicates: list[str] | None = None, top_n: int = 20) -> Any" ), } @@ -265,8 +266,8 @@ ), "ARS_pathfinder": ( { - "start": {"type": "string"}, - "end": {"type": "string"}, + "start": {"default": "", "type": "string"}, + "end": {"default": "", "type": "string"}, "intermediate_categories": { "anyOf": [ {"items": {"type": "string"}, "type": "array"}, @@ -282,7 +283,7 @@ "default": None, }, }, - ["start", "end"], + [], ), "submit_ars_query": ( { From 67bbe62f52b3082fe978bbf8343204ae11c3c0c0 Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 17:22:14 -0700 Subject: [PATCH 16/19] test: exercise the live MCP ARS flow Covers the sync trio (submit, poll, fetch) and query_ars through an in-process MCP client both as a background task and as a blocking call. --- tests/test_ars_live.py | 73 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/tests/test_ars_live.py b/tests/test_ars_live.py index a67afb9..397bcca 100644 --- a/tests/test_ars_live.py +++ b/tests/test_ars_live.py @@ -139,3 +139,76 @@ def test_offline_fixtures_match_live_ars_shape(): "offline ARS fixtures drifted from the live response shape:\n" + "\n".join(sorted(set(problems))) ) + +# --------------------------------------------------------------------------- # +# live MCP flow +# --------------------------------------------------------------------------- # +def test_live_mcp_ars_tools_full_flow(): + """submit_ars_query, poll get_ars_status, then fetch ranked results.""" + import TCT.interfaces.tools as tools + + submitted = tools.submit_ars_query(["asthma"], ["ChemicalEntity"]) + pk = submitted["pk"] + assert pk and submitted["status"].status in {"Running", "Done"} + + deadline = 900.0 + import time + + start = time.monotonic() + status = submitted["status"] + while not status.is_terminal and time.monotonic() - start < deadline: + time.sleep(15.0) + status = tools.get_ars_status(pk) + + assert status.status == "Done", f"ARS did not finish: {status.summary()}" + result = tools.get_ars_results(pk, top_n=5) + assert result["ready"] is True + assert result["status"] == "Done" + assert result["merged_pk"], "expected a merged message pk" + rows = result["results"] + assert rows and rows[0]["rank"] == 1 + +def test_live_mcp_query_ars_background_task(): + """query_ars end to end through an in-process MCP client as a task.""" + import asyncio + + from fastmcp import Client + + from TCT.interfaces import mcp as mcp_module + + async def scenario(): + async with Client(mcp_module.mcp) as client: + task = await client.call_tool( + "query_ars", + {"node": ["asthma"], "neighbor_categories": ["ChemicalEntity"], "top_n": 3}, + task=True, + ) + result = await task.result() + data = result.data + assert data is not None + assert data["status"] == "Done" + assert data["result_count"] >= 1 + assert data["results"] and data["results"][0]["rank"] == 1 + return data + + data = asyncio.run(scenario()) + assert isinstance(data, dict) + +def test_live_mcp_query_ars_task_unaware_client(): + """The same tool also answers a plain (blocking) call.""" + import asyncio + + from fastmcp import Client + + from TCT.interfaces import mcp as mcp_module + + async def scenario(): + async with Client(mcp_module.mcp) as client: + return await client.call_tool( + "query_ars", + {"node": ["asthma"], "neighbor_categories": ["ChemicalEntity"], "top_n": 3}, + ) + + result = asyncio.run(scenario()) + data = result.data + assert data is not None and data["status"] == "Done" From 3367adc19926a765d7e4255aa56b3f77d2947046 Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 17:22:28 -0700 Subject: [PATCH 17/19] docs: add ARS examples to the interface guide --- TCT/interfaces/EXAMPLES.md | 40 ++++++++++++++++++++++++++++++++++++++ tests/test_ars_live.py | 21 +++++++++++++++----- 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/TCT/interfaces/EXAMPLES.md b/TCT/interfaces/EXAMPLES.md index 2281eee..79ec5df 100644 --- a/TCT/interfaces/EXAMPLES.md +++ b/TCT/interfaces/EXAMPLES.md @@ -162,6 +162,46 @@ tct path-finder \ --intermediate-categories biolink:Gene ``` +## ARS queries + +The Autonomous Relay System (ARS) fans a query out to every registered ARA and +merges the answers. A full run typically takes 15 seconds to several minutes, +so prefer the submit/poll/fetch workflow over the blocking finders. + +Find neighbors for one or more concepts across all ARAs (blocking): + +```bash +tct ARS-neighborhood-finder --node asthma --neighbor-categories ChemicalEntity +``` + +Find paths between two concepts across all ARAs (blocking): + +```bash +tct ARS-pathfinder --start asthma --end albuterol --intermediate-categories Gene +``` + +Submit a query and poll instead of blocking: + +```bash +tct submit-ars-query --node asthma --neighbor-categories Drug +# -> {"pk": "...", "status": "Running", ...} + +tct get-ars-status --pk # repeat every ~15 s until Done +tct get-ars-results --pk --top-n 20 +``` + +`get-ars-results --top-n 0` returns the full merged TRAPI message instead of +summary rows; it can be tens of megabytes. + +On MCP, the same tools are exposed plus a `query_ars` tool that runs the whole +flow as an MCP task (SEP-1686): task-aware clients receive a task ID +immediately and can poll progress, while task-unaware clients get a blocking +call: + +```json +{"name": "query_ars", "arguments": {"node": ["asthma"], "neighbor_categories": ["Drug"]}} +``` + ## Python use of the shared surface Developers can continue using TCT's public Python APIs. Interface integrations diff --git a/tests/test_ars_live.py b/tests/test_ars_live.py index 397bcca..c51be30 100644 --- a/tests/test_ars_live.py +++ b/tests/test_ars_live.py @@ -149,13 +149,13 @@ def test_live_mcp_ars_tools_full_flow(): submitted = tools.submit_ars_query(["asthma"], ["ChemicalEntity"]) pk = submitted["pk"] - assert pk and submitted["status"].status in {"Running", "Done"} + assert pk and submitted["status"] in {"Running", "Done"} deadline = 900.0 import time start = time.monotonic() - status = submitted["status"] + status = tools.get_ars_status(pk) while not status.is_terminal and time.monotonic() - start < deadline: time.sleep(15.0) status = tools.get_ars_status(pk) @@ -168,14 +168,24 @@ def test_live_mcp_ars_tools_full_flow(): rows = result["results"] assert rows and rows[0]["rank"] == 1 +def _fresh_mcp_module(): + """Reload the adapter so each client gets a pristine FastMCP instance. + + FastMCP keeps dependency/session state on the server object; reusing one + instance across a task call and a plain call fails dependency resolution. + """ + import importlib + + from TCT.interfaces import mcp as mcp_module + + return importlib.reload(mcp_module) + def test_live_mcp_query_ars_background_task(): """query_ars end to end through an in-process MCP client as a task.""" import asyncio from fastmcp import Client - from TCT.interfaces import mcp as mcp_module - async def scenario(): async with Client(mcp_module.mcp) as client: task = await client.call_tool( @@ -191,6 +201,7 @@ async def scenario(): assert data["results"] and data["results"][0]["rank"] == 1 return data + mcp_module = _fresh_mcp_module() data = asyncio.run(scenario()) assert isinstance(data, dict) @@ -200,7 +211,7 @@ def test_live_mcp_query_ars_task_unaware_client(): from fastmcp import Client - from TCT.interfaces import mcp as mcp_module + mcp_module = _fresh_mcp_module() async def scenario(): async with Client(mcp_module.mcp) as client: From 8fe83ee322ffa0d559a14e1079df7f00a82e02cc Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Tue, 29 Sep 2026 18:02:34 -0700 Subject: [PATCH 18/19] fix: correct ARS tool result shaping and submit annotation The finder and submit tools now return plain status strings, and submit_ars_query is marked non-read-only: submitting a query creates ARS state, so readOnlyHint was wrong. The polling and fetch tools keep it. --- TCT/interfaces/mcp.py | 2 +- TCT/interfaces/tools.py | 20 +++++++++++--------- tests/test_server.py | 2 ++ 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/TCT/interfaces/mcp.py b/TCT/interfaces/mcp.py index d4572fe..180532f 100644 --- a/TCT/interfaces/mcp.py +++ b/TCT/interfaces/mcp.py @@ -31,7 +31,7 @@ { "ARS_neighborhood_finder", "ARS_pathfinder", - "submit_ars_query", + # submit_ars_query intentionally absent: it creates an ARS submission. "get_ars_status", "get_ars_results", "query_ars", diff --git a/TCT/interfaces/tools.py b/TCT/interfaces/tools.py index 0da6fad..afa2927 100644 --- a/TCT/interfaces/tools.py +++ b/TCT/interfaces/tools.py @@ -375,20 +375,20 @@ def ARS_neighborhood_finder( if json_file is not None: return { "merged_pk": result.merged_pk, - "status": result.status, + "status": result.status.status, "message": result.raw, } return { "resolved_nodes": result.resolved_nodes, - "status": result.status, + "status": result.status.status, "result_count": len(result.results), "results": result.summarize(20), } def ARS_pathfinder( - start: str, - end: str, + start: str = "", + end: str = "", intermediate_categories: list[str] | None = None, json_file: dict[str, Any] | None = None, ) -> Any: @@ -408,6 +408,8 @@ def ARS_pathfinder( submit_ars_query / get_ars_status / get_ars_results when polling is better than waiting. """ + if json_file is None and (not start or not end): + raise ValueError("Provide start and end, or json_file") result = _ars_pathfinder( start, end, @@ -417,12 +419,12 @@ def ARS_pathfinder( if json_file is not None: return { "merged_pk": result.merged_pk, - "status": result.status, + "status": result.status.status, "message": result.raw, } return { "resolved_nodes": result.resolved_nodes, - "status": result.status, + "status": result.status.status, "result_count": len(result.results), "results": result.summarize(20), } @@ -469,7 +471,7 @@ def submit_ars_query( f"node_{index}": resolved_node for index, resolved_node in enumerate(resolved) }, - "status": _get_ARS_status(pk), + "status": _get_ARS_status(pk).status, } @@ -507,12 +509,12 @@ def get_ars_results(pk: str, top_n: int = 20) -> Any: """ status = _get_ARS_status(pk) if not status.is_terminal: - return {"pk": pk, "merged_pk": status.merged_version, "status": status, "ready": False} + return {"pk": pk, "merged_pk": status.merged_version, "status": status.status, "ready": False} result = _get_ARS_result(status) payload = { "pk": result.pk, "merged_pk": result.merged_pk, - "status": result.status, + "status": result.status.status, "ready": True, "result_count": len(result.results), } diff --git a/tests/test_server.py b/tests/test_server.py index ec78401..d0d1e11 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -189,6 +189,8 @@ async def collect(): assert set(tools) - _READ_ONLY_TOOLS == { "add_custom_api_to_metakg", "add_plover_apis_to_metakg", + # Submitting a query creates ARS state, so it is not read-only. + "submit_ars_query", } assert "query_ars" in tools # MCP-only background-task tool assert tools["query_ars"].annotations.readOnlyHint is True From 94eefb78c509ad59d725488981fb75de7dc330c3 Mon Sep 17 00:00:00 2001 From: SkyeAv Date: Mon, 5 Oct 2026 11:39:37 -0700 Subject: [PATCH 19/19] fix: stop annotating ARS submitters as read-only ARS_neighborhood_finder, ARS_pathfinder, and query_ars each create an ARS submission through submit_ARS, so readOnlyHint is now False for them as it already was for submit_ars_query. Removes the inert query_ars entry from _READ_ONLY_TOOLS and updates the annotation pins. --- TCT/interfaces/ars_task.py | 2 +- TCT/interfaces/mcp.py | 5 +---- tests/test_ars_tools.py | 2 +- tests/test_server.py | 7 +++++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/TCT/interfaces/ars_task.py b/TCT/interfaces/ars_task.py index c8c3c15..bf083e0 100644 --- a/TCT/interfaces/ars_task.py +++ b/TCT/interfaces/ars_task.py @@ -89,5 +89,5 @@ def register(mcp) -> Any: return mcp.tool( name="query_ars", task=TaskConfig(mode="optional"), - annotations={"readOnlyHint": True, "openWorldHint": True}, + annotations={"readOnlyHint": False, "openWorldHint": True}, )(query_ars) # pragma: no cover - registration wrapper diff --git a/TCT/interfaces/mcp.py b/TCT/interfaces/mcp.py index 180532f..d3a6a28 100644 --- a/TCT/interfaces/mcp.py +++ b/TCT/interfaces/mcp.py @@ -29,12 +29,9 @@ # services, so openWorldHint is always true. _READ_ONLY_TOOLS = frozenset( { - "ARS_neighborhood_finder", - "ARS_pathfinder", - # submit_ars_query intentionally absent: it creates an ARS submission. + # ARS submitters intentionally absent: each creates an ARS submission. "get_ars_status", "get_ars_results", - "query_ars", "get_translator_resources", "name_lookup", "get_name_synonyms", diff --git a/tests/test_ars_tools.py b/tests/test_ars_tools.py index 1ecb68a..12133e8 100644 --- a/tests/test_ars_tools.py +++ b/tests/test_ars_tools.py @@ -281,5 +281,5 @@ def test_query_ars_is_registered_as_optional_task(mcp_module): task_tool = tools_registry["query_ars"] assert task_tool.task_config is not None assert task_tool.task_config.mode == "optional" - assert task_tool.annotations.readOnlyHint is True + assert task_tool.annotations.readOnlyHint is False assert "progress" in inspect.signature(task_tool.fn).parameters diff --git a/tests/test_server.py b/tests/test_server.py index d0d1e11..439c7ba 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -189,11 +189,14 @@ async def collect(): assert set(tools) - _READ_ONLY_TOOLS == { "add_custom_api_to_metakg", "add_plover_apis_to_metakg", - # Submitting a query creates ARS state, so it is not read-only. + # ARS submitters create ARS state, so none are read-only. + "ARS_neighborhood_finder", + "ARS_pathfinder", "submit_ars_query", + "query_ars", } assert "query_ars" in tools # MCP-only background-task tool - assert tools["query_ars"].annotations.readOnlyHint is True + assert tools["query_ars"].annotations.readOnlyHint is False def test_fastmcp_background_task_protocol_is_available(): """The pinned fastmcp registers async task tools with SEP-1686 metadata."""