feat!: add support for transforming CIViC prognostic evidence (#407)

closes #242

Previously only Therapeutic Response evidence items were supported. This
adds support for Prognostic evidence.

src/metakb/load_data.py

@@ -376,6 +376,7 @@ def _add_obj_id_to_set(obj: dict, ids_set: set[str]) -> None:
             statement.get("reportedIn"),
             statement.get("subjectVariant"),
             statement.get("objectTherapeutic"),
+            statement.get("objectCondition"),
             statement.get("conditionQualifier"),
             statement.get("geneContextQualifier"),
         ]:
@@ -444,11 +445,14 @@ def _add_statement(tx: ManagedTransaction, statement_in: dict) -> None:
     match_line += f"MERGE (v:Variation {{ id: '{variant_id}' }})\n"
     rel_line += "MERGE (s) -[:HAS_VARIANT] -> (v)\n"
 
-    therapeutic_id = statement["objectTherapeutic"]["id"]
-    match_line += f"MERGE (t:TherapeuticProcedure {{ id: '{therapeutic_id}' }})\n"
-    rel_line += "MERGE (s) -[:HAS_THERAPEUTIC] -> (t)\n"
+    therapeutic = statement.get("objectTherapeutic")
+    if therapeutic:
+        therapeutic_id = therapeutic["id"]
+        match_line += f"MERGE (t:TherapeuticProcedure {{ id: '{therapeutic_id}' }})\n"
+        rel_line += "MERGE (s) -[:HAS_THERAPEUTIC] -> (t)\n"
 
-    tumor_type_id = statement["conditionQualifier"]["id"]
+    tumor_type = statement.get("conditionQualifier") or statement.get("objectCondition")
+    tumor_type_id = tumor_type["id"]
     match_line += f"MERGE (tt:Condition {{ id: '{tumor_type_id}' }})\n"
     rel_line += "MERGE (s) -[:HAS_TUMOR_TYPE] -> (tt)\n"
 

src/metakb/query.py

@@ -15,6 +15,7 @@
 )
 from ga4gh.core.entity_models import Coding, Document, Extension, Method
 from ga4gh.va_spec.profiles.var_study_stmt import (
+    VariantPrognosticStudyStatement,
     VariantTherapeuticResponseStudyStatement,
 )
 from ga4gh.vrs.models import Expression, Variation
@@ -492,16 +493,28 @@ def _get_nested_stmts(self, statement_nodes: list[Node]) -> list[dict]:
 
     def _get_nested_stmt(self, stmt_node: Node) -> dict:
         """Get information related to a statement
-        Only VariantTherapeuticResponseStudyStatement are supported at the moment
+        Only VariantTherapeuticResponseStudyStatement and VariantPrognosticStudyStatement
+        are supported at the moment
 
         :param stmt_node: Neo4j graph node for statement
         :return: Nested statement
         """
-        if stmt_node["type"] != "VariantTherapeuticResponseStudyStatement":
+        study_stmt_type = stmt_node["type"]
+        if study_stmt_type not in {
+            "VariantTherapeuticResponseStudyStatement",
+            "VariantPrognosticStudyStatement",
+        }:
             return {}
 
+        if study_stmt_type == "VariantPrognosticStudyStatement":
+            study_stmt_cls = VariantPrognosticStudyStatement
+            condition_key = "objectCondition"
+        else:
+            study_stmt_cls = VariantTherapeuticResponseStudyStatement
+            condition_key = "conditionQualifier"
+
         params = {
-            "conditionQualifier": None,
+            condition_key: None,
             "subjectVariant": None,
             "strength": None,
             "reportedIn": [],
@@ -526,7 +539,7 @@ def _get_nested_stmt(self, stmt_node: Node) -> dict:
             node = data["n"]
 
             if rel_type == "HAS_TUMOR_TYPE":
-                params["conditionQualifier"] = self._get_disease(node)
+                params[condition_key] = self._get_disease(node)
             elif rel_type == "HAS_VARIANT":
                 params["subjectVariant"] = self._get_cat_var(node)
             elif rel_type == "HAS_GENE_CONTEXT":
@@ -546,7 +559,7 @@ def _get_nested_stmt(self, stmt_node: Node) -> dict:
             else:
                 logger.warning("relation type not supported: %s", rel_type)
 
-        return VariantTherapeuticResponseStudyStatement(**params).model_dump()
+        return study_stmt_cls(**params).model_dump()
 
     @staticmethod
     def _get_vicc_normalizer_extension(node: dict) -> ViccNormalizerDataExtension:
@@ -905,6 +918,9 @@ async def batch_search_statements(
         response.statement_ids = [n["id"] for n in statement_nodes]
         stmts = self._get_nested_stmts(statement_nodes)
         response.statements = [
-            VariantTherapeuticResponseStudyStatement(**s) for s in stmts
+            VariantTherapeuticResponseStudyStatement(**s)
+            if s["type"] == "VariantTherapeuticResponseStudyStatement"
+            else VariantPrognosticStudyStatement(**s)
+            for s in stmts
         ]
         return response

src/metakb/schemas/api.py

@@ -3,6 +3,7 @@
 from typing import Literal
 
 from ga4gh.va_spec.profiles.var_study_stmt import (
+    VariantPrognosticStudyStatement,
     VariantTherapeuticResponseStudyStatement,
 )
 from pydantic import BaseModel, ConfigDict, StrictStr
@@ -46,7 +47,9 @@ class SearchStatementsService(BaseModel):
     query: SearchStatementsQuery
     warnings: list[StrictStr] = []
     statement_ids: list[StrictStr] = []
-    statements: list[VariantTherapeuticResponseStudyStatement] = []
+    statements: list[
+        VariantTherapeuticResponseStudyStatement | VariantPrognosticStudyStatement
+    ] = []
     service_meta_: ServiceMeta
 
 
@@ -69,5 +72,7 @@ class BatchSearchStatementsService(BaseModel):
     query: BatchSearchStatementsQuery
     warnings: list[StrictStr] = []
     statement_ids: list[StrictStr] = []
-    statements: list[VariantTherapeuticResponseStudyStatement] = []
+    statements: list[
+        VariantTherapeuticResponseStudyStatement | VariantPrognosticStudyStatement
+    ] = []
     service_meta_: ServiceMeta

src/metakb/transformers/base.py

@@ -25,6 +25,7 @@
 )
 from ga4gh.core.entity_models import Coding, Document, Extension, Method
 from ga4gh.va_spec.profiles.var_study_stmt import (
+    VariantPrognosticStudyStatement,
     VariantTherapeuticResponseStudyStatement,
 )
 from ga4gh.vrs.models import Allele
@@ -109,7 +110,9 @@ class ViccConceptVocab(BaseModel):
 class TransformedData(BaseModel):
     """Define model for transformed data"""
 
-    statements: list[VariantTherapeuticResponseStudyStatement] = []
+    statements: list[
+        VariantTherapeuticResponseStudyStatement | VariantPrognosticStudyStatement
+    ] = []
     categorical_variants: list[CategoricalVariant] = []
     variations: list[Allele] = []
     genes: list[Gene] = []