Merge pull request #231 from AllenNeuralDynamics/refactor-against-aind-behavior-services-breaking

Refactor against version 0.13 of aind-behavior-services

Author: bruno-f-cruz

examples/_mocks.py

@@ -5,9 +5,9 @@
 
 import git
 from aind_behavior_curriculum import Stage, TrainerState
-from aind_behavior_services.rig import AindBehaviorRigModel
-from aind_behavior_services.session import AindBehaviorSessionModel
-from aind_behavior_services.task_logic import AindBehaviorTaskLogicModel
+from aind_behavior_services.rig import Rig
+from aind_behavior_services.session import Session
+from aind_behavior_services.task import Task
 from pydantic import Field
 
 from clabe.data_mapper import DataMapper
@@ -19,20 +19,20 @@
 
 
 ### Task-specific definitions
-class RigModel(AindBehaviorRigModel):
+class RigModel(Rig):
     rig_name: str = Field(default="TestRig", description="Rig name")
     version: Literal["0.0.0"] = "0.0.0"
 
 
-class TaskLogicModel(AindBehaviorTaskLogicModel):
+class MockTask(Task):
     version: Literal["0.0.0"] = "0.0.0"
     name: Literal[TASK_NAME] = TASK_NAME
 
 
 mock_trainer_state = TrainerState[Any](
     curriculum=None,
     is_on_curriculum=False,
-    stage=Stage(name="TestStage", task=TaskLogicModel(name=TASK_NAME, task_parameters={"foo": "bar"})),
+    stage=Stage(name="TestStage", task=MockTask(name=TASK_NAME, task_parameters={"foo": "bar"})),
 )
 
 
@@ -55,8 +55,8 @@ class DemoAindDataSchemaSessionDataMapper(DataMapper[MockAindDataSchemaSession])
     def __init__(
         self,
         rig_model: RigModel,
-        session_model: AindBehaviorSessionModel,
-        task_logic_model: TaskLogicModel,
+        session_model: Session,
+        task_model: MockTask,
         repository: Union[os.PathLike, git.Repo],
         script_path: os.PathLike,
         session_end_time: Optional[datetime.datetime] = None,
@@ -65,7 +65,7 @@ def __init__(
         super().__init__()
         self.session_model = session_model
         self.rig_model = rig_model
-        self.task_logic_model = task_logic_model
+        self.task_model = task_model
         self.repository = repository
         self.script_path = script_path
         self.session_end_time = session_end_time
@@ -74,7 +74,7 @@ def __init__(
 
     def map(self) -> MockAindDataSchemaSession:
         self._mapped = MockAindDataSchemaSession(
-            computer_name=self.rig_model.computer_name, repository=self.repository, task_name=self.task_logic_model.name
+            computer_name=self.rig_model.computer_name, repository=self.repository, task_name=self.task_model.name
         )
         print("#" * 50)
         print("THIS IS MAPPED DATA!")
@@ -87,8 +87,8 @@ def create_fake_subjects():
     subjects = ["00000", "123456"]
     for subject in subjects:
         os.makedirs(f"{LIB_CONFIG}/Subjects/{subject}", exist_ok=True)
-        with open(f"{LIB_CONFIG}/Subjects/{subject}/task_logic.json", "w", encoding="utf-8") as f:
-            f.write(TaskLogicModel(task_parameters={"subject": subject}).model_dump_json(indent=2))
+        with open(f"{LIB_CONFIG}/Subjects/{subject}/task.json", "w", encoding="utf-8") as f:
+            f.write(MockTask(task_parameters={"subject": subject}).model_dump_json(indent=2))
         with open(f"{LIB_CONFIG}/Subjects/{subject}/trainer_state.json", "w", encoding="utf-8") as f:
             f.write(mock_trainer_state.model_dump_json(indent=2))
 
@@ -97,4 +97,4 @@ def create_fake_rig():
     computer_name = os.getenv("COMPUTERNAME")
     os.makedirs(_dir := f"{LIB_CONFIG}/Rig/{computer_name}", exist_ok=True)
     with open(f"{_dir}/rig1.json", "w", encoding="utf-8") as f:
-        f.write(RigModel(data_directory=r"./local/data").model_dump_json(indent=2))
+        f.write(RigModel(data_directory=r"./local/data", computer_name="mock_pc").model_dump_json(indent=2))

examples/behavior_launcher.py

@@ -4,10 +4,10 @@
 
 from _mocks import (
     LIB_CONFIG,
-    AindBehaviorSessionModel,
     DemoAindDataSchemaSessionDataMapper,
+    MockTask,
     RigModel,
-    TaskLogicModel,
+    Session,
     create_fake_rig,
     create_fake_subjects,
 )
@@ -30,10 +30,10 @@ async def demo_experiment(launcher: Launcher) -> None:
         experimenter_validator=lambda _: True,
     )
 
-    session = picker.pick_session(AindBehaviorSessionModel)
+    session = picker.pick_session(Session)
     rig = picker.pick_rig(RigModel)
     launcher.register_session(session, rig.data_directory)
-    trainer_state, task_logic = picker.pick_trainer_state(TaskLogicModel)
+    trainer_state, task = picker.pick_trainer_state(MockTask)
     _temp_trainer_state_path = launcher.save_temp_model(trainer_state)
 
     resource_monitor.ResourceMonitor(
@@ -62,7 +62,7 @@ def fmt(value: str) -> str:
     DemoAindDataSchemaSessionDataMapper(
         rig,
         session,
-        task_logic,
+        task,
         repository=launcher.repository,
         script_path=Path("./mock/script.py"),
         output_parameters={"suggestion": suggestion.model_dump()},

examples/client_behavior_launcher.py

@@ -3,9 +3,9 @@
 
 from _mocks import (
     LIB_CONFIG,
-    AindBehaviorSessionModel,
+    MockTask,
     RigModel,
-    TaskLogicModel,
+    Session,
     create_fake_rig,
     create_fake_subjects,
 )
@@ -29,10 +29,10 @@ async def client_experiment(launcher: Launcher) -> None:
         experimenter_validator=lambda _: True,
     )
 
-    session = picker.pick_session(AindBehaviorSessionModel)
+    session = picker.pick_session(Session)
     rig = picker.pick_rig(RigModel)
     launcher.register_session(session, rig.data_directory)
-    trainer_state, task_logic = picker.pick_trainer_state(TaskLogicModel)
+    trainer_state, task = picker.pick_trainer_state(MockTask)
 
     resource_monitor.ResourceMonitor(
         constrains=[
@@ -45,10 +45,10 @@ async def client_experiment(launcher: Launcher) -> None:
     bonsai_root = Path(r"C:\git\AllenNeuralDynamics\Aind.Behavior.VrForaging")
     session_response = xml_rpc_client.upload_model(session, "session.json")
     rig_response = xml_rpc_client.upload_model(rig, "rig.json")
-    task_logic_response = xml_rpc_client.upload_model(task_logic, "task_logic.json")
+    task_response = xml_rpc_client.upload_model(task, "task.json")
     assert rig_response.path is not None
     assert session_response.path is not None
-    assert task_logic_response.path is not None
+    assert task_response.path is not None
 
     bonsai_app_result = await xml_rpc_client.run_async(
         BonsaiApp(
@@ -57,7 +57,7 @@ async def client_experiment(launcher: Launcher) -> None:
             additional_externalized_properties={
                 "RigPath": rig_response.path,
                 "SessionPath": session_response.path,
-                "TaskLogicPath": task_logic_response.path,
+                "TaskPath": task_response.path,
             },
         ).command
     )

pyproject.toml

@@ -26,8 +26,8 @@ dependencies = [
     "requests",
     "semver",
     "rich",
-    "aind_behavior_services >= 0.13rc0",
     "questionary",
+    "aind_behavior_services>=0.13.0",
 ]
 
 [project.urls]

src/clabe/_typing.py

@@ -1,7 +1,7 @@
 from typing import TypeVar
 
-from aind_behavior_services import AindBehaviorRigModel, AindBehaviorSessionModel, AindBehaviorTaskLogicModel
+from aind_behavior_services import Rig, Session, Task
 
-TSession = TypeVar("TSession", bound=AindBehaviorSessionModel)
-TRig = TypeVar("TRig", bound=AindBehaviorRigModel)
-TTaskLogic = TypeVar("TTaskLogic", bound=AindBehaviorTaskLogicModel)
+TSession = TypeVar("TSession", bound=Session)
+TRig = TypeVar("TRig", bound=Rig)
+TTask = TypeVar("TTask", bound=Task)

src/clabe/apps/_bonsai.py

@@ -7,7 +7,7 @@
 from typing import Dict, List, Optional
 
 import pydantic
-from aind_behavior_services import AindBehaviorRigModel, AindBehaviorSessionModel, AindBehaviorTaskLogicModel
+from aind_behavior_services import Rig, Session, Task
 
 from ..constants import TMP_DIR
 from ._base import Command, CommandResult, ExecutableApp, identity_parser
@@ -170,7 +170,7 @@ class AindBehaviorServicesBonsaiApp(BonsaiApp):
     Specialized Bonsai application for AIND behavior services integration.
 
     This class extends the base BonsaiApp to provide specific functionality for
-    AIND behavior experiments, including automatic configuration of task logic,
+    AIND behavior experiments, including automatic configuration of task,
     session, and rig paths for the Bonsai workflow.
 
     Example:
@@ -186,15 +186,15 @@ def __init__(
         workflow: os.PathLike,
         *,
         temp_directory: Optional[os.PathLike] = None,
-        rig: Optional[AindBehaviorRigModel] = None,
-        session: Optional[AindBehaviorSessionModel] = None,
-        task_logic: Optional[AindBehaviorTaskLogicModel] = None,
+        rig: Optional[Rig] = None,
+        session: Optional[Session] = None,
+        task: Optional[Task] = None,
         **kwargs,
     ) -> None:
         """
         Initializes the AIND behavior services Bonsai app with automatic model configuration.
 
-        Automatically configures RigPath, SessionPath, and TaskLogicPath properties
+        Automatically configures RigPath, SessionPath, and TaskPath properties
         for the Bonsai workflow by saving provided models to temporary files and
         passing their paths as externalized properties.
 
@@ -204,38 +204,38 @@ def __init__(
             launcher: The launcher instance for saving temporary models
             rig: Optional rig model to configure. Defaults to None
             session: Optional session model to configure. Defaults to None
-            task_logic: Optional task logic model to configure. Defaults to None
+            task: Optional task model to configure. Defaults to None
             **kwargs: Additional keyword arguments passed to BonsaiApp (executable,
                 is_editor_mode, is_start_flag, additional_properties, cwd, timeout,
                 additional_externalized_properties)
 
         Example:
             ```python
             from aind_behavior_services import (
-                AindBehaviorRigModel,
-                AindBehaviorSessionModel,
-                AindBehaviorTaskLogicModel
+                Rig,
+                Session,
+                Task
             )
 
             # Create models
-            rig = AindBehaviorRigModel(...)
-            session = AindBehaviorSessionModel(...)
-            task_logic = AindBehaviorTaskLogicModel(...)
+            rig = Rig(...)
+            session = Session(...)
+            task = Task(...)
 
             # Create app with automatic configuration
             app = AindBehaviorServicesBonsaiApp(
                 workflow="behavior_workflow.bonsai",
                 launcher=my_launcher,
                 rig=rig,
                 session=session,
-                task_logic=task_logic
+                task=task
             )
             app.run()
 
             # The workflow will receive:
             # -p:"RigPath"="/tmp/rig_temp.json"
             # -p:"SessionPath"="/tmp/session_temp.json"
-            # -p:"TaskLogicPath"="/tmp/task_logic_temp.json"
+            # -p:"TaskPath"="/tmp/task_temp.json"
             ```
         """
         self._temp_directory = Path(temp_directory or TMP_DIR)
@@ -245,10 +245,8 @@ def __init__(
             additional_externalized_properties["RigPath"] = os.path.abspath(self._save_temp_model(model=rig))
         if session:
             additional_externalized_properties["SessionPath"] = os.path.abspath(self._save_temp_model(model=session))
-        if task_logic:
-            additional_externalized_properties["TaskLogicPath"] = os.path.abspath(
-                self._save_temp_model(model=task_logic)
-            )
+        if task:
+            additional_externalized_properties["TaskPath"] = os.path.abspath(self._save_temp_model(model=task))
         super().__init__(
             workflow=workflow, additional_externalized_properties=additional_externalized_properties, **kwargs
         )

src/clabe/constants.py

@@ -34,10 +34,10 @@ class ByAnimalFiles(enum.StrEnum):
 
     Example:
         ```python
-        # Use the task logic file type
-        filename = f"{ByAnimalFiles.TASK_LOGIC}.json"
+        # Use the task file type
+        filename = f"{ByAnimalFiles.task}.json"
         ```
     """
 
-    TASK_LOGIC = "task_logic"
+    TASK = "task"
     TRAINER_STATE = "trainer_state"

src/clabe/data_mapper/helpers.py

@@ -6,16 +6,14 @@
 from typing import Dict, List, Union
 
 import pydantic
-from aind_behavior_services import AindBehaviorRigModel
+from aind_behavior_services import Rig
 from aind_behavior_services.rig.cameras import CameraController, CameraTypes
 from aind_behavior_services.utils import get_fields_of_type
 
 logger = logging.getLogger(__name__)
 
 
-def get_cameras(
-    rig_instance: AindBehaviorRigModel, exclude_without_video_writer: bool = True
-) -> Dict[str, CameraTypes]:
+def get_cameras(rig_instance: Rig, exclude_without_video_writer: bool = True) -> Dict[str, CameraTypes]:
     """
     Retrieves cameras from a rig instance.
 

src/clabe/data_transfer/aind_watchdog.py

@@ -11,7 +11,7 @@
 import pydantic
 import requests
 import yaml
-from aind_behavior_services import AindBehaviorSessionModel
+from aind_behavior_services import Session
 from aind_watchdog_service.models import (
     ManifestConfig,
 )
@@ -76,7 +76,7 @@ def __init__(
         self,
         source: PathLike | list[PathLike],
         settings: WatchdogSettings,
-        session: AindBehaviorSessionModel,
+        session: Session,
         *,
         validate: bool = True,
         email_from_experimenter_builder: Optional[
@@ -210,7 +210,7 @@ def is_valid_project_name(self) -> bool:
         project_names = self._get_project_names()
         return self._settings.project_name in project_names
 
-    def _create_manifest_from_session(self, session: AindBehaviorSessionModel) -> ManifestConfig:
+    def _create_manifest_from_session(self, session: Session) -> ManifestConfig:
         """
         Creates a ManifestConfig from an aind-behavior-services session.
 

src/clabe/launcher/_base.py

@@ -10,9 +10,7 @@
 
 import git.exc
 import pydantic
-from aind_behavior_services import (
-    AindBehaviorSessionModel,
-)
+from aind_behavior_services import Session
 
 from .. import __version__, logging_helper
 from ..constants import TMP_DIR
@@ -90,7 +88,7 @@ def __init__(
 
         self._logger = _logger
 
-        self._session: Optional[AindBehaviorSessionModel] = None
+        self._session: Optional[Session] = None
         self._has_copied_logs = False
 
     @property
@@ -119,7 +117,7 @@ def session_directory(self) -> Path:
         else:
             return Path(self.data_directory) / (session.session_name if session.session_name is not None else "")
 
-    def register_session(self, session: AindBehaviorSessionModel, data_directory: os.PathLike) -> Self:
+    def register_session(self, session: Session, data_directory: os.PathLike) -> Self:
         """
         Registers the session model with the launcher and creates the session
         data directory structure.
@@ -144,12 +142,12 @@ def register_session(self, session: AindBehaviorSessionModel, data_directory: os
         return self
 
     @property
-    def session(self) -> AindBehaviorSessionModel:
+    def session(self) -> Session:
         """
         Returns the registered session model.
 
         Returns:
-            AindBehaviorSessionModel: The session model
+            Session: The session model
 
         Raises:
             ValueError: If session is not set