pickers.default_behavior

DefaultBehaviorPickerSettings

Bases: ServiceSettings

Settings for the default behavior picker.

Attributes:

Name	Type	Description
config_library_dir	PathLike	The directory where configuration files are stored.

settings_customise_sources classmethod

settings_customise_sources(
    settings_cls: Type[BaseSettings],
    init_settings: PydanticBaseSettingsSource,
    env_settings: PydanticBaseSettingsSource,
    dotenv_settings: PydanticBaseSettingsSource,
    file_secret_settings: PydanticBaseSettingsSource,
) -> Tuple[PydanticBaseSettingsSource, ...]

Customizes the settings sources to include the safe YAML settings source.

Parameters:

Name	Type	Description	Default
settings_cls	Type[BaseSettings]	The settings class	required
init_settings	PydanticBaseSettingsSource	The initial settings source	required
env_settings	PydanticBaseSettingsSource	The environment settings source	required
dotenv_settings	PydanticBaseSettingsSource	The dotenv settings source	required
file_secret_settings	PydanticBaseSettingsSource	The file secret settings source	required

Returns:

Type	Description
Tuple[PydanticBaseSettingsSource, ...]	Tuple[PydanticBaseSettingsSource, ...]: A tuple of settings sources

Source code in src/clabe/services.py

@classmethod
def settings_customise_sources(
    cls,
    settings_cls: t.Type[ps.BaseSettings],
    init_settings: ps.PydanticBaseSettingsSource,
    env_settings: ps.PydanticBaseSettingsSource,
    dotenv_settings: ps.PydanticBaseSettingsSource,
    file_secret_settings: ps.PydanticBaseSettingsSource,
) -> t.Tuple[ps.PydanticBaseSettingsSource, ...]:
    """
    Customizes the settings sources to include the safe YAML settings source.

    Args:
        settings_cls: The settings class
        init_settings: The initial settings source
        env_settings: The environment settings source
        dotenv_settings: The dotenv settings source
        file_secret_settings: The file secret settings source

    Returns:
        Tuple[PydanticBaseSettingsSource, ...]: A tuple of settings sources
    """
    return (
        init_settings,
        *(
            _SafeYamlSettingsSource(settings_cls, yaml_file=p, yaml_config_section=cls.__yml_section__)
            for p in KNOWN_CONFIG_FILES
        ),
        env_settings,
        dotenv_settings,
        file_secret_settings,
    )

DefaultBehaviorPicker

DefaultBehaviorPicker(
    settings: DefaultBehaviorPickerSettings,
    launcher: Launcher,
    ui_helper: Optional[IUiHelper] = None,
    experimenter_validator: Optional[
        Callable[[str], bool]
    ] = validate_username,
    rig_validator: Optional[
        Callable[[Rig], Rig]
    ] = validate_rig_computer_name,
    use_cache: bool = True,
)

A picker class for selecting rig, session, and task configurations.

Provides methods to initialize directories, pick configurations, and prompt user inputs for various components of the experiment setup. Manages the configuration library structure and user interactions for selecting experiment parameters.

Properties

ui_helper: Helper for user interface interactions trainer_state: The current trainer state config_library_dir: Path to the configuration library directory rig_dir: Path to the rig configurations directory subject_dir: Path to the subject configurations directory task_dir: Path to the task configurations directory

Methods:

Name	Description
pick_rig	Picks the rig configuration
pick_session	Picks the session configuration
pick_task	Picks the task configuration
pick_trainer_state	Picks the trainer state configuration
choose_subject	Allows the user to choose a subject
prompt_experimenter	Prompts for experimenter information
dump_model	Saves a Pydantic model to a file

Initializes the DefaultBehaviorPicker.

Parameters:

Name	Type	Description	Default
settings	DefaultBehaviorPickerSettings	Settings containing configuration including config_library_dir. By default, attempts to rely on DefaultBehaviorPickerSettings to automatic loading from yaml files	required
launcher	Launcher	The launcher instance for managing experiment execution	required
ui_helper	Optional[IUiHelper]	Helper for user interface interactions. If None, uses launcher's ui_helper. Defaults to None	None
experimenter_validator	Optional[Callable[[str], bool]]	Function to validate the experimenter's username. If None, no validation is performed. Defaults to validate_username	validate_username
rig_validator	Optional[Callable[[Rig], Rig]]	Function to validate the rig configuration. If None, no validation is performed. Defaults to validate_rig_computer_name	validate_rig_computer_name
use_cache	bool	Whether to use caching for selections. Defaults to True	True

Source code in src/clabe/pickers/default_behavior.py

def __init__(
    self,
    settings: DefaultBehaviorPickerSettings,
    launcher: Launcher,
    ui_helper: Optional[ui.IUiHelper] = None,
    experimenter_validator: Optional[Callable[[str], bool]] = validate_username,
    rig_validator: Optional[Callable[[Rig], Rig]] = validate_rig_computer_name,
    use_cache: bool = True,
):
    """
    Initializes the DefaultBehaviorPicker.

    Args:
        settings: Settings containing configuration including config_library_dir. By default, attempts to rely on DefaultBehaviorPickerSettings to automatic loading from yaml files
        launcher: The launcher instance for managing experiment execution
        ui_helper: Helper for user interface interactions. If None, uses launcher's ui_helper. Defaults to None
        experimenter_validator: Function to validate the experimenter's username. If None, no validation is performed. Defaults to validate_username
        rig_validator: Function to validate the rig configuration. If None, no validation is performed. Defaults to validate_rig_computer_name
        use_cache: Whether to use caching for selections. Defaults to True
    """
    self._launcher = launcher
    self._ui_helper = launcher.ui_helper if ui_helper is None else ui_helper
    self._settings = settings
    self._ensure_directories()
    self._experimenter_validator = experimenter_validator
    self._rig_validator = rig_validator
    self._trainer_state: Optional[TrainerState] = None
    self._session: Optional[Session] = None
    self._cache_manager = CacheManager.get_instance()
    self._use_cache = use_cache

ui_helper property

ui_helper: IUiHelper

Retrieves the registered UI helper.

Returns:

Name	Type	Description
UiHelper	IUiHelper	The registered UI helper

Raises:

Type	Description
ValueError	If no UI helper is registered

trainer_state property

trainer_state: TrainerState

Returns the current trainer state.

Returns:

Name	Type	Description
TrainerState	TrainerState	The current trainer state

Raises:

Type	Description
ValueError	If the trainer state is not set

session_directory property

session_directory: Path

Returns the directory path for the current session.

session property

session: Session

Returns the current session model.

config_library_dir property

config_library_dir: Path

Returns the path to the configuration library directory.

Returns:

Name	Type	Description
Path	Path	The configuration library directory

rig_dir property

rig_dir: Path

Returns the path to the rig configuration directory.

Returns:

Name	Type	Description
Path	Path	The rig configuration directory

subject_dir property

subject_dir: Path

Returns the path to the subject configuration directory.

Returns:

Name	Type	Description
Path	Path	The subject configuration directory

task_dir property

task_dir: Path

Returns the path to the task configuration directory.

Returns:

Name	Type	Description
Path	Path	The task configuration directory

pick_rig

pick_rig(model: Type[TRig]) -> TRig

Prompts the user to select a rig configuration file.

Searches for available rig configuration files and either automatically selects a single file or prompts the user to choose from multiple options.

Parameters:

Name	Type	Description	Default
model	Type[TRig]	The rig model type to validate against	required

Returns:

Name	Type	Description
TRig	TRig	The selected rig configuration

Raises:

Type	Description
ValueError	If no rig configuration files are found or an invalid choice is made

Source code in src/clabe/pickers/default_behavior.py

def pick_rig(self, model: Type[TRig]) -> TRig:
    """
    Prompts the user to select a rig configuration file.

    Searches for available rig configuration files and either automatically
    selects a single file or prompts the user to choose from multiple options.

    Args:
        model: The rig model type to validate against

    Returns:
        TRig: The selected rig configuration

    Raises:
        ValueError: If no rig configuration files are found or an invalid choice is made
    """
    rig: TRig | None = None
    rig_path: str | None = None

    # Check cache for previously used rigs
    if self._use_cache:
        cache = self._cache_manager.try_get_cache(model.__name__)
    else:
        cache = None

    if cache:
        cache.sort()
        rig_path = self.ui_helper.prompt_pick_from_list(
            cache,
            prompt=f"Choose a rig for {model.__name__}:",
            allow_0_as_none=True,
            zero_label="Select from library",
        )
        if rig_path is not None:
            rig = self._load_rig_from_path(Path(rig_path), model)

    # Prompt user to select a rig if not already selected
    while rig_path is None:
        available_rigs = glob.glob(os.path.join(self.rig_dir, "*.json"))
        # We raise if no rigs are found to prevent an infinite loop
        if len(available_rigs) == 0:
            logger.error("No rig config files found.")
            raise ValueError("No rig config files found.")
        # Use the single available rig config file
        elif len(available_rigs) == 1:
            logger.info("Found a single rig config file. Using %s.", {available_rigs[0]})
            rig_path = available_rigs[0]
            rig = model_from_json_file(rig_path, model)
        else:
            rig_path = self.ui_helper.prompt_pick_from_list(
                available_rigs, prompt=f"Choose a rig for {model.__name__}:"
            )
            if rig_path is not None:
                rig = self._load_rig_from_path(Path(rig_path), model)
    assert rig_path is not None
    assert rig is not None
    if self._rig_validator:
        rig = self._rig_validator(rig)
    # Add the selected rig path to the cache
    self._cache_manager.add_to_cache("rigs", rig_path)
    return rig

pick_session

pick_session(model: Type[TSession] = Session) -> TSession

Prompts the user to select or create a session configuration.

Collects experimenter information, subject selection, and session notes to create a new session configuration with appropriate metadata.

Parameters:

Name	Type	Description	Default
model	Type[TSession]	The session model type to instantiate. Defaults to Session	Session

Returns:

Name	Type	Description
TSession	TSession	The created or selected session configuration

Source code in src/clabe/pickers/default_behavior.py

def pick_session(self, model: Type[TSession] = Session) -> TSession:
    """
    Prompts the user to select or create a session configuration.

    Collects experimenter information, subject selection, and session notes
    to create a new session configuration with appropriate metadata.

    Args:
        model: The session model type to instantiate. Defaults to Session

    Returns:
        TSession: The created or selected session configuration
    """

    experimenter = self.prompt_experimenter(strict=True)
    subject = self.choose_subject(self.subject_dir)

    if not (self.subject_dir / subject).exists():
        logger.info("Directory for subject %s does not exist. Creating a new one.", subject)
        os.makedirs(self.subject_dir / subject)

    notes = self.ui_helper.prompt_text("Enter notes: ")
    session = model(
        subject=subject,
        notes=notes,
        experimenter=experimenter if experimenter is not None else [],
        commit_hash=self._launcher.repository.head.commit.hexsha,
        allow_dirty_repo=self._launcher.settings.debug_mode or self._launcher.settings.allow_dirty,
        skip_hardware_validation=self._launcher.settings.skip_hardware_validation,
    )
    self._session = session
    return session

pick_task

pick_task(model: Type[TTask]) -> TTask

Prompts the user to select or create a task configuration.

Attempts to load task in the following order: 1. From CLI if already set 2. From subject-specific folder 3. From user selection in task library

Parameters:

Name	Type	Description	Default
model	Type[TTask]	The task model type to validate against	required

Returns:

Name	Type	Description
TTask	TTask	The created or selected task configuration

Raises:

Type	Description
ValueError	If no valid task file is found

Source code in src/clabe/pickers/default_behavior.py

def pick_task(self, model: Type[TTask]) -> TTask:
    """
    Prompts the user to select or create a task configuration.

    Attempts to load task in the following order:
    1. From CLI if already set
    2. From subject-specific folder
    3. From user selection in task library

    Args:
        model: The task model type to validate against

    Returns:
        TTask: The created or selected task configuration

    Raises:
        ValueError: If no valid task file is found
    """
    task: Optional[TTask] = None
    if self._session is None:
        raise ValueError("Session must be picked (pick_session) before picking task.")

    try:
        f = self.subject_dir / self._session.subject / (ByAnimalFiles.TASK.value + ".json")
        logger.info("Attempting to load task from subject folder: %s", f)
        task = model_from_json_file(f, model)
    except (ValueError, FileNotFoundError, pydantic.ValidationError) as e:
        logger.warning("Failed to find a valid task file. %s", e)
    else:
        logger.info("Found a valid task file in subject folder!")
        _is_manual = not self.ui_helper.prompt_yes_no_question("Would you like to use this task?")
        if not _is_manual:
            if task is not None:
                return task
            else:
                logger.error("No valid task file found in subject folder.")
                raise ValueError("No valid task file found.")
        else:
            task = None

    # If not found, we prompt the user to choose/enter a task file
    while task is None:
        try:
            _path = Path(os.path.join(self.config_library_dir, self.task_dir))
            available_files = glob.glob(os.path.join(_path, "*.json"))
            if len(available_files) == 0:
                break
            path = self.ui_helper.prompt_pick_from_list(
                available_files, prompt=f"Choose a task for {model.__name__}:"
            )
            if not isinstance(path, str):
                raise ValueError("Invalid choice.")
            if not os.path.isfile(path):
                raise FileNotFoundError(f"File not found: {path}")
            task = model_from_json_file(path, model)
            logger.info("User entered: %s.", path)
        except pydantic.ValidationError as e:
            logger.error("Failed to validate pydantic model. Try again. %s", e)
        except (ValueError, FileNotFoundError) as e:
            logger.info("Invalid choice. Try again. %s", e)
    if task is None:
        logger.error("No task file found.")
        raise ValueError("No task file found.")

    return task

pick_trainer_state

pick_trainer_state(
    task_model: Type[TTask],
) -> tuple[TrainerState, TTask]

Prompts the user to select or create a trainer state configuration.

Attempts to load trainer state in the following order: 1. If task already exists in launcher, will return an empty TrainerState 2. From subject-specific folder

It will launcher.set_task if the deserialized TrainerState is valid.

Parameters:

Name	Type	Description	Default
task_model	Type[TTask]	The task model type to validate against	required

Returns:

Type	Description
tuple[TrainerState, TTask]	tuple[TrainerState, TTask]: The deserialized TrainerState object and validated task

Raises:

Type	Description
ValueError	If no valid task file is found or session is not set

Source code in src/clabe/pickers/default_behavior.py

def pick_trainer_state(self, task_model: Type[TTask]) -> tuple[TrainerState, TTask]:
    """
    Prompts the user to select or create a trainer state configuration.

    Attempts to load trainer state in the following order:
    1. If task already exists in launcher, will return an empty TrainerState
    2. From subject-specific folder

    It will launcher.set_task if the deserialized TrainerState is valid.

    Args:
        task_model: The task model type to validate against

    Returns:
        tuple[TrainerState, TTask]: The deserialized TrainerState object and validated task

    Raises:
        ValueError: If no valid task file is found or session is not set
    """

    if self._session is None:
        raise ValueError("Session must be picked (pick_session) before picking trainer state.")
    try:
        f = self.subject_dir / self._session.subject / (ByAnimalFiles.TRAINER_STATE.value + ".json")
        logger.info("Attempting to load trainer state from subject folder: %s", f)
        trainer_state = model_from_json_file(f, TrainerState)
        if trainer_state.stage is None:
            raise ValueError("Trainer state stage is None, cannot use this trainer state.")
    except (ValueError, FileNotFoundError, pydantic.ValidationError) as e:
        logger.error("Failed to find a valid task file. %s", e)
        raise
    else:
        self._trainer_state = trainer_state

    if not self._trainer_state.is_on_curriculum:
        logging.warning("Deserialized TrainerState is NOT on curriculum.")

    assert self._trainer_state.stage is not None
    return (
        self.trainer_state,
        task_model.model_validate_json(self.trainer_state.stage.task.model_dump_json()),
    )

choose_subject

choose_subject(directory: str | PathLike) -> str

Prompts the user to select or manually enter a subject name.

Allows the user to either type a new subject name or select from existing subject directories.

Parameters:

Name	Type	Description	Default
directory	str | PathLike	Path to the directory containing subject folders	required

Returns:

Name	Type	Description
str	str	The selected or entered subject name.

Example

# Choose a subject from the subjects directory
subject = picker.choose_subject("Subjects")

Source code in src/clabe/pickers/default_behavior.py

def choose_subject(self, directory: str | os.PathLike) -> str:
    """
    Prompts the user to select or manually enter a subject name.

    Allows the user to either type a new subject name or select from
    existing subject directories.

    Args:
        directory: Path to the directory containing subject folders

    Returns:
        str: The selected or entered subject name.

    Example:
        ```python
        # Choose a subject from the subjects directory
        subject = picker.choose_subject("Subjects")
        ```
    """
    if self._use_cache:
        subjects = self._cache_manager.try_get_cache("subjects")
    else:
        subjects = None
    if subjects:
        subjects.sort()
        subject = self.ui_helper.prompt_pick_from_list(
            subjects,
            prompt="Choose a subject:",
            allow_0_as_none=True,
            zero_label="Enter manually",
        )
    else:
        subject = None

    while subject is None:
        subject = self.ui_helper.input("Enter subject name: ")
        if subject == "":
            subject = None
    self._cache_manager.add_to_cache("subjects", subject)
    return subject

prompt_experimenter

prompt_experimenter(
    strict: bool = True,
) -> Optional[List[str]]

Prompts the user to enter the experimenter's name(s).

Accepts multiple experimenter names separated by commas or spaces. Validates names using the configured validator function if provided.

Parameters:

Name	Type	Description	Default
strict	bool	Whether to enforce non-empty input. Defaults to True	True

Returns:

Type	Description
Optional[List[str]]	Optional[List[str]]: List of experimenter names

Example

# Prompt for experimenter with validation
names = picker.prompt_experimenter(strict=True)
print("Experimenters:", names)

Source code in src/clabe/pickers/default_behavior.py

def prompt_experimenter(self, strict: bool = True) -> Optional[List[str]]:
    """
    Prompts the user to enter the experimenter's name(s).

    Accepts multiple experimenter names separated by commas or spaces.
    Validates names using the configured validator function if provided.

    Args:
        strict: Whether to enforce non-empty input. Defaults to True

    Returns:
        Optional[List[str]]: List of experimenter names

    Example:
        ```python
        # Prompt for experimenter with validation
        names = picker.prompt_experimenter(strict=True)
        print("Experimenters:", names)
        ```
    """
    if self._use_cache:
        experimenters_cache = self._cache_manager.try_get_cache("experimenters")
    else:
        experimenters_cache = None
    experimenter: Optional[List[str]] = None
    _picked: str | None = None
    while experimenter is None:
        if experimenters_cache:
            experimenters_cache.sort()
            _picked = self.ui_helper.prompt_pick_from_list(
                experimenters_cache,
                prompt="Choose an experimenter:",
                allow_0_as_none=True,
                zero_label="Enter manually",
            )
        if _picked is None:
            _input = self.ui_helper.prompt_text("Experimenter name: ")
        else:
            _input = _picked
        experimenter = _input.replace(",", " ").split()
        if strict & (len(experimenter) == 0):
            logger.info("Experimenter name is not valid. Try again.")
            experimenter = None
        else:
            if self._experimenter_validator:
                for name in experimenter:
                    if not self._experimenter_validator(name):
                        logger.warning("Experimenter name: %s, is not valid. Try again", name)
                        experimenter = None
                        break
    self._cache_manager.add_to_cache("experimenters", ",".join(experimenter))
    return experimenter

dump_model

dump_model(
    model: Union[Rig, Task, TrainerState],
) -> Optional[Path]

Saves the provided model to the appropriate configuration file.

Parameters:

Name	Type	Description	Default
model	Union[Rig, Task, TrainerState]	The model instance to save	required

Returns:

Type	Description
Optional[Path]	Optional[Path]: The path to the saved model file, or None if not saved

Source code in src/clabe/pickers/default_behavior.py

def dump_model(
    self,
    model: Union[Rig, Task, TrainerState],
) -> Optional[Path]:
    """
    Saves the provided model to the appropriate configuration file.

    Args:
        model: The model instance to save

    Returns:
        Optional[Path]: The path to the saved model file, or None if not saved
    """

    path: Path
    if isinstance(model, Rig):
        path = self.rig_dir / ("rig.json")
    elif isinstance(model, Task):
        if self._session is None:
            raise ValueError("Session must be picked (pick_session) before dumping task.")
        path = Path(self.subject_dir) / self._session.subject / (ByAnimalFiles.TASK.value + ".json")
    elif isinstance(model, TrainerState):
        if self._session is None:
            raise ValueError("Session must be picked (pick_session) before dumping trainer state.")
        path = Path(self.subject_dir) / self._session.subject / (ByAnimalFiles.TRAINER_STATE.value + ".json")
    else:
        raise ValueError("Model type not supported for dumping.")

    os.makedirs(path.parent, exist_ok=True)
    if path.exists():
        overwrite = self.ui_helper.prompt_yes_no_question(f"File {path} already exists. Overwrite?")
        if not overwrite:
            logger.info("User chose not to overwrite the existing file: %s", path)
            return None
    with open(path, "w", encoding="utf-8") as f:
        f.write(model.model_dump_json(indent=2))
        logger.info("Saved model to %s", path)
    return path