Fix launcher use of data mappers

Author: bruno-f-cruz

src/aind_behavior_vr_foraging/data_mappers/__init__.py

@@ -2,15 +2,6 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from aind_behavior_services.session import AindBehaviorSessionModel
-from aind_behavior_services.utils import model_from_json_file
-from clabe.apps import BonsaiApp, CurriculumSuggestion
-from git import Repo
-
-from aind_behavior_vr_foraging.data_contract.utils import calculate_consumed_water
-from aind_behavior_vr_foraging.rig import AindVrForagingRig
-from aind_behavior_vr_foraging.task_logic import AindVrForagingTaskLogic
-
 from ._rig import AindRigDataMapper
 from ._session import AindSessionDataMapper
 

src/aind_behavior_vr_foraging/data_mappers/_session.py

@@ -33,7 +33,7 @@ def __init__(
         self,
         data_path: os.PathLike,
         repo_path: os.PathLike,
-        curriculum_suggestion_path: Optional[os.PathLike] = None,
+        curriculum_suggestion: Optional[os.PathLike] | CurriculumSuggestion = None,
         session_end_time: Optional[AwareDatetime] = None,
     ):
         self._data_path = data_path
@@ -45,8 +45,11 @@ def __init__(
         self.rig_model = model_from_json_file(abs_schemas_path / "rig_input.json", AindVrForagingRig)
         self.task_logic_model = model_from_json_file(abs_schemas_path / "tasklogic_input.json", AindVrForagingTaskLogic)
 
-        if curriculum_suggestion_path is not None:
-            curriculum_suggestion = model_from_json_file(Path(curriculum_suggestion_path), CurriculumSuggestion)
+        if curriculum_suggestion is not None:
+            if isinstance(curriculum_suggestion, CurriculumSuggestion):
+                pass
+            else:
+                curriculum_suggestion = model_from_json_file(Path(curriculum_suggestion), CurriculumSuggestion)
         else:
             try:
                 curriculum_suggestion = model_from_json_file(

src/aind_behavior_vr_foraging/launcher.py

@@ -18,7 +18,6 @@
 from contraqctor.contract.json import SoftwareEvents
 from pydantic_settings import CliApp
 
-import aind_behavior_vr_foraging.data_contract.utils
 from aind_behavior_vr_foraging import data_contract
 from aind_behavior_vr_foraging.data_mappers import AindRigDataMapper, AindSessionDataMapper
 from aind_behavior_vr_foraging.rig import AindVrForagingRig
@@ -96,25 +95,15 @@ async def experiment(launcher: Launcher) -> None:
         # Push updated trainer state back to the database
         picker.push_new_suggestion(suggestion.trainer_state)
 
-    try:
-        total_water_consumed = aind_behavior_vr_foraging.data_contract.utils.calculate_consumed_water(
-            launcher.session_directory
-        )
-    except Exception as e:
-        logger.warning(f"Could not calculate consumed water: {e}")
-        total_water_consumed = None
-
     # Mappers
+    assert launcher.repository.working_tree_dir is not None
     ads_session = AindSessionDataMapper(
-        rig=rig,
-        session=session,
-        task_logic=task_logic,
+        data_path=launcher.session_directory,
+        repo_path=launcher.repository.working_tree_dir,  # type: ignore[arg-type]
         curriculum_suggestion=suggestion,
-        bonsai_app=bonsai_app,
-        water_consumed_ml=total_water_consumed,
     ).map()
     ads_session.write_standard_file(launcher.session_directory)
-    ads_rig = AindRigDataMapper(rig=rig).map()
+    ads_rig = AindRigDataMapper(data_path=launcher.session_directory).map()
     ads_rig.write_standard_file(launcher.session_directory)
 
     # Run data qc