Arax pathfinder

compose.yml

@@ -273,6 +273,20 @@ services:
     volumes:
       - ./logs:/app/logs
       - ./.env:/app/.env
+  arax_pathfinder:
+    container_name: arax_pathfinder
+    build:
+      context: .
+      dockerfile: workers/arax_pathfinder/Dockerfile
+    restart: unless-stopped
+    depends_on:
+      shepherd_db:
+        condition: service_healthy
+      shepherd_broker:
+        condition: service_healthy
+    volumes:
+      - ./logs:/app/logs
+      - ./.env:/app/.env
 
 ######### BTE
   bte:

shepherd_server/main.py

@@ -1,6 +1,5 @@
 import uvicorn
 
-
 if __name__ == "__main__":
     uvicorn.run(
         "shepherd_server.server:APP",

shepherd_utils/config.py

@@ -23,7 +23,23 @@ class Settings(BaseSettings):
     sync_kg_retrieval_url: str = "https://strider.renci.org/query"
     default_data_tier: int = 0
     omnicorp_url: str = "https://aragorn-ranker.renci.org/omnicorp_overlay"
+
+    # ARAX configs
     arax_url: str = "https://arax.ncats.io/shepherd/api/arax/v1.4/query"
+    plover_url: str = "https://kg2cplover3.rtx.ai:9990"
+    curie_ngd_addr: str = (
+        "mysql:arax-databases-mysql.rtx.ai:public_ro:curie_ngd_v1_0_kg2_10_2"
+    )
+    node_degree_addr: str = (
+        "mysql:arax-databases-mysql.rtx.ai:public_ro:kg2c_v1_0_kg2_10_2"
+    )
+    arax_biolink_version: str = "4.2.5"
+    arax_blocked_list_url: str = (
+        "https://raw.githubusercontent.com/RTXteam/RTX/master/"
+        "code/ARAX/KnowledgeSources/general_concepts.json"
+    )
+    # End of ARAX configs
+
     node_norm: str = "https://biothings.ci.transltr.io/nodenorm/api/"
 
     pathfinder_redis_host: str = "host.docker.internal"

shepherd_utils/inject_shepherd_arax_provenance.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+
+SHEPHERD_ARAX_SOURCE = {
+    "resource_id": "infores:shepherd-arax",
+    "resource_role": "aggregator_knowledge_source",
+    "source_record_urls": None,
+    "upstream_resource_ids": ["infores:arax"],
+}
+
+
+def add_shepherd_arax_to_edge_sources(trapi_response: Dict[str, Any]) -> Dict[str, Any]:
+    message = trapi_response.get("message")
+    if not isinstance(message, dict):
+        return trapi_response
+
+    kg = message.get("knowledge_graph")
+    if not isinstance(kg, dict):
+        return trapi_response
+
+    edges = kg.get("edges")
+    if not isinstance(edges, dict):
+        return trapi_response
+
+    for edge_id, edge_obj in edges.items():
+        if not isinstance(edge_obj, dict):
+            continue
+
+        sources = edge_obj.get("sources")
+        if sources is None:
+            edge_obj["sources"] = [dict(SHEPHERD_ARAX_SOURCE)]
+            continue
+
+        if not isinstance(sources, list):
+            continue
+
+        already_present = any(
+            isinstance(s, dict) and s.get("resource_id") == "infores:shepherd-arax"
+            for s in sources
+        )
+        if not already_present:
+            sources.append(dict(SHEPHERD_ARAX_SOURCE))
+
+    return trapi_response

workers/arax/worker.py

@@ -9,6 +9,7 @@
 from shepherd_utils.db import get_message, save_message
 from shepherd_utils.shared import get_tasks, wrap_up_task
 from shepherd_utils.otel import setup_tracer
+from shepherd_utils.inject_shepherd_arax_provenance import add_shepherd_arax_to_edge_sources
 
 # Queue name
 STREAM = "arax"
@@ -18,33 +19,50 @@
 tracer = setup_tracer(STREAM)
 
 
-async def arax(task, logger: logging.Logger):
+def is_pathfinder_query(message):
     try:
-        start = time.time()
-        query_id = task[1]["query_id"]
-        logger.info(f"Getting message from db for query id {query_id}")
-        message = await get_message(query_id, logger)
-        message["submitter"] = "Shepherd"
-        logger.info(f"Get the message from db {message}")
-
-        headers = {"Content-Type": "application/json"}
-        response = requests.post(settings.arax_url, json=message, headers=headers)
-
-        logger.info(f"Status Code from ARAX response: {response.status_code}")
-        result = response.json()
-
-    except Exception as e:
-        logger.error(f"Error occurred in ARAX entry module: {e}")
-        result = {"status": "error", "error": str(e)}
-
-    response_id = task[1]["response_id"]
+        # this can still fail if the input looks like e.g.:
+        #  "query_graph": None
+        qedges = message.get("message", {}).get("query_graph", {}).get("edges", {})
+    except:
+        qedges = {}
+    try:
+        # this can still fail if the input looks like e.g.:
+        #  "query_graph": None
+        qpaths = message.get("message", {}).get("query_graph", {}).get("paths", {})
+    except:
+        qpaths = {}
+    if len(qpaths) > 1:
+        raise Exception("Only a single path is supported", 400)
+    if (len(qpaths) > 0) and (len(qedges) > 0):
+        raise Exception("Mixed mode pathfinder queries are not supported", 400)
+    return len(qpaths) == 1
 
-    await save_message(response_id, result, logger)
 
-    workflow = [{"id": "arax"}]
+async def arax(task, logger: logging.Logger):
+    start = time.time()
+    query_id = task[1]["query_id"]
+    logger.info(f"Getting message from db for query id {query_id}")
+    message = await get_message(query_id, logger)
+    if is_pathfinder_query(message):
+        workflow = [{"id": "arax.pathfinder"}]
+    else:
+        try:
+            workflow = [{"id": "arax"}]
+            message["submitter"] = "Shepherd"
+            logger.info(f"Get the message from db {message}")
+            headers = {"Content-Type": "application/json"}
+            response = requests.post(settings.arax_url, json=message, headers=headers)
+            logger.info(f"Status Code from ARAX response: {response.status_code}")
+            result = response.json()
+            result = add_shepherd_arax_to_edge_sources(result)
+        except Exception as e:
+            logger.error(f"Error occurred calling ARAX service: {e}")
+            result = {"status": "error", "error": str(e)}
+        response_id = task[1]["response_id"]
+        await save_message(response_id, result, logger)
 
     await wrap_up_task(STREAM, GROUP, task, workflow, logger)
-
     logger.info(f"Finished task {task[0]} in {time.time() - start}")
 
 

workers/arax_pathfinder/Dockerfile

@@ -0,0 +1,34 @@
+# Use RENCI python base image
+FROM ghcr.io/translatorsri/renci-python-image:3.11.5
+
+# Add image info
+LABEL org.opencontainers.image.source https://github.com/BioPack-team/shepherd
+
+ENV PYTHONHASHSEED=0
+
+# set up requirements
+WORKDIR /app
+
+# make sure all is writeable for the nru USER later on
+RUN chmod -R 777 .
+
+# Install requirements
+COPY ./shepherd_utils ./shepherd_utils
+COPY ./pyproject.toml .
+RUN pip install .
+
+COPY ./workers/arax_pathfinder/requirements.txt .
+RUN pip install -r requirements.txt
+
+# switch to the non-root user (nru). defined in the base image
+USER nru
+
+# Copy in files
+COPY ./workers/arax_pathfinder ./
+
+# Set up base for command and any variables
+# that shouldn't be modified
+# ENTRYPOINT ["uvicorn", "shepherd_server.server:APP"]
+
+# Variables that can be overriden
+CMD ["python", "worker.py"]

workers/arax_pathfinder/requirements.txt

@@ -0,0 +1,2 @@
+catrax-pathfinder==1.1.1
+biolink-helper-pkg==1.0.0