Merge pull request #59 from BioPack-team/merge_message_bugs

Merge message bugs

Author: maximusunc

shepherd_broker/redis.conf

@@ -2,12 +2,13 @@ requirepass supersecretpassword
 port 6379
 maxmemory 4gb
 # evict least frequently used keys when memory cap is hit
-maxmemory-policy volatile-lfu
+maxmemory-policy allkeys-lru
 loglevel notice
 # If we want to log to a file
 logfile /data/shepherd_broker.log
 save 60 50
 stop-writes-on-bgsave-error no
+dir /data
 dbfilename shepherd_broker.rdb
 
 # enable larger entry writes

shepherd_server/openapi-config.yaml

@@ -4,7 +4,7 @@ contact:
   x-id: https://github.com/maximusunc
   x-role: responsible developer
 description: '<img src="/static/favicon.png" width="200px"><br /><br />Shepherd: Translator Autonomous Relay Agent Platform'
-version: 0.5.1
+version: 0.5.2
 servers:
 - description: Default server
   url: https://shepherd.renci.org

shepherd_utils/config.py

@@ -23,7 +23,7 @@ class Settings(BaseSettings):
     sync_kg_retrieval_url: str = "https://strider.renci.org/query"
     omnicorp_url: str = "https://aragorn-ranker.renci.org/omnicorp_overlay"
     arax_url: str = "https://arax.ncats.io/shepherd/api/arax/v1.4/query"
-    node_norm: str = "https://nodenormalization-sri.renci.org/"
+    node_norm: str = "https://biothings.ci.transltr.io/nodenorm/api/"
 
     pathfinder_redis_host: str = "host.docker.internal"
     pathfinder_redis_port: int = 6383
@@ -35,6 +35,9 @@ class Settings(BaseSettings):
     jaeger_host: str = "http://jaeger"
     jaeger_port: int = 4317
 
+    # ttl in seconds
+    redis_ttl: int = 1210000
+
     class Config:
         env_file = ".env"
 

shepherd_utils/db.py

@@ -105,8 +105,16 @@ async def add_query(
             connection_pool=data_db_pool,
         )
         # print(f"Putting {query_id} on {ara_target} stream")
-        await client.set(query_id, zstandard.compress(orjson.dumps(query)))
-        await client.set(response_id, zstandard.compress(orjson.dumps(query)))
+        await client.set(
+            query_id,
+            zstandard.compress(orjson.dumps(query)),
+            ex=settings.redis_ttl,
+        )
+        await client.set(
+            response_id,
+            zstandard.compress(orjson.dumps(query)),
+            ex=settings.redis_ttl,
+        )
         await client.aclose()
     except Exception as e:
         # failed to put message in db
@@ -153,7 +161,11 @@ async def save_message(
             connection_pool=data_db_pool,
         )
         # print(f"Putting {query_id} on {ara_target} stream")
-        await client.set(callback_id, compressed)
+        await client.set(
+            callback_id,
+            compressed,
+            ex=settings.redis_ttl,
+        )
         await client.aclose()
         logger.debug(f"Saving message took {time.time() - start} seconds")
     except Exception as e:
@@ -226,7 +238,9 @@ async def save_logs(
             new_logs = list(handler.contents())
             new_logs.reverse()
             existing_logs.extend(new_logs)
-        await client.set(response_id, orjson.dumps(existing_logs))
+        await client.set(
+            response_id, orjson.dumps(existing_logs), ex=settings.redis_ttl
+        )
         await client.aclose()
     except Exception as e:
         logger.error(f"Failed to save logs for response {response_id}: {e}")

shepherd_utils/shared.py

@@ -162,6 +162,16 @@ def recursive_get_auxgraph_edges(
     return edges, auxgraphs, nodes
 
 
+def is_support_edge(edge) -> bool:
+    """Checks if a given edge is a support edge."""
+    if "attributes" not in edge:
+        return False
+    for attribute in edge["attributes"]:
+        if attribute["attribute_type_id"] == "biolink:support_graphs":
+            return True
+    return False
+
+
 def validate_message(message, logger):
     """Validate a given message for missing nodes."""
     valid = True
@@ -231,7 +241,7 @@ def make_hashable(d):
     return result
 
 
-def merge_kgraph(og_message, new_message, logger: logging.Logger):
+def merge_kgraph(og_message, new_message, source, logger: logging.Logger):
     """Merge two TRAPI kgraphs together."""
     merged_kgraph = copy.deepcopy(og_message)
     for key, value in new_message["nodes"].items():
@@ -299,4 +309,18 @@ def merge_kgraph(og_message, new_message, logger: logging.Logger):
         else:
             merged_kgraph["edges"][key] = value
 
+            if value["sources"] and not is_support_edge(value):
+                new_sources = combine_unique_dicts(
+                    value["sources"],
+                    [
+                        {
+                            "resource_id": source,
+                            "resource_role": "aggregator_knowledge_source",
+                            "upstream_resource_ids": ["infores:retriever"],
+                        }
+                    ],
+                    logger,
+                )
+                merged_kgraph["edges"][key]["sources"] = new_sources
+
     return merged_kgraph

tests/unit/test_merge_message.py

@@ -26,9 +26,9 @@ async def test_message_merge(redis_mock, mocker):
 
     merged_message = merge_messages(
         "test_ara",
-        lookup_query_graph,
         original_query_graph,
-        result_messages=[original_response, callback_response],
+        original_response,
+        callback_response,
         logger=logger,
     )
 

workers/aragorn_lookup/worker.py

@@ -128,10 +128,6 @@ async def aragorn_lookup(task, logger: logging.Logger):
         # Put callback UID and query ID in postgres
         callback_id = str(uuid.uuid4())[:8]
         await add_callback_id(query_id, callback_id, logger)
-        # put lookup query graph in redis
-        await save_message(
-            f"{query_id}_lookup_query_graph", message["message"]["query_graph"], logger
-        )
         message["callback"] = f"{settings.callback_host}/aragorn/callback/{callback_id}"
 
         async with httpx.AsyncClient(timeout=100) as client:
@@ -142,12 +138,6 @@ async def aragorn_lookup(task, logger: logging.Logger):
     else:
         expanded_messages = expand_aragorn_query(message)
         requests = []
-        # put lookup query graph in redis
-        await save_message(
-            f"{query_id}_lookup_query_graph",
-            expanded_messages[0]["message"]["query_graph"],
-            logger,
-        )
         # send all messages to lookup service
         async with httpx.AsyncClient(timeout=20) as client:
             for expanded_message in expanded_messages:

workers/aragorn_pathfinder/worker.py

@@ -136,6 +136,8 @@ async def shadowfax(task, logger: logging.Logger):
     parameters["tiers"] = parameters.get("tiers") or [0]
     message["parameters"] = parameters
 
+    source = "infores:shepherd-aragorn"
+
     qgraph = message["message"]["query_graph"]
     pinned_node_keys = []
     pinned_node_ids = []
@@ -202,7 +204,9 @@ async def shadowfax(task, logger: logging.Logger):
     pairwise_pubs = get_the_pmids([source_node, target_node])
     if source_pubs == 0 or target_pubs == 0 or len(pairwise_pubs) == 0:
         logger.info("No publications found.")
-        return message, 200
+        await wrap_up_task(STREAM, GROUP, task, workflow, logger)
+        logger.info(f"Task took {time.time() - start}")
+        return
 
     # Find other nodes from those shared publications
     curies = set()
@@ -218,12 +222,16 @@ async def shadowfax(task, logger: logging.Logger):
 
     if len(curies) == 0:
         logger.info("No curies found.")
-        return message, 200
+        await wrap_up_task(STREAM, GROUP, task, workflow, logger)
+        logger.info(f"Task took {time.time() - start}")
+        return
 
     normalizer_response = await get_normalized_curies(list(curies), logger)
     if normalizer_response is None:
         logger.error("Failed to get a good response from Node Normalizer")
-        return message, 500
+        await wrap_up_task(STREAM, GROUP, task, workflow, logger)
+        logger.info(f"Task took {time.time() - start}")
+        return
 
     curie_info = defaultdict(dict)
     for curie, normalizer_info in normalizer_response.items():
@@ -306,7 +314,7 @@ async def shadowfax(task, logger: logging.Logger):
         try:
             logger.debug(f"Got back {len(lookup_message.get('results', 0))} results.")
             merged_kgraph = merge_kgraph(
-                merged_kgraph, lookup_message["knowledge_graph"], logger
+                merged_kgraph, lookup_message["knowledge_graph"], source, logger
             )
             merged_aux_graphs.update(lookup_message["auxiliary_graphs"])
         except KeyError as e:
@@ -475,7 +483,7 @@ async def shadowfax(task, logger: logging.Logger):
             aux_edges_keys.append(aux_graph_key)
 
         analysis = {
-            "resource_id": "infores:aragorn",
+            "resource_id": source,
             "path_bindings": {
                 path_key: [{"id": aux_graph_key, "attributes": []}],
             },

workers/bte_lookup/worker.py

@@ -131,10 +131,6 @@ async def bte_lookup(task, logger: logging.Logger):
         # Put callback UID and query ID in postgres
         callback_id = str(uuid.uuid4())[:8]
         await add_callback_id(query_id, callback_id, logger)
-        # put lookup query graph in redis
-        await save_message(
-            f"{query_id}_lookup_query_graph", message["message"]["query_graph"], logger
-        )
         message["callback"] = f"{settings.callback_host}/bte/callback/{callback_id}"
 
         async with httpx.AsyncClient(timeout=100) as client:
@@ -146,16 +142,6 @@ async def bte_lookup(task, logger: logging.Logger):
         expanded_messages = expand_bte_query(message, logger)
         logger.info(f"Expanded to {len(expanded_messages)} messages")
         requests = []
-        if len(expanded_messages) == 0:
-            lookup_query_graph = message["message"]["query_graph"]
-        else:
-            lookup_query_graph = expanded_messages[0]["message"]["query_graph"]
-        # put lookup query graph in redis
-        await save_message(
-            f"{query_id}_lookup_query_graph",
-            lookup_query_graph,
-            logger,
-        )
         # send all messages to retriever
         async with httpx.AsyncClient(timeout=20) as client:
             for expanded_message in expanded_messages: