Feat 188 integrate protocols (#220)

* protocols mapping and integration * workaround to avoid multiple smartsheet api calls * fix surgery protocol mapping, unit tests * linters * fixes linters, docstrings * upper bound for pydantic * upper bound 2.6.4 * switch order * pins pydantic version * feat: reverts changes to endpoint args --------- Co-authored-by: jtyoung84 <[email protected]>
AllenNeuralDynamics · Apr 19, 2024 · 72a2cdb · 72a2cdb
1 parent e68e3fd
commit 72a2cdb
Show file tree

Hide file tree

Showing 5 changed files with 355 additions and 8 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -18,7 +18,8 @@ dynamic = ["version"]
 
 dependencies = [
     'aind-metadata-mapper==0.6.1',
-    'pydantic-settings>=2.0'
+    'pydantic-settings>=2.0',
+    'pydantic<=2.6.4'
 ]
 
 [project.optional-dependencies]

diff --git a/src/aind_metadata_service/server.py b/src/aind_metadata_service/server.py
@@ -31,7 +31,10 @@
 from aind_metadata_service.smartsheet.perfusions.mapping import (
     PerfusionsMapper,
 )
-from aind_metadata_service.smartsheet.protocols.mapping import ProtocolsMapper
+from aind_metadata_service.smartsheet.protocols.mapping import (
+    ProtocolsIntegrator,
+    ProtocolsMapper,
+)
 from aind_metadata_service.tars.client import AzureSettings, TarsClient
 from aind_metadata_service.tars.mapping import TarsResponseHandler
 
@@ -133,7 +136,10 @@ async def retrieve_rig(rig_id, pickle: bool = False):
 
 
 @app.get("/protocols/{protocol_name}")
-async def retrieve_protocols(protocol_name, pickle: bool = False):
+async def retrieve_protocols(
+    protocol_name,
+    pickle: bool = False,
+):
     """Retrieves perfusion information from smartsheet"""
 
     # TODO: We can probably cache the response if it's 200
@@ -245,6 +251,7 @@ async def retrieve_procedures(subject_id, pickle: bool = False):
     merged_response = sharepoint_client.merge_responses(
         [lb_response, sp2019_response, sp2023_response]
     )
+    # integrate TARS response
     mapper = TarsResponseHandler()
     viruses = mapper.get_virus_strains(merged_response)
     tars_mapping = {}
@@ -256,6 +263,29 @@ async def retrieve_procedures(subject_id, pickle: bool = False):
     integrated_response = mapper.integrate_injection_materials(
         response=merged_response, tars_mapping=tars_mapping
     )
+    # integrate protocols from smartsheet
+    smart_sheet_response = await run_in_threadpool(
+        protocols_smart_sheet_client.get_sheet
+    )
+    protocols_integrator = ProtocolsIntegrator()
+    protocols_list = protocols_integrator.get_protocols_list(
+        integrated_response
+    )
+    protocols_mapping = {}
+    for protocol_name in protocols_list:
+        mapper = ProtocolsMapper(
+            smart_sheet_response=smart_sheet_response, input_id=protocol_name
+        )
+        model_response = mapper.get_model_response()
+        # smartsheet_response = await retrieve_protocols(
+        #     protocol_name=protocol_name
+        # )
+        protocols_mapping[protocol_name] = (
+            model_response.map_to_json_response()
+        )
+    integrated_response = protocols_integrator.integrate_protocols(
+        response=integrated_response, protocols_mapping=protocols_mapping
+    )
     if pickle:
         return integrated_response.map_to_pickled_response()
     else:

diff --git a/src/aind_metadata_service/smartsheet/protocols/mapping.py b/src/aind_metadata_service/smartsheet/protocols/mapping.py
@@ -1,9 +1,18 @@
 """Module that handles the methods to map the SmartSheet response to a
 protocol_id."""
 
+import json
 import logging
-from typing import List, Optional
+from typing import Dict, List, Optional
 
+from aind_data_schema.core.procedures import (
+    Craniotomy,
+    IontophoresisInjection,
+    NanojectInjection,
+    Perfusion,
+    ProtectiveMaterial,
+    Surgery,
+)
 from pydantic import ValidationError
 
 from aind_metadata_service.client import StatusCodes
@@ -12,6 +21,7 @@
 from aind_metadata_service.smartsheet.mapper import SmartSheetMapper
 from aind_metadata_service.smartsheet.models import SheetRow
 from aind_metadata_service.smartsheet.protocols.models import (
+    ProtocolNames,
     ProtocolsColumnNames,
 )
 
@@ -20,7 +30,7 @@ class ProtocolsMapper(SmartSheetMapper):
     """Primary class to handle mapping data models and returning a response"""
 
     def _map_row_to_protocol(
-        self, row: SheetRow, input_protocol_name: str
+        self, row: SheetRow, input_protocol_name: Optional[str]
     ) -> Optional[ProtocolInformation]:
         """
         Map a row to an optional funding model.
@@ -110,3 +120,102 @@ def _get_model_response(self) -> ModelResponse:
         except Exception as e:
             logging.error(repr(e))
             return ModelResponse.internal_server_error_response()
+
+
+class ProtocolsIntegrator:
+    """Methods to integrate Protocols into Procedures"""
+
+    @staticmethod
+    def _get_protocol_name(procedure):
+        """Gets protocol name based on procedure type"""
+        if isinstance(procedure, NanojectInjection):
+            return ProtocolNames.INJECTION_NANOJECT.value
+        elif isinstance(procedure, IontophoresisInjection):
+            return ProtocolNames.INJECTION_IONTOPHORESIS.value
+        elif isinstance(procedure, Perfusion):
+            return ProtocolNames.PERFUSION.value
+        elif isinstance(procedure, Craniotomy):
+            if procedure.protective_material == ProtectiveMaterial.DURAGEL:
+                return ProtocolNames.DURAGEL_APPLICATION.value
+        else:
+            return None
+
+    def get_protocols_list(self, response: ModelResponse) -> List:
+        """Creates a list of protocol names from procedures list"""
+        protocol_list = []
+        if len(response.aind_models) > 0:
+            procedures = response.aind_models[0]
+            for subject_procedure in procedures.subject_procedures:
+                if isinstance(subject_procedure, Surgery):
+                    protocol_list.append(ProtocolNames.SURGERY.value)
+                for procedure in subject_procedure.procedures:
+                    protocol_name = self._get_protocol_name(
+                        procedure=procedure
+                    )
+                    protocol_list.append(protocol_name)
+        return protocol_list
+
+    def integrate_protocols(
+        self, response: ModelResponse, protocols_mapping: Dict
+    ) -> ModelResponse:
+        """
+        Merges protocols responses with procedures response
+        Parameters
+        ----------
+        response: ModelResponse
+             Merged response from procedures endpoints
+        protocols_mapping: dict
+             Dictionary mapping protocol names to info from smartsheet
+        Returns
+        -------
+        Procedures response with protocols
+        """
+        output_aind_models = []
+        status_code = response.status_code
+        if len(response.aind_models) > 0:
+            pre_procedures = response.aind_models[0]
+            for subject_procedure in pre_procedures.subject_procedures:
+                if isinstance(subject_procedure, Surgery):
+                    protocol_name = ProtocolNames.SURGERY.value
+                    smartsheet_response = protocols_mapping.get(protocol_name)
+                    if (
+                        smartsheet_response.status_code
+                        == StatusCodes.DB_RESPONDED.value
+                        or smartsheet_response.status_code
+                        == StatusCodes.VALID_DATA.value
+                        or smartsheet_response.status_code
+                        == StatusCodes.INVALID_DATA.value
+                    ):
+                        data = json.loads(smartsheet_response.body)["data"]
+                        subject_procedure.protocol_id = data["doi"]
+                    elif (
+                        smartsheet_response.status_code
+                        == StatusCodes.NO_DATA_FOUND.value
+                    ):
+                        pass
+                    else:
+                        status_code = StatusCodes.MULTI_STATUS
+                for procedure in subject_procedure.procedures:
+                    protocol_name = self._get_protocol_name(procedure)
+                    smartsheet_response = protocols_mapping.get(protocol_name)
+                    if (
+                        smartsheet_response.status_code
+                        == StatusCodes.DB_RESPONDED.value
+                        or smartsheet_response.status_code
+                        == StatusCodes.VALID_DATA.value
+                        or smartsheet_response.status_code
+                        == StatusCodes.INVALID_DATA.value
+                    ):
+                        data = json.loads(smartsheet_response.body)["data"]
+                        procedure.protocol_id = data["doi"]
+                    elif (
+                        smartsheet_response.status_code
+                        == StatusCodes.NO_DATA_FOUND.value
+                    ):
+                        pass
+                    else:
+                        status_code = StatusCodes.MULTI_STATUS
+                output_aind_models = [pre_procedures]
+        return ModelResponse(
+            aind_models=output_aind_models, status_code=status_code
+        )
diff --git a/src/aind_metadata_service/smartsheet/protocols/models.py b/src/aind_metadata_service/smartsheet/protocols/models.py
@@ -13,3 +13,37 @@ class ProtocolsColumnNames(str, Enum):
     VERSION = "Version"
     PROTOCOL_COLLECTION = "Protocol collection"
     PRIORITY = "Priority "
+
+
+class ProtocolNames(Enum):
+    """Enum of Protocol Names in Smartsheet"""
+
+    IMMUNOLABELING = "Immunolabeling of a Whole Mouse Brain"
+    DELIPIDATION = (
+        "Tetrahydrofuran and Dichloromethane Delipidation of a"
+        " Whole Mouse Brain"
+    )
+    SBIP_DELIPADATION = "Aqueous (SBiP) Delipidation of a Whole Mouse Brain"
+    GELATIN_PREVIOUS = (
+        "Whole Mouse Brain Delipidation, Immunolabeling,"
+        " and Expansion Microscopy"
+    )
+    INJECTION_NANOJECT = "Injection of Viral Tracers by Nanoject V.4"
+    INJECTION_IONTOPHORESIS = (
+        "Stereotaxic Surgery for Delivery of Tracers by Iontophoresis V.3"
+    )
+    PERFUSION = "Mouse Cardiac Perfusion Fixation and Brain Collection V.5"
+    SMARTSPIM_IMAGING = "Imaging cleared mouse brains on SmartSPIM"
+    SMARTSPIM_SETUP = "SmartSPIM setup and alignment"
+    SURGERY = "General Set-Up and Take-Down for Rodent Neurosurgery"
+    PROTOCOL_COLLECTION = (
+        "Protocol Collection: Perfusing, Sectioning, IHC,"
+        " Mounting and Coverslipping Mouse Brain Specimens"
+    )
+    SECTIONING = "Sectioning Mouse Brain with Sliding Microtome"
+    MOUNTING_COVERSLIPPING = "Mounting and Coverslipping Mouse Brain Sections"
+    IHC_SECTIONS = "Immunohistochemistry (IHC) Staining Mouse Brain Sections"
+    DAPI_STAINING = "DAPI Staining Mouse Brain Sections"
+    DURAGEL_APPLICATION = (
+        "Duragel application for acute electrophysiological recordings"
+    )