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/TCT/Query_ARS.py b/TCT/Query_ARS.py new file mode 100644 index 0000000..b7cd642 --- /dev/null +++ b/TCT/Query_ARS.py @@ -0,0 +1,776 @@ +"""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 _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/TCT/config.py b/TCT/config.py index ad4c2f9..f1ff96e 100644 --- a/TCT/config.py +++ b/TCT/config.py @@ -33,8 +33,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/", @@ -59,6 +60,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/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/TCT/interfaces/ars_task.py b/TCT/interfaces/ars_task.py new file mode 100644 index 0000000..bf083e0 --- /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": False, "openWorldHint": True}, + )(query_ars) # pragma: no cover - registration wrapper diff --git a/TCT/interfaces/mcp.py b/TCT/interfaces/mcp.py index 5af3982..d3a6a28 100644 --- a/TCT/interfaces/mcp.py +++ b/TCT/interfaces/mcp.py @@ -17,7 +17,7 @@ from fastmcp.server.middleware import Middleware, MiddlewareContext 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 @@ -29,6 +29,9 @@ # services, so openWorldHint is always true. _READ_ONLY_TOOLS = frozenset( { + # ARS submitters intentionally absent: each creates an ARS submission. + "get_ars_status", + "get_ars_results", "get_translator_resources", "name_lookup", "get_name_synonyms", @@ -110,6 +113,10 @@ def invoke(*args: Any, **kwargs: Any) -> Any: globals()[_tool.__name__] = _register_tool(_tool) +# MCP-only background-task surface (see TCT.interfaces.ars_task). +query_ars = ars_task.register(mcp) + + def main() -> None: """Entry point for the installed ``tct-server`` command.""" try: @@ -119,4 +126,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]] diff --git a/TCT/interfaces/tools.py b/TCT/interfaces/tools.py index eb92800..afa2927 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,186 @@ 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.status, + "message": result.raw, + } + return { + "resolved_nodes": result.resolved_nodes, + "status": result.status.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. + """ + 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, + intermediate_categories=intermediate_categories, + json_file=json_file, + ) + if json_file is not None: + return { + "merged_pk": result.merged_pk, + "status": result.status.status, + "message": result.raw, + } + return { + "resolved_nodes": result.resolved_nodes, + "status": result.status.status, + "result_count": len(result.results), + "results": result.summarize(20), + } + + +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).status, + } + + +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.status, "ready": False} + result = _get_ARS_result(status) + payload = { + "pk": result.pk, + "merged_pk": result.merged_pk, + "status": result.status.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, @@ -353,6 +542,11 @@ def path_finder( trapi_query_endpoint, neighborhood_finder, path_finder, + ARS_neighborhood_finder, + ARS_pathfinder, + submit_ars_query, + get_ars_status, + get_ars_results, ) __all__ = [tool.__name__ for tool in TOOLS] + ["TOOLS"] 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/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/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") 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 new file mode 100644 index 0000000..de68a11 --- /dev/null +++ b/tests/ars_fixtures.py @@ -0,0 +1,115 @@ +"""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. + + ``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": f"[['{merged}', 'ars']]" if merged else "[]", + "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": [ + { + "normalized_score": 0.9, + "node_bindings": {"n00": [{"id": "MONDO:1"}], "n01": [{"id": "CHEBI:1"}]}, + "analyses": [ + { + "resource_id": "infores:arax", + "edge_bindings": {"e00": [{"id": "e1"}]}, + } + ], + }, + {"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_ars_live.py b/tests/test_ars_live.py new file mode 100644 index 0000000..c51be30 --- /dev/null +++ b/tests/test_ars_live.py @@ -0,0 +1,225 @@ +"""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(): + # 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) + # 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(): + 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 == "*": + # 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) + 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"], + ), 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, + ) + 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" + + "\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"] in {"Running", "Done"} + + deadline = 900.0 + import time + + start = time.monotonic() + 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) + + 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 _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 + + 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 + + mcp_module = _fresh_mcp_module() + 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 + + mcp_module = _fresh_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" diff --git a/tests/test_ars_tools.py b/tests/test_ars_tools.py new file mode 100644 index 0000000..12133e8 --- /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 False + assert "progress" in inspect.signature(task_tool.fn).parameters 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") diff --git a/tests/test_mcp_contract.py b/tests/test_mcp_contract.py index 4665b22..1b60d71 100644 --- a/tests/test_mcp_contract.py +++ b/tests/test_mcp_contract.py @@ -60,6 +60,27 @@ "(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'", + # 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" + ), } @@ -223,6 +244,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": {"default": "", "type": "string"}, + "end": {"default": "", "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, + }, + }, + [], + ), + "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 +333,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_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 diff --git a/tests/test_query_ars.py b/tests/test_query_ars.py new file mode 100644 index 0000000..7071fe1 --- /dev/null +++ b/tests/test_query_ars.py @@ -0,0 +1,388 @@ +"""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]["nodes"]["n01"][0]["name"] == "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(): + # 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"]["n01"][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) == [] diff --git a/tests/test_server.py b/tests/test_server.py index d82a3aa..439c7ba 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -189,7 +189,14 @@ async def collect(): assert set(tools) - _READ_ONLY_TOOLS == { "add_custom_api_to_metakg", "add_plover_apis_to_metakg", + # 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 False 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" """