fix: handle KeyError for missing ontology terms in CellGuide pipeline (#7685)

Author: kaloster

backend/cellguide/pipeline/metadata/metadata_generator.py

@@ -1,7 +1,12 @@
 import logging
 
 from backend.cellguide.pipeline.metadata.types import CellMetadata, TissueMetadata
-from backend.common.census_cube.utils import ontology_parser
+from backend.common.census_cube.data.ontology_labels import (
+    is_ontology_term_deprecated,
+    ontology_term_description,
+    ontology_term_label,
+    ontology_term_synonyms,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -27,20 +32,20 @@ def generate_cellguide_card_metadata(all_cell_type_ids_in_corpus: list[str]) ->
 
     for id in all_cell_type_ids_in_corpus:
 
-        if ontology_parser.is_term_deprecated(id):
+        if is_ontology_term_deprecated(id):
             obsolete_cell_ids.append(id)
         else:
-            description = ontology_parser.get_term_description(id)
+            description = ontology_term_description(id)
             if description is not None:
                 cell_ids_with_cl_description += 1
             else:
                 cell_ids_without_cl_description += 1
 
             metadata = CellMetadata(
-                name=ontology_parser.get_term_label(id),
+                name=ontology_term_label(id),
                 id=id,
                 clDescription=description,
-                synonyms=ontology_parser.get_term_synonyms(id),
+                synonyms=ontology_term_synonyms(id),
             )
             cellguide_card_metadata[id] = metadata
 
@@ -71,20 +76,20 @@ def generate_cellguide_tissue_card_metadata(all_tissue_ids_in_corpus: list[str])
     uberon_ids_without_description = 0
 
     for id in all_tissue_ids_in_corpus:
-        if ontology_parser.is_term_deprecated(id):
+        if is_ontology_term_deprecated(id):
             obsolete_uberon_ids.append(id)
         else:
-            description = ontology_parser.get_term_description(id)
+            description = ontology_term_description(id)
             if description is not None:
                 uberon_ids_with_description += 1
             else:
                 uberon_ids_without_description += 1
 
             metadata = TissueMetadata(
-                name=ontology_parser.get_term_label(id),
+                name=ontology_term_label(id),
                 id=id,
                 uberonDescription=description,
-                synonyms=ontology_parser.get_term_synonyms(id),
+                synonyms=ontology_term_synonyms(id),
             )
             cellguide_tissue_card_metadata[id] = metadata
 

backend/cellguide/pipeline/ontology_tree/tree_builder.py

@@ -9,6 +9,7 @@
 
 from backend.cellguide.pipeline.constants import CELLGUIDE_PIPELINE_NUM_CPUS
 from backend.cellguide.pipeline.ontology_tree.types import OntologyTree, OntologyTreeState
+from backend.common.census_cube.data.ontology_labels import ontology_term_label
 from backend.common.census_cube.utils import ontology_parser, rollup_across_cell_type_descendants, to_dict
 
 logger = logging.getLogger(__name__)
@@ -98,7 +99,7 @@ def __init__(self, cell_counts_df, root_node="CL:0000000"):
         self.all_cell_type_ids_to_labels_in_corpus = dict(
             zip(
                 self.all_cell_type_ids_in_corpus,
-                [self.ontology.get_term_label(c) for c in self.all_cell_type_ids_in_corpus],
+                [ontology_term_label(c) for c in self.all_cell_type_ids_in_corpus],
                 strict=False,
             )
         )
@@ -390,7 +391,7 @@ def _process_tissue__parallel(self, tissueId: str) -> OntologyTreeState:
                     The number of cells is a dictionary with keys "n_cells" and "n_cells_rollup".
         """
         end_nodes = self.uberon_by_celltype[tissueId]
-        tissue_label = self.ontology.get_term_label(tissueId)
+        tissue_label = ontology_term_label(tissueId)
         uberon_ancestors = self.ontology.get_term_ancestors(tissueId, include_self=True)
 
         # filter out hemaotoietic cell types from non-whitelisted tissues

backend/cellguide/pipeline/source_collections/source_collections_generator.py

@@ -2,7 +2,8 @@
 from pandas import DataFrame
 
 from backend.cellguide.pipeline.source_collections.types import SourceCollectionsData
-from backend.common.census_cube.utils import descendants, ontology_parser
+from backend.common.census_cube.data.ontology_labels import ontology_term_label
+from backend.common.census_cube.utils import descendants
 
 
 def generate_source_collections_data(
@@ -40,7 +41,7 @@ def generate_source_collections_data(
             # We need tissue, disease, and organism labels AND ontology term ids for each cell type id
             df_dict = {
                 df_agg.index[i]: [
-                    {"label": ontology_parser.get_term_label(cell_type_id), "ontology_term_id": cell_type_id}
+                    {"label": ontology_term_label(cell_type_id), "ontology_term_id": cell_type_id}
                     for cell_type_id in df_agg.values[i][0].split(",")
                 ]
                 for i in range(len(df_agg))

backend/common/census_cube/data/ontology_labels.py

@@ -20,12 +20,46 @@ def ontology_term_label(ontology_term_id: str) -> Optional[str]:
     """
     try:
         return ontology_parser.get_term_label(ontology_term_id)
-    # If the ontology term id is invalid, return the ontology term id itself
-    # This is useful for cases like publication citation.
-    except ValueError:
+    # If the ontology term id is invalid or not found in the ontology schema,
+    # return the ontology term id itself. This is useful for cases like publication
+    # citation and newly added cell types not yet in the ontology schema.
+    except (ValueError, KeyError):
         return ontology_term_id
 
 
+def is_ontology_term_deprecated(ontology_term_id: str) -> bool:
+    """
+    Returns whether an ontology term is deprecated. Returns False for unknown terms.
+    """
+    try:
+        return ontology_parser.is_term_deprecated(ontology_term_id)
+    except (ValueError, KeyError):
+        # Assume unknown terms are not deprecated
+        return False
+
+
+def ontology_term_description(ontology_term_id: str) -> Optional[str]:
+    """
+    Returns the description for an ontology term, given its id. Returns None if the term is not found.
+    """
+    try:
+        return ontology_parser.get_term_description(ontology_term_id)
+    except (ValueError, KeyError):
+        # Return None for unknown terms (no description available)
+        return None
+
+
+def ontology_term_synonyms(ontology_term_id: str) -> list[str]:
+    """
+    Returns the synonyms for an ontology term, given its id. Returns empty list if the term is not found.
+    """
+    try:
+        return ontology_parser.get_term_synonyms(ontology_term_id)
+    except (ValueError, KeyError):
+        # Return empty list for unknown terms (no synonyms available)
+        return []
+
+
 def gene_term_label(gene_ontology_term_id: str) -> Optional[str]:
     """
     Returns the label for a gene ontology term, given its id. Return None if ontology term id is invalid.

tests/unit/wmg_processing/test_ontology_labels.py

@@ -1,6 +1,13 @@
 import unittest
+from unittest.mock import patch
 
-from backend.common.census_cube.data.ontology_labels import gene_term_label, ontology_term_label
+from backend.common.census_cube.data.ontology_labels import (
+    gene_term_label,
+    is_ontology_term_deprecated,
+    ontology_term_description,
+    ontology_term_label,
+    ontology_term_synonyms,
+)
 
 
 class OntologyLabelTests(unittest.TestCase):
@@ -36,6 +43,97 @@ def test__gene_label(self):
             with self.subTest(gene_id):
                 self.assertEqual(gene_term_label(gene_id), expected_gene_label)
 
+    @patch("backend.common.census_cube.data.ontology_labels.ontology_parser")
+    def test_ontology_term_label_handles_key_error(self, mock_ontology_parser):
+        """Test that ontology_term_label handles KeyError for missing terms in ontology."""
+        # Simulate the error that occurred in production with CL:4033085
+        mock_ontology_parser.get_term_label.side_effect = KeyError("CL:4033085")
+
+        result = ontology_term_label("CL:4033085")
+
+        # Should return the term ID itself as a fallback instead of crashing
+        self.assertEqual(result, "CL:4033085")
+
+    @patch("backend.common.census_cube.data.ontology_labels.ontology_parser")
+    def test_ontology_term_label_handles_value_error(self, mock_ontology_parser):
+        """Test that ontology_term_label handles ValueError gracefully."""
+        mock_ontology_parser.get_term_label.side_effect = ValueError("Invalid term")
+
+        result = ontology_term_label("CL:INVALID")
+
+        # Should return the term ID itself as a fallback
+        self.assertEqual(result, "CL:INVALID")
+
+    @patch("backend.common.census_cube.data.ontology_labels.ontology_parser")
+    def test_ontology_term_label_success_path(self, mock_ontology_parser):
+        """Test that ontology_term_label returns the correct label when term exists."""
+        mock_ontology_parser.get_term_label.return_value = "native cell"
+
+        result = ontology_term_label("CL:0000003")
+
+        self.assertEqual(result, "native cell")
+        mock_ontology_parser.get_term_label.assert_called_once_with("CL:0000003")
+
+    @patch("backend.common.census_cube.data.ontology_labels.ontology_parser")
+    def test_is_ontology_term_deprecated_handles_key_error(self, mock_ontology_parser):
+        """Test that is_ontology_term_deprecated handles KeyError for missing terms."""
+        mock_ontology_parser.is_term_deprecated.side_effect = KeyError("CL:4033085")
+
+        result = is_ontology_term_deprecated("CL:4033085")
+
+        # Should return False for unknown terms (assume not deprecated)
+        self.assertEqual(result, False)
+
+    @patch("backend.common.census_cube.data.ontology_labels.ontology_parser")
+    def test_is_ontology_term_deprecated_success_path(self, mock_ontology_parser):
+        """Test that is_ontology_term_deprecated returns the correct value when term exists."""
+        mock_ontology_parser.is_term_deprecated.return_value = True
+
+        result = is_ontology_term_deprecated("CL:0000001")
+
+        self.assertEqual(result, True)
+        mock_ontology_parser.is_term_deprecated.assert_called_once_with("CL:0000001")
+
+    @patch("backend.common.census_cube.data.ontology_labels.ontology_parser")
+    def test_ontology_term_description_handles_key_error(self, mock_ontology_parser):
+        """Test that ontology_term_description handles KeyError for missing terms."""
+        mock_ontology_parser.get_term_description.side_effect = KeyError("CL:4033085")
+
+        result = ontology_term_description("CL:4033085")
+
+        # Should return None for unknown terms (no description available)
+        self.assertIsNone(result)
+
+    @patch("backend.common.census_cube.data.ontology_labels.ontology_parser")
+    def test_ontology_term_description_success_path(self, mock_ontology_parser):
+        """Test that ontology_term_description returns the correct description when term exists."""
+        mock_ontology_parser.get_term_description.return_value = "A cell description"
+
+        result = ontology_term_description("CL:0000003")
+
+        self.assertEqual(result, "A cell description")
+        mock_ontology_parser.get_term_description.assert_called_once_with("CL:0000003")
+
+    @patch("backend.common.census_cube.data.ontology_labels.ontology_parser")
+    def test_ontology_term_synonyms_handles_key_error(self, mock_ontology_parser):
+        """Test that ontology_term_synonyms handles KeyError for missing terms."""
+        mock_ontology_parser.get_term_synonyms.side_effect = KeyError("CL:4033085")
+
+        result = ontology_term_synonyms("CL:4033085")
+
+        # Should return empty list for unknown terms (no synonyms available)
+        self.assertEqual(result, [])
+
+    @patch("backend.common.census_cube.data.ontology_labels.ontology_parser")
+    def test_ontology_term_synonyms_success_path(self, mock_ontology_parser):
+        """Test that ontology_term_synonyms returns the correct synonyms when term exists."""
+        mock_ontology_parser.get_term_synonyms.return_value = ["synonym1", "synonym2"]
+
+        result = ontology_term_synonyms("CL:0000003")
+
+        self.assertEqual(result, ["synonym1", "synonym2"])
+        mock_ontology_parser.get_term_synonyms.assert_called_once_with("CL:0000003")
+
 
 if __name__ == "__main__":
     unittest.main()