diff --git a/code/ARAX/ARAXQuery/ARAX_expander.py b/code/ARAX/ARAXQuery/ARAX_expander.py index 197831546..4ecc59b97 100644 --- a/code/ARAX/ARAXQuery/ARAX_expander.py +++ b/code/ARAX/ARAXQuery/ARAX_expander.py @@ -353,7 +353,8 @@ def apply(self, response, input_parameters, mode: str = "ARAX"): if inferred_qedge_keys and len(query_graph.edges) == 1: for edge in query_sub_graph.edges.keys(): query_sub_graph.edges[edge].knowledge_type = 'lookup' - # Expand the query graph edge-by-edge + + # Expand the query graph edge-by-edge (in regular 'lookup' fashion) for qedge_key in ordered_qedge_keys_to_expand: log.debug(f"Expanding qedge {qedge_key}") response.update_query_plan(qedge_key, 'edge_properties', 'status', 'Expanding') @@ -583,9 +584,6 @@ def apply(self, response, input_parameters, mode: str = "ARAX"): decorator = ARAXDecorator() decorator.decorate_nodes(response) decorator.decorate_edges(response, kind="RTX-KG2") - - # Override node types to only include descendants of what was asked for in the QG (where applicable) #1360 - self._override_node_categories(message.knowledge_graph, message.query_graph, log) elif mode == "RTXKG2": decorator = ARAXDecorator() decorator.decorate_edges(response, kind="SEMMEDDB") @@ -1400,31 +1398,6 @@ def _load_fda_approved_drug_ids() -> Set[str]: fda_approved_drug_ids = pickle.load(fda_pickle) return fda_approved_drug_ids - def _override_node_categories(self, kg: KnowledgeGraph, qg: QueryGraph, log: ARAXResponse): - # Clean up what we list as the TRAPI node.categories; list descendants of what was asked for in the QG - log.debug(f"Overriding node categories to better align with what's in the QG") - qnode_descendant_categories_map = {qnode_key: set(self.bh.get_descendants(qnode.categories)) - for qnode_key, qnode in qg.nodes.items() if qnode.categories} - for node_key, node in kg.nodes.items(): - final_categories = set() - for qnode_key in node.qnode_keys: - # If qnode has categories specified, use node's all_categories that are descendants of qnode categories - if qnode_key in qnode_descendant_categories_map: - all_categories_attributes = [attribute for attribute in eu.convert_to_list(node.attributes) - if attribute.attribute_type_id == "biolink:category"] - node_categories = all_categories_attributes[0].value if all_categories_attributes else node.categories - relevant_categories = set(node_categories).intersection(qnode_descendant_categories_map[qnode_key]) - # Otherwise just use what's already in the node's categories (for KG2 this is the 'preferred' category) - else: - relevant_categories = set(node.categories) - final_categories = final_categories.union(relevant_categories) - if final_categories: - node.categories = list(final_categories) - else: - # Leave categories as they are but issue a warning - log.warning(f"None of the categories KPs gave node {node_key} ({node.categories}) are descendants of " - f"those asked for in the QG (for qnode {node.qnode_keys})") - @staticmethod def _map_back_to_input_curies(kg: KnowledgeGraph, qg: QueryGraph, log: ARAXResponse): """ diff --git a/code/ARAX/ARAXQuery/Expand/kp_selector.py b/code/ARAX/ARAXQuery/Expand/kp_selector.py index 8c556d086..692520716 100644 --- a/code/ARAX/ARAXQuery/Expand/kp_selector.py +++ b/code/ARAX/ARAXQuery/Expand/kp_selector.py @@ -42,10 +42,7 @@ def _load_cached_kp_info(self) -> tuple: self.log.error(f"Failed to load KP info caches due to {e}", error_code="LoadKPCachesFailed") return None, None, None, None - # Record None URLs for our local KPs - allowed_kp_urls = smart_api_info["allowed_kp_urls"] - - return (meta_map, allowed_kp_urls, smart_api_info["kps_excluded_by_version"], + return (meta_map, smart_api_info["allowed_kp_urls"], smart_api_info["kps_excluded_by_version"], smart_api_info["kps_excluded_by_maturity"]) def get_kps_for_single_hop_qg(self, qg: QueryGraph) -> Optional[Set[str]]: