Arax pathfinder

compose.yml

@@ -338,6 +338,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
 
   arax_rank:
     container_name: arax_rank

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://kg2cploverdb.ci.transltr.io"
+    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"

workers/arax/worker.py

@@ -3,14 +3,18 @@
 import asyncio
 import json
 import logging
-import requests
 import time
 import uuid
+
+import requests
+
 from shepherd_utils.config import settings
 from shepherd_utils.db import get_message, save_message
-from shepherd_utils.shared import get_tasks, handle_task_failure, wrap_up_task
+from shepherd_utils.inject_shepherd_arax_provenance import (
+    add_shepherd_arax_to_edge_sources,
+)
 from shepherd_utils.otel import setup_tracer
-from inject_shepherd_arax_provenance import add_shepherd_arax_to_edge_sources
+from shepherd_utils.shared import get_tasks, handle_task_failure, wrap_up_task
 
 # Queue name
 STREAM = "arax"
@@ -20,30 +24,50 @@
 tracer = setup_tracer(STREAM)
 
 
-async def arax(task, logger: logging.Logger):
+def is_pathfinder_query(message):
     try:
-        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()
-        result = add_shepherd_arax_to_edge_sources(result)
-
-    except Exception as e:
-        logger.error(f"Error occurred in ARAX entry module: {e}")
-        result = {"status": "error", "error": str(e)}
+        # 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
 
-    response_id = task[1]["response_id"]
 
-    await save_message(response_id, result, logger)
+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):
+        task[1]["workflow"] = json.dumps([{"id": "arax.pathfinder"}])
+    else:
+        try:
+            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)
+        task[1]["workflow"] = json.dumps([{"id": "arax"}])
 
-    task[1]["workflow"] = json.dumps([{"id": "arax"}])
+    logger.info(f"Finished task {task[0]} in {time.time() - start}")
 
 
 async def process_task(task, parent_ctx, logger: logging.Logger, limiter):

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.2.2
+biolink-helper-pkg==1.0.0