From 9247f5fabc67cb62fcf61c2b515f9cb7010131ff Mon Sep 17 00:00:00 2001
From: bruno-f-cruz <7049351+bruno-f-cruz@users.noreply.github.com>
Date: Mon, 17 Jun 2024 09:56:45 -0700
Subject: [PATCH 1/9] Bump curriculum core package dependency

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 66eaf2d0..29aadc0b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,7 +18,7 @@ dependencies = [
     'gitpython>=3.1, <4.0',
     'scikit-learn',
     'semver',
-    'aind-behavior-curriculum@git+https://github.com/AllenNeuralDynamics/aind-behavior-curriculum@v0.0.2',
+    'aind-behavior-curriculum@git+https://github.com/AllenNeuralDynamics/aind-behavior-curriculum@v0.0.27',
 
 ]
 

From ebcc9145572044129fc50009cb7eb5b170ec1696 Mon Sep 17 00:00:00 2001
From: bruno-f-cruz <7049351+bruno-f-cruz@users.noreply.github.com>
Date: Mon, 17 Jun 2024 09:56:56 -0700
Subject: [PATCH 2/9] Linting

---
 .../aind_behavior_services/launcher.py        | 21 ++++++++++++-------
 .../aind_behavior_services/utils.py           | 15 +++++++++----
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/src/DataSchemas/aind_behavior_services/launcher.py b/src/DataSchemas/aind_behavior_services/launcher.py
index 3d9b6642..0d86025c 100644
--- a/src/DataSchemas/aind_behavior_services/launcher.py
+++ b/src/DataSchemas/aind_behavior_services/launcher.py
@@ -132,7 +132,9 @@ def _print_diagnosis(self) -> None:
         """
         Prints the diagnosis information for the launcher.
 
-        This method prints the diagnosis information, including the computer name, data directory, and config library directory.
+        This method prints the diagnosis information,
+        including the computer name, data directory,
+        and config library directory.
 
         Parameters:
             None
@@ -171,11 +173,14 @@ def _print_header(self) -> None:
         print("-------------------------------")
         print(self._HEADER)
         print(
-            f"TaskLogic ({self.task_logic_schema.__name__}) Schema Version: {self.task_logic_schema.model_construct().version}"
+            f"TaskLogic ({self.task_logic_schema.__name__}) \
+                Schema Version: {self.task_logic_schema.model_construct().version}"
         )
-        print(f"Rig ({self.rig_schema.__name__}) Schema Version: {self.rig_schema.model_construct().version}")
+        print(f"Rig ({self.rig_schema.__name__}) \
+              Schema Version: {self.rig_schema.model_construct().version}")
         print(
-            f"Session ({self.session_schema.__name__}) Schema Version: {self.session_schema.model_construct().version}"
+            f"Session ({self.session_schema.__name__}) \
+                Schema Version: {self.session_schema.model_construct().version}"
         )
         print("-------------------------------")
         if self._dev_mode:
@@ -187,8 +192,9 @@ def save_temp_model(self, model: Union[TRig, TSession, TTaskLogic], folder: Opti
 
         Args:
             model (BaseModel): The model to be saved.
-            folder (Optional[os.PathLike | str]): The folder where the model should be saved. If not provided, the default
-                temporary folder will be used.
+            folder (Optional[os.PathLike | str]):
+              The folder where the model should be saved.
+              If not provided, the default temporary folder will be used.
 
         Returns:
             str: The file path of the saved model.
@@ -221,7 +227,8 @@ def _validate_dependencies(self) -> None:
             raise FileNotFoundError(f"Config library not found! Expected {self.config_library_dir}.")
         if not (os.path.isdir(os.path.join(self.config_library_dir, "Rig", self.computer_name))):
             raise FileNotFoundError(
-                f"Rig configuration not found! Expected {os.path.join(self.config_library_dir, self.RIG_DIR, self.computer_name)}."
+                f"Rig configuration not found! \
+                    Expected {os.path.join(self.config_library_dir, self.RIG_DIR, self.computer_name)}."
             )
         if not (os.path.isfile(os.path.join(self.default_workflow))):
             raise FileNotFoundError(f"Bonsai workflow file not found! Expected {self.default_workflow}.")
diff --git a/src/DataSchemas/aind_behavior_services/utils.py b/src/DataSchemas/aind_behavior_services/utils.py
index 9a6f86d5..f02e028b 100644
--- a/src/DataSchemas/aind_behavior_services/utils.py
+++ b/src/DataSchemas/aind_behavior_services/utils.py
@@ -203,10 +203,17 @@ def bonsai_sgen(
         CompletedProcess: The result of running the command.
     Args:
         schema_path (PathLike): Target Json Schema file
-        output_path (PathLike): Specifies the name of the file containing the generated code.
-        namespace (Optional[str], optional): Specifies the namespace to use for all generated serialization classes. Defaults to DataSchema.
-        root_element (Optional[str], optional):  Specifies the name of the class used to represent the schema root element. If None, it will use the json schema root element. Defaults to None.
-        serializer (Optional[List[BonsaiSgenSerializers]], optional): Specifies the serializer data annotations to include in the generated classes. Defaults to None.
+        output_path (PathLike): Specifies the name of the
+          file containing the generated code.
+        namespace (Optional[str], optional): Specifies the
+          namespace to use for all generated serialization
+          classes. Defaults to DataSchema.
+        root_element (Optional[str], optional):  Specifies the
+          name of the class used to represent the schema root element.
+          If None, it will use the json schema root element. Defaults to None.
+        serializer (Optional[List[BonsaiSgenSerializers]], optional):
+          Specifies the serializer data annotations to include in the generated classes.
+          Defaults to None.
     """
 
     if serializer is None:

From 60f3c22fd063b2bb8fd7f266e7cda329f9a59f00 Mon Sep 17 00:00:00 2001
From: bruno-f-cruz <7049351+bruno-f-cruz@users.noreply.github.com>
Date: Mon, 17 Jun 2024 13:27:20 -0700
Subject: [PATCH 3/9] Regenerate schemas to latest curriculum spec

---
 .../aind_manipulator_calibration_logic.json   | 21 +++++++++--
 .../schemas/load_cells_calibration_logic.json | 21 +++++++++--
 .../olfactometer_calibration_logic.json       | 21 +++++++++--
 .../water_valve_calibration_logic.json        | 21 +++++++++--
 .../AindManipulatorCalibrationLogic.cs        | 37 +++++++++++++++----
 src/Extensions/LoadCellsCalibrationLogic.cs   | 37 +++++++++++++++----
 .../OlfactometerCalibrationLogic.cs           | 37 +++++++++++++++----
 7 files changed, 155 insertions(+), 40 deletions(-)

diff --git a/src/DataSchemas/schemas/aind_manipulator_calibration_logic.json b/src/DataSchemas/schemas/aind_manipulator_calibration_logic.json
index 5b06a1ac..68cb27bc 100644
--- a/src/DataSchemas/schemas/aind_manipulator_calibration_logic.json
+++ b/src/DataSchemas/schemas/aind_manipulator_calibration_logic.json
@@ -47,19 +47,32 @@
       "title": "Description",
       "type": "string"
     },
+    "task_parameters": {
+      "allOf": [
+        {
+          "$ref": "#/definitions/CalibrationParameters"
+        }
+      ],
+      "title": "Task parameters"
+    },
     "version": {
       "const": "0.2.0",
       "default": "0.2.0",
       "title": "Version",
       "type": "string"
     },
-    "task_parameters": {
-      "allOf": [
+    "stage_name": {
+      "default": null,
+      "description": "Optional stage name the `Task` object instance represents.",
+      "oneOf": [
         {
-          "$ref": "#/definitions/CalibrationParameters"
+          "type": "string"
+        },
+        {
+          "type": "null"
         }
       ],
-      "title": "Task parameters"
+      "title": "Stage Name"
     }
   },
   "required": [
diff --git a/src/DataSchemas/schemas/load_cells_calibration_logic.json b/src/DataSchemas/schemas/load_cells_calibration_logic.json
index 816cf73b..4a8661ca 100644
--- a/src/DataSchemas/schemas/load_cells_calibration_logic.json
+++ b/src/DataSchemas/schemas/load_cells_calibration_logic.json
@@ -76,19 +76,32 @@
       "title": "Description",
       "type": "string"
     },
+    "task_parameters": {
+      "allOf": [
+        {
+          "$ref": "#/definitions/CalibrationParameters"
+        }
+      ],
+      "title": "Task parameters"
+    },
     "version": {
       "const": "0.4.0",
       "default": "0.4.0",
       "title": "Version",
       "type": "string"
     },
-    "task_parameters": {
-      "allOf": [
+    "stage_name": {
+      "default": null,
+      "description": "Optional stage name the `Task` object instance represents.",
+      "oneOf": [
         {
-          "$ref": "#/definitions/CalibrationParameters"
+          "type": "string"
+        },
+        {
+          "type": "null"
         }
       ],
-      "title": "Task parameters"
+      "title": "Stage Name"
     }
   },
   "required": [
diff --git a/src/DataSchemas/schemas/olfactometer_calibration_logic.json b/src/DataSchemas/schemas/olfactometer_calibration_logic.json
index 33042f69..35a8ff89 100644
--- a/src/DataSchemas/schemas/olfactometer_calibration_logic.json
+++ b/src/DataSchemas/schemas/olfactometer_calibration_logic.json
@@ -156,19 +156,32 @@
       "title": "Description",
       "type": "string"
     },
+    "task_parameters": {
+      "allOf": [
+        {
+          "$ref": "#/definitions/CalibrationParameters"
+        }
+      ],
+      "title": "Task parameters"
+    },
     "version": {
       "const": "0.4.0",
       "default": "0.4.0",
       "title": "Version",
       "type": "string"
     },
-    "task_parameters": {
-      "allOf": [
+    "stage_name": {
+      "default": null,
+      "description": "Optional stage name the `Task` object instance represents.",
+      "oneOf": [
         {
-          "$ref": "#/definitions/CalibrationParameters"
+          "type": "string"
+        },
+        {
+          "type": "null"
         }
       ],
-      "title": "Task parameters"
+      "title": "Stage Name"
     }
   },
   "required": [
diff --git a/src/DataSchemas/schemas/water_valve_calibration_logic.json b/src/DataSchemas/schemas/water_valve_calibration_logic.json
index 817418bc..c64ae46e 100644
--- a/src/DataSchemas/schemas/water_valve_calibration_logic.json
+++ b/src/DataSchemas/schemas/water_valve_calibration_logic.json
@@ -75,19 +75,32 @@
       "title": "Description",
       "type": "string"
     },
+    "task_parameters": {
+      "allOf": [
+        {
+          "$ref": "#/definitions/CalibrationParameters"
+        }
+      ],
+      "title": "Task parameters"
+    },
     "version": {
       "const": "0.4.0",
       "default": "0.4.0",
       "title": "Version",
       "type": "string"
     },
-    "task_parameters": {
-      "allOf": [
+    "stage_name": {
+      "default": null,
+      "description": "Optional stage name the `Task` object instance represents.",
+      "oneOf": [
         {
-          "$ref": "#/definitions/CalibrationParameters"
+          "type": "string"
+        },
+        {
+          "type": "null"
         }
       ],
-      "title": "Task parameters"
+      "title": "Stage Name"
     }
   },
   "required": [
diff --git a/src/Extensions/AindManipulatorCalibrationLogic.cs b/src/Extensions/AindManipulatorCalibrationLogic.cs
index e68a5d6e..e4800bf2 100644
--- a/src/Extensions/AindManipulatorCalibrationLogic.cs
+++ b/src/Extensions/AindManipulatorCalibrationLogic.cs
@@ -106,9 +106,11 @@ public partial class CalibrationLogic
     
         private string _description = "";
     
+        private CalibrationParameters _taskParameters = new CalibrationParameters();
+    
         private string _version = "0.2.0";
     
-        private CalibrationParameters _taskParameters = new CalibrationParameters();
+        private string _stageName;
     
         public CalibrationLogic()
         {
@@ -118,8 +120,9 @@ protected CalibrationLogic(CalibrationLogic other)
         {
             _name = other._name;
             _description = other._description;
-            _version = other._version;
             _taskParameters = other._taskParameters;
+            _version = other._version;
+            _stageName = other._stageName;
         }
     
         [Newtonsoft.Json.JsonPropertyAttribute("name")]
@@ -152,6 +155,20 @@ public string Description
             }
         }
     
+        [System.Xml.Serialization.XmlIgnoreAttribute()]
+        [Newtonsoft.Json.JsonPropertyAttribute("task_parameters", Required=Newtonsoft.Json.Required.Always)]
+        public CalibrationParameters TaskParameters
+        {
+            get
+            {
+                return _taskParameters;
+            }
+            set
+            {
+                _taskParameters = value;
+            }
+        }
+    
         [Newtonsoft.Json.JsonPropertyAttribute("version")]
         public string Version
         {
@@ -165,17 +182,20 @@ public string Version
             }
         }
     
-        [System.Xml.Serialization.XmlIgnoreAttribute()]
-        [Newtonsoft.Json.JsonPropertyAttribute("task_parameters", Required=Newtonsoft.Json.Required.Always)]
-        public CalibrationParameters TaskParameters
+        /// <summary>
+        /// Optional stage name the `Task` object instance represents.
+        /// </summary>
+        [Newtonsoft.Json.JsonPropertyAttribute("stage_name")]
+        [System.ComponentModel.DescriptionAttribute("Optional stage name the `Task` object instance represents.")]
+        public string StageName
         {
             get
             {
-                return _taskParameters;
+                return _stageName;
             }
             set
             {
-                _taskParameters = value;
+                _stageName = value;
             }
         }
     
@@ -193,8 +213,9 @@ protected virtual bool PrintMembers(System.Text.StringBuilder stringBuilder)
         {
             stringBuilder.Append("name = " + _name + ", ");
             stringBuilder.Append("description = " + _description + ", ");
+            stringBuilder.Append("task_parameters = " + _taskParameters + ", ");
             stringBuilder.Append("version = " + _version + ", ");
-            stringBuilder.Append("task_parameters = " + _taskParameters);
+            stringBuilder.Append("stage_name = " + _stageName);
             return true;
         }
     
diff --git a/src/Extensions/LoadCellsCalibrationLogic.cs b/src/Extensions/LoadCellsCalibrationLogic.cs
index 30950e40..8499c27e 100644
--- a/src/Extensions/LoadCellsCalibrationLogic.cs
+++ b/src/Extensions/LoadCellsCalibrationLogic.cs
@@ -154,9 +154,11 @@ public partial class CalibrationLogic
     
         private string _description = "";
     
+        private CalibrationParameters _taskParameters = new CalibrationParameters();
+    
         private string _version = "0.4.0";
     
-        private CalibrationParameters _taskParameters = new CalibrationParameters();
+        private string _stageName;
     
         public CalibrationLogic()
         {
@@ -166,8 +168,9 @@ protected CalibrationLogic(CalibrationLogic other)
         {
             _name = other._name;
             _description = other._description;
-            _version = other._version;
             _taskParameters = other._taskParameters;
+            _version = other._version;
+            _stageName = other._stageName;
         }
     
         [Newtonsoft.Json.JsonPropertyAttribute("name")]
@@ -200,6 +203,20 @@ public string Description
             }
         }
     
+        [System.Xml.Serialization.XmlIgnoreAttribute()]
+        [Newtonsoft.Json.JsonPropertyAttribute("task_parameters", Required=Newtonsoft.Json.Required.Always)]
+        public CalibrationParameters TaskParameters
+        {
+            get
+            {
+                return _taskParameters;
+            }
+            set
+            {
+                _taskParameters = value;
+            }
+        }
+    
         [Newtonsoft.Json.JsonPropertyAttribute("version")]
         public string Version
         {
@@ -213,17 +230,20 @@ public string Version
             }
         }
     
-        [System.Xml.Serialization.XmlIgnoreAttribute()]
-        [Newtonsoft.Json.JsonPropertyAttribute("task_parameters", Required=Newtonsoft.Json.Required.Always)]
-        public CalibrationParameters TaskParameters
+        /// <summary>
+        /// Optional stage name the `Task` object instance represents.
+        /// </summary>
+        [Newtonsoft.Json.JsonPropertyAttribute("stage_name")]
+        [System.ComponentModel.DescriptionAttribute("Optional stage name the `Task` object instance represents.")]
+        public string StageName
         {
             get
             {
-                return _taskParameters;
+                return _stageName;
             }
             set
             {
-                _taskParameters = value;
+                _stageName = value;
             }
         }
     
@@ -241,8 +261,9 @@ protected virtual bool PrintMembers(System.Text.StringBuilder stringBuilder)
         {
             stringBuilder.Append("name = " + _name + ", ");
             stringBuilder.Append("description = " + _description + ", ");
+            stringBuilder.Append("task_parameters = " + _taskParameters + ", ");
             stringBuilder.Append("version = " + _version + ", ");
-            stringBuilder.Append("task_parameters = " + _taskParameters);
+            stringBuilder.Append("stage_name = " + _stageName);
             return true;
         }
     
diff --git a/src/Extensions/OlfactometerCalibrationLogic.cs b/src/Extensions/OlfactometerCalibrationLogic.cs
index bd76c0db..cf078baf 100644
--- a/src/Extensions/OlfactometerCalibrationLogic.cs
+++ b/src/Extensions/OlfactometerCalibrationLogic.cs
@@ -382,9 +382,11 @@ public partial class CalibrationLogic
     
         private string _description = "";
     
+        private CalibrationParameters _taskParameters = new CalibrationParameters();
+    
         private string _version = "0.4.0";
     
-        private CalibrationParameters _taskParameters = new CalibrationParameters();
+        private string _stageName;
     
         public CalibrationLogic()
         {
@@ -394,8 +396,9 @@ protected CalibrationLogic(CalibrationLogic other)
         {
             _name = other._name;
             _description = other._description;
-            _version = other._version;
             _taskParameters = other._taskParameters;
+            _version = other._version;
+            _stageName = other._stageName;
         }
     
         [Newtonsoft.Json.JsonPropertyAttribute("name")]
@@ -428,6 +431,20 @@ public string Description
             }
         }
     
+        [System.Xml.Serialization.XmlIgnoreAttribute()]
+        [Newtonsoft.Json.JsonPropertyAttribute("task_parameters", Required=Newtonsoft.Json.Required.Always)]
+        public CalibrationParameters TaskParameters
+        {
+            get
+            {
+                return _taskParameters;
+            }
+            set
+            {
+                _taskParameters = value;
+            }
+        }
+    
         [Newtonsoft.Json.JsonPropertyAttribute("version")]
         public string Version
         {
@@ -441,17 +458,20 @@ public string Version
             }
         }
     
-        [System.Xml.Serialization.XmlIgnoreAttribute()]
-        [Newtonsoft.Json.JsonPropertyAttribute("task_parameters", Required=Newtonsoft.Json.Required.Always)]
-        public CalibrationParameters TaskParameters
+        /// <summary>
+        /// Optional stage name the `Task` object instance represents.
+        /// </summary>
+        [Newtonsoft.Json.JsonPropertyAttribute("stage_name")]
+        [System.ComponentModel.DescriptionAttribute("Optional stage name the `Task` object instance represents.")]
+        public string StageName
         {
             get
             {
-                return _taskParameters;
+                return _stageName;
             }
             set
             {
-                _taskParameters = value;
+                _stageName = value;
             }
         }
     
@@ -469,8 +489,9 @@ protected virtual bool PrintMembers(System.Text.StringBuilder stringBuilder)
         {
             stringBuilder.Append("name = " + _name + ", ");
             stringBuilder.Append("description = " + _description + ", ");
+            stringBuilder.Append("task_parameters = " + _taskParameters + ", ");
             stringBuilder.Append("version = " + _version + ", ");
-            stringBuilder.Append("task_parameters = " + _taskParameters);
+            stringBuilder.Append("stage_name = " + _stageName);
             return true;
         }
     

From 362f479fbc3348b7dae0e5309a28c929b28431ba Mon Sep 17 00:00:00 2001
From: bruno-f-cruz <7049351+bruno-f-cruz@users.noreply.github.com>
Date: Mon, 17 Jun 2024 15:22:24 -0700
Subject: [PATCH 4/9] Isolate parser method

---
 .../aind_behavior_services/launcher.py        | 80 ++++++++++---------
 1 file changed, 43 insertions(+), 37 deletions(-)

diff --git a/src/DataSchemas/aind_behavior_services/launcher.py b/src/DataSchemas/aind_behavior_services/launcher.py
index 0d86025c..7babcbda 100644
--- a/src/DataSchemas/aind_behavior_services/launcher.py
+++ b/src/DataSchemas/aind_behavior_services/launcher.py
@@ -526,43 +526,8 @@ def __init__(
         **launcher_kwargs,
     ) -> None:
 
-        parser = argparse.ArgumentParser()
-
-        parser.add_argument("--data_dir", help="Specify the data directory")
-        parser.add_argument("--remote_data_dir", help="Specify the remote data directory")
-        parser.add_argument("--repository_dir", help="Specify the repository directory")
-        parser.add_argument("--config_library_dir", help="Specify the configuration library directory")
-        parser.add_argument("--workflow", help="Specify the workflow")
-        parser.add_argument(
-            "--force_create_directories",
-            help="Specify whether to force create directories",
-            action="store_true",
-            default=False,
-        )
-        parser.add_argument("--dev_mode", help="Specify whether to run in dev mode", action="store_true", default=False)
-        parser.add_argument(
-            "--bonsai_is_editor_mode",
-            help="Specify whether to run in Bonsai editor mode",
-            action="store_false",
-            default=True,
-        )
-        parser.add_argument(
-            "--bonsai_is_start_flag",
-            help="Specify whether to start the Bonsai workflow",
-            action="store_false",
-            default=True,
-        )
-        parser.add_argument(
-            "--allow_dirty_repo", help="Specify whether to allow a dirty repository", action="store_true", default=False
-        )
-        parser.add_argument(
-            "--skip_hardware_validation",
-            help="Specify whether to skip hardware validation",
-            action="store_true",
-            default=False,
-        )
-
-        args = parser.parse_args()
+        parser = self._get_default_arg_parser()
+        args, _ = parser.parse_known_args()
 
         # optional parameters that override the defaults
         data_dir = args.data_dir if args.data_dir is not None else data_dir
@@ -607,3 +572,44 @@ def make_folder_structure(self) -> None:
 
     def _validate_dependencies(self) -> None:
         self.launcher._validate_dependencies()
+
+    @staticmethod
+    def _get_default_arg_parser() -> argparse.ArgumentParser:
+
+        parser = argparse.ArgumentParser()
+
+        parser.add_argument("--data_dir", help="Specify the data directory")
+        parser.add_argument("--remote_data_dir", help="Specify the remote data directory")
+        parser.add_argument("--repository_dir", help="Specify the repository directory")
+        parser.add_argument("--config_library_dir", help="Specify the configuration library directory")
+        parser.add_argument("--workflow", help="Specify the workflow")
+        parser.add_argument(
+            "--force_create_directories",
+            help="Specify whether to force create directories",
+            action="store_true",
+            default=False,
+        )
+        parser.add_argument("--dev_mode", help="Specify whether to run in dev mode", action="store_true", default=False)
+        parser.add_argument(
+            "--bonsai_is_editor_mode",
+            help="Specify whether to run in Bonsai editor mode",
+            action="store_false",
+            default=True,
+        )
+        parser.add_argument(
+            "--bonsai_is_start_flag",
+            help="Specify whether to start the Bonsai workflow",
+            action="store_false",
+            default=True,
+        )
+        parser.add_argument(
+            "--allow_dirty_repo", help="Specify whether to allow a dirty repository", action="store_true", default=False
+        )
+        parser.add_argument(
+            "--skip_hardware_validation",
+            help="Specify whether to skip hardware validation",
+            action="store_true",
+            default=False,
+        )
+
+        return parser
\ No newline at end of file

From 4cf27e81690b0413c139556d3772132592c6b81e Mon Sep 17 00:00:00 2001
From: bruno-f-cruz <7049351+bruno-f-cruz@users.noreply.github.com>
Date: Mon, 17 Jun 2024 15:22:45 -0700
Subject: [PATCH 5/9] Add version to Task

---
 src/DataSchemas/aind_behavior_services/task_logic/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/DataSchemas/aind_behavior_services/task_logic/__init__.py b/src/DataSchemas/aind_behavior_services/task_logic/__init__.py
index 939ad1ba..ea340b10 100644
--- a/src/DataSchemas/aind_behavior_services/task_logic/__init__.py
+++ b/src/DataSchemas/aind_behavior_services/task_logic/__init__.py
@@ -13,6 +13,7 @@ class TaskParameters(curriculum_task.TaskParameters):
 class AindBehaviorTaskLogicModel(curriculum_task.Task):
 
     task_parameters: TaskParameters = Field(..., description="Parameters of the task logic", validate_default=True)
+    version: str = Field(..., pattern=curriculum_task.SEMVER_REGEX, description="task schema version")
 
     @field_validator("version", mode="before")
     @classmethod

From ab9548ac1da79a92104b6dc7d3063e54141b95b3 Mon Sep 17 00:00:00 2001
From: bruno-f-cruz <7049351+bruno-f-cruz@users.noreply.github.com>
Date: Mon, 17 Jun 2024 15:23:22 -0700
Subject: [PATCH 6/9] Regenerate schema

---
 src/Extensions/WaterValveCalibrationLogic.cs | 37 +++++++++++++++-----
 1 file changed, 29 insertions(+), 8 deletions(-)

diff --git a/src/Extensions/WaterValveCalibrationLogic.cs b/src/Extensions/WaterValveCalibrationLogic.cs
index af22adcc..4f9450a7 100644
--- a/src/Extensions/WaterValveCalibrationLogic.cs
+++ b/src/Extensions/WaterValveCalibrationLogic.cs
@@ -175,9 +175,11 @@ public partial class CalibrationLogic
     
         private string _description = "";
     
+        private CalibrationParameters _taskParameters = new CalibrationParameters();
+    
         private string _version = "0.4.0";
     
-        private CalibrationParameters _taskParameters = new CalibrationParameters();
+        private string _stageName;
     
         public CalibrationLogic()
         {
@@ -187,8 +189,9 @@ protected CalibrationLogic(CalibrationLogic other)
         {
             _name = other._name;
             _description = other._description;
-            _version = other._version;
             _taskParameters = other._taskParameters;
+            _version = other._version;
+            _stageName = other._stageName;
         }
     
         [Newtonsoft.Json.JsonPropertyAttribute("name")]
@@ -221,6 +224,20 @@ public string Description
             }
         }
     
+        [System.Xml.Serialization.XmlIgnoreAttribute()]
+        [Newtonsoft.Json.JsonPropertyAttribute("task_parameters", Required=Newtonsoft.Json.Required.Always)]
+        public CalibrationParameters TaskParameters
+        {
+            get
+            {
+                return _taskParameters;
+            }
+            set
+            {
+                _taskParameters = value;
+            }
+        }
+    
         [Newtonsoft.Json.JsonPropertyAttribute("version")]
         public string Version
         {
@@ -234,17 +251,20 @@ public string Version
             }
         }
     
-        [System.Xml.Serialization.XmlIgnoreAttribute()]
-        [Newtonsoft.Json.JsonPropertyAttribute("task_parameters", Required=Newtonsoft.Json.Required.Always)]
-        public CalibrationParameters TaskParameters
+        /// <summary>
+        /// Optional stage name the `Task` object instance represents.
+        /// </summary>
+        [Newtonsoft.Json.JsonPropertyAttribute("stage_name")]
+        [System.ComponentModel.DescriptionAttribute("Optional stage name the `Task` object instance represents.")]
+        public string StageName
         {
             get
             {
-                return _taskParameters;
+                return _stageName;
             }
             set
             {
-                _taskParameters = value;
+                _stageName = value;
             }
         }
     
@@ -262,8 +282,9 @@ protected virtual bool PrintMembers(System.Text.StringBuilder stringBuilder)
         {
             stringBuilder.Append("name = " + _name + ", ");
             stringBuilder.Append("description = " + _description + ", ");
+            stringBuilder.Append("task_parameters = " + _taskParameters + ", ");
             stringBuilder.Append("version = " + _version + ", ");
-            stringBuilder.Append("task_parameters = " + _taskParameters);
+            stringBuilder.Append("stage_name = " + _stageName);
             return true;
         }
     

From 9c8ece251042b59435635b2280abb51e17f93819 Mon Sep 17 00:00:00 2001
From: bruno-f-cruz <7049351+bruno-f-cruz@users.noreply.github.com>
Date: Mon, 17 Jun 2024 15:33:22 -0700
Subject: [PATCH 7/9] Ensure output folder is created if needed

---
 examples/aind_manipulator.py | 2 ++
 examples/load_cells.py       | 3 +++
 examples/olfactometer.py     | 3 +++
 examples/water_valve.py      | 3 +++
 4 files changed, 11 insertions(+)

diff --git a/examples/aind_manipulator.py b/examples/aind_manipulator.py
index 851591d9..76cd8f90 100644
--- a/examples/aind_manipulator.py
+++ b/examples/aind_manipulator.py
@@ -1,4 +1,5 @@
 import datetime
+import os
 
 from aind_behavior_services.base import get_commit_hash
 from aind_behavior_services.calibration import aind_manipulator as m
@@ -41,6 +42,7 @@
 )
 
 seed_path = "local/aind_manipulator_{suffix}.json"
+os.makedirs(os.path.dirname(seed_path), exist_ok=True)
 with open(seed_path.format(suffix="calibration_logic"), "w") as f:
     f.write(calibration_logic.model_dump_json(indent=3))
 with open(seed_path.format(suffix="session"), "w") as f:
diff --git a/examples/load_cells.py b/examples/load_cells.py
index 5db0a5c3..47b65a3f 100644
--- a/examples/load_cells.py
+++ b/examples/load_cells.py
@@ -1,4 +1,5 @@
 import datetime
+import os
 
 from aind_behavior_services.base import get_commit_hash
 from aind_behavior_services.calibration import load_cells as lc
@@ -42,6 +43,8 @@
 )
 
 seed_path = "local/load_cells_{suffix}.json"
+os.makedirs(os.path.dirname(seed_path), exist_ok=True)
+
 with open(seed_path.format(suffix="calibration_logic"), "w") as f:
     f.write(calibration_logic.model_dump_json(indent=3))
 with open(seed_path.format(suffix="session"), "w") as f:
diff --git a/examples/olfactometer.py b/examples/olfactometer.py
index 675a3f56..b2bdc9a8 100644
--- a/examples/olfactometer.py
+++ b/examples/olfactometer.py
@@ -1,4 +1,5 @@
 import datetime
+import os
 
 from aind_behavior_services.base import get_commit_hash
 from aind_behavior_services.calibration import olfactometer as olf
@@ -65,6 +66,8 @@
 
 
 seed_path = "local/olfactometer_{suffix}.json"
+os.makedirs(os.path.dirname(seed_path), exist_ok=True)
+
 with open(seed_path.format(suffix="calibration_logic"), "w") as f:
     f.write(calibration_logic.model_dump_json(indent=3))
 with open(seed_path.format(suffix="session"), "w") as f:
diff --git a/examples/water_valve.py b/examples/water_valve.py
index c9d55786..edf7bafe 100644
--- a/examples/water_valve.py
+++ b/examples/water_valve.py
@@ -1,4 +1,5 @@
 import datetime
+import os
 
 from aind_behavior_services.base import get_commit_hash
 from aind_behavior_services.calibration import water_valve as wv
@@ -55,6 +56,8 @@ def linear_model(time, slope, offset):
 rig = wv.CalibrationRig(rig_name="WaterValveRig")
 
 seed_path = "local/water_valve_{suffix}.json"
+os.makedirs(os.path.dirname(seed_path), exist_ok=True)
+
 with open(seed_path.format(suffix="calibration_logic"), "w") as f:
     f.write(calibration_logic.model_dump_json(indent=3))
 with open(seed_path.format(suffix="session"), "w") as f:

From a49038803cab32a2fdeee8342a2ea60ddca6d342 Mon Sep 17 00:00:00 2001
From: bruno-f-cruz <7049351+bruno-f-cruz@users.noreply.github.com>
Date: Mon, 17 Jun 2024 15:34:01 -0700
Subject: [PATCH 8/9] Fix test suite to correctly regenerate examples across
 tests

---
 .../aind_behavior_services/launcher.py        |  8 ++++---
 tests/__init__.py                             | 21 +++++++++++++++++++
 tests/test_bonsai.py                          |  4 +++-
 tests/test_examples.py                        | 11 ++--------
 tests/test_water_valve.py                     | 10 +++++----
 5 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/src/DataSchemas/aind_behavior_services/launcher.py b/src/DataSchemas/aind_behavior_services/launcher.py
index 7babcbda..0f825d94 100644
--- a/src/DataSchemas/aind_behavior_services/launcher.py
+++ b/src/DataSchemas/aind_behavior_services/launcher.py
@@ -176,8 +176,10 @@ def _print_header(self) -> None:
             f"TaskLogic ({self.task_logic_schema.__name__}) \
                 Schema Version: {self.task_logic_schema.model_construct().version}"
         )
-        print(f"Rig ({self.rig_schema.__name__}) \
-              Schema Version: {self.rig_schema.model_construct().version}")
+        print(
+            f"Rig ({self.rig_schema.__name__}) \
+              Schema Version: {self.rig_schema.model_construct().version}"
+        )
         print(
             f"Session ({self.session_schema.__name__}) \
                 Schema Version: {self.session_schema.model_construct().version}"
@@ -612,4 +614,4 @@ def _get_default_arg_parser() -> argparse.ArgumentParser:
             default=False,
         )
 
-        return parser
\ No newline at end of file
+        return parser
diff --git a/tests/__init__.py b/tests/__init__.py
index e69de29b..5f2fe603 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -0,0 +1,21 @@
+import glob
+import importlib.util
+from pathlib import Path
+from types import ModuleType
+
+EXAMPLES_DIR = Path(__file__).parents[1] / "examples"
+
+
+def build_example(script_path: str) -> ModuleType:
+    module_name = Path(script_path).stem
+    spec = importlib.util.spec_from_file_location(module_name, script_path)
+    if spec is None:
+        raise ImportError(f"Can't find {script_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def build_examples(examples_dir: Path = EXAMPLES_DIR):
+    for script_path in glob.glob(str(examples_dir / "*.py")):
+        _ = build_example(script_path)
diff --git a/tests/test_bonsai.py b/tests/test_bonsai.py
index 496da5a6..fd27122e 100644
--- a/tests/test_bonsai.py
+++ b/tests/test_bonsai.py
@@ -3,9 +3,10 @@
 from pathlib import Path
 from typing import Dict
 
-import test_examples
 from aind_behavior_services.utils import run_bonsai_process
 
+from . import build_examples
+
 
 class BonsaiTests(unittest.TestCase):
 
@@ -18,6 +19,7 @@ def test_deserialization(self):
             "water_valve",
         ]
 
+        build_examples()
         JSON_ROOT = Path("./local").resolve()
 
         workflow_props: Dict[str, str] = {}
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 55462c2d..22547964 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -2,24 +2,17 @@
 
 import glob
 import unittest
-from pathlib import Path
 
-EXAMPLES_DIR = Path(__file__).parents[1] / "examples"
+from . import EXAMPLES_DIR, build_example
 
 
 class ExampleTests(unittest.TestCase):
     """tests for examples"""
 
     def test_examples(self):
-        import importlib.util
-
         for script_path in glob.glob(str(EXAMPLES_DIR / "*.py")):
             with self.subTest(script_path=script_path):
-                module_name = Path(script_path).stem
-                spec = importlib.util.spec_from_file_location(module_name, script_path)
-                module = importlib.util.module_from_spec(spec)
-                spec.loader.exec_module(module)
-
+                module = build_example(script_path)
                 # Check if the module executed successfully
                 if hasattr(module, "__name__") and module.__name__ == "__main__":
                     self.assertEqual(module.__name__, "__main__", f"Script {script_path} failed to execute")
diff --git a/tests/test_water_valve.py b/tests/test_water_valve.py
index d9ba924c..117a37b3 100644
--- a/tests/test_water_valve.py
+++ b/tests/test_water_valve.py
@@ -19,8 +19,7 @@ def test_calibration(self):
         _delta_times = [0.1, 0.2, 0.3, 0.4, 0.5]
         _slope = 10.1
         _offset = -0.3
-        _linear_model = lambda time: _slope * time + _offset
-        _water_weights = [_linear_model(x) for x in _delta_times]
+        _water_weights = [water_mock_model(x, _slope, _offset) for x in _delta_times]
         _inputs = [
             Measurement(valve_open_interval=0.5, valve_open_time=t[0], water_weight=[t[1]], repeat_count=1)
             for t in zip(_delta_times, _water_weights)
@@ -52,8 +51,7 @@ def test_calibration_on_null_output(self):
         _delta_times = [0.1, 0.2, 0.3, 0.4, 0.5]
         _slope = 10.1
         _offset = -0.3
-        _linear_model = lambda time: _slope * time + _offset
-        _water_weights = [_linear_model(x) for x in _delta_times]
+        _water_weights = [water_mock_model(x, _slope, _offset) for x in _delta_times]
         _inputs = WaterValveCalibrationInput(
             measurements=[
                 Measurement(valve_open_interval=0.5, valve_open_time=t[0], water_weight=[t[1]], repeat_count=1)
@@ -73,5 +71,9 @@ def test_calibration_on_null_output(self):
         self.assertAlmostEqual(1.0, calibration.output.r2, 2, "R2 is not almost equal")
 
 
+def water_mock_model(time: float, slope: float, offset: float) -> float:
+    return slope * time + offset
+
+
 if __name__ == "__main__":
     unittest.main()

From f0b1098515fd82a42cbc7f7b150bec8136958e88 Mon Sep 17 00:00:00 2001
From: bruno-f-cruz <7049351+bruno-f-cruz@users.noreply.github.com>
Date: Mon, 17 Jun 2024 15:34:12 -0700
Subject: [PATCH 9/9] Bump package version

---
 src/DataSchemas/aind_behavior_services/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/DataSchemas/aind_behavior_services/__init__.py b/src/DataSchemas/aind_behavior_services/__init__.py
index 60fe81f3..02c8292d 100644
--- a/src/DataSchemas/aind_behavior_services/__init__.py
+++ b/src/DataSchemas/aind_behavior_services/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.7.7"
+__version__ = "0.7.8"
 
 from .rig import AindBehaviorRigModel  # noqa: F401
 from .session import AindBehaviorSessionModel  # noqa: F401