pickers.dataverse¶

DataverseSuggestion ¶

Bases: BaseModel

Internal representation of a suggestion entry in Dataverse.

validate_trainer_state `classmethod` ¶

validate_trainer_state(value)

Validate and convert the trainer_state field from a JSON string to a TrainerState object.

Source code in src\clabe\pickers\dataverse.py

@field_validator("trainer_state", mode="before")
@classmethod
def validate_trainer_state(cls, value):
    """
    Validate and convert the trainer_state field from a JSON string to a TrainerState object.
    """
    if value is None:
        return value
    if isinstance(value, str):
        return TrainerState.model_validate_json(value)
    return value

from_request_output `classmethod` ¶

from_request_output(
    subject: str, request_output: dict
) -> DataverseSuggestion

Create a _Suggestion instance from a dictionary of data.

Source code in src\clabe\pickers\dataverse.py

@classmethod
def from_request_output(cls, subject: str, request_output: dict) -> "DataverseSuggestion":
    """
    Create a _Suggestion instance from a dictionary of data.
    """
    trainer_state = request_output.get("aibs_trainer_state", None)
    trainer_state = TrainerState.model_validate_json(cls._strip_html(trainer_state)) if trainer_state else None
    return cls(
        subject_id=subject,
        trainer_state=trainer_state,
        task_name=request_output.get("aibs_task_name", None),
        stage_name=request_output.get("aibs_stage_name", None),
        modified_on=request_output.get("modifiedon", None),
        created_on=request_output.get("createdon", None),
    )

from_trainer_state `classmethod` ¶

from_trainer_state(
    subject: str, trainer_state: TrainerState
) -> DataverseSuggestion

Create a _Suggestion instance from a TrainerState object.

Source code in src\clabe\pickers\dataverse.py

@classmethod
def from_trainer_state(cls, subject: str, trainer_state: TrainerState) -> "DataverseSuggestion":
    """
    Create a _Suggestion instance from a TrainerState object.
    """
    if trainer_state is None:
        raise ValueError("trainer_state cannot be None")
    if trainer_state.stage is None:
        raise ValueError("trainer_state.stage cannot be None")
    return cls(
        subject_id=subject,
        trainer_state=trainer_state,
        task_name=trainer_state.stage.task.name,
        stage_name=trainer_state.stage.name,
    )

DataversePicker ¶

DataversePicker(
    *,
    dataverse_client: Optional[_DataverseRestClient] = None,
    settings: DefaultBehaviorPickerSettings,
    ui_helper: Optional[UiHelper] = None,
    experimenter_validator: Optional[
        Callable[[str], bool]
    ] = validate_aind_username,
)

Bases: DefaultBehaviorPicker, Generic[TRig, TSession, TTaskLogic]

Picker that integrates with Dataverse to fetch and push trainer state suggestions.

Initializes the DataversePicker.

Parameters:

Name	Type	Description	Default
`dataverse_client`	`Optional[_DataverseRestClient]`	Optional Dataverse REST client for making API calls. If not provided, a new client will be created using settings from KeePass.	`None`
`settings`	`DefaultBehaviorPickerSettings`	Settings containing configuration including config_library_dir	required
`ui_helper`	`Optional[UiHelper]`	Helper for user interface interactions	`None`
`experimenter_validator`	`Optional[Callable[[str], bool]]`	Function to validate the experimenter's username. If None, no validation is performed	`validate_aind_username`

Source code in src\clabe\pickers\dataverse.py

def __init__(
    self,
    *,
    dataverse_client: Optional[_DataverseRestClient] = None,
    settings: DefaultBehaviorPickerSettings,
    ui_helper: Optional[ui.UiHelper] = None,
    experimenter_validator: Optional[Callable[[str], bool]] = validate_aind_username,
):
    """
    Initializes the DataversePicker.

    Args:
        dataverse_client: Optional Dataverse REST client for making API calls. If not provided, a new client will be created using settings from KeePass.
        settings: Settings containing configuration including config_library_dir
        ui_helper: Helper for user interface interactions
        experimenter_validator: Function to validate the experimenter's username. If None, no validation is performed
    """
    super().__init__(settings=settings, ui_helper=ui_helper, experimenter_validator=experimenter_validator)
    self._dataverse_client = (
        dataverse_client
        if dataverse_client is not None
        else _DataverseRestClient(_DataverseRestClientSettings.from_keepass())
    )
    self._dataverse_suggestion: Optional[DataverseSuggestion] = None

has_ui_helper `property` ¶

has_ui_helper: bool

Checks if a UI helper is registered.

Returns:

Name	Type	Description
`bool`	`bool`	True if a UI helper is registered, False otherwise

ui_helper `property` ¶

ui_helper: UiHelper

Retrieves the registered UI helper.

Returns:

Name	Type	Description
`DefaultUIHelper`	`UiHelper`	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.

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_logic_dir `property` ¶

task_logic_dir: Path

Returns the path to the task logic configuration directory.

Returns:

Name	Type	Description
`Path`	`Path`	The task logic configuration directory.

pick_trainer_state ¶

pick_trainer_state(
    launcher: Launcher[TRig, TSession, TTaskLogic],
) -> TrainerState

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

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

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

Returns:

Name	Type	Description
`TrainerState`	`TrainerState`	The deserialized TrainerState object.

Raises:

Type	Description
`ValueError`	If no valid task logic file is found.

Source code in src\clabe\pickers\dataverse.py

def pick_trainer_state(self, launcher: Launcher[TRig, TSession, TTaskLogic]) -> TrainerState:
    """
    Prompts the user to select or create a trainer state configuration.

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

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

    Returns:
        TrainerState: The deserialized TrainerState object.

    Raises:
        ValueError: If no valid task logic file is found.
    """
    if (launcher.get_task_logic()) is not None:
        logger.debug("Task logic already set in launcher. Cannot inject a trainer state.")
        self._trainer_state = TrainerState(curriculum=None, stage=None, is_on_curriculum=False)
    else:
        if launcher.subject is None:
            logger.error("No subject set in launcher. Cannot load trainer state.")
            raise ValueError("No subject set in launcher.")
        task_logic_name = launcher.get_task_logic_model().model_fields["name"].default
        if not task_logic_name:
            raise ValueError("Task logic model does not have a default name.")
        try:
            logger.debug("Attempting to load trainer state dataverse")
            last_suggestions = _get_last_suggestions(self._dataverse_client, launcher.subject, task_logic_name, 1)
        except requests.exceptions.HTTPError as e:
            logger.error("Failed to fetch suggestions from Dataverse: %s", e)
            raise
        except pydantic.ValidationError as e:
            logger.error("Failed to validate suggestion from Dataverse: %s", e)
            raise
        if len(last_suggestions) == 0:
            raise ValueError(
                f"No valid suggestions found in Dataverse for subject {launcher.subject} with task {task_logic_name}."
            )

        _dataverse_suggestion = last_suggestions[0]

        assert _dataverse_suggestion is not None
        if _dataverse_suggestion.trainer_state is None:
            raise ValueError("No trainer state found in the latest suggestion.")
        if _dataverse_suggestion.trainer_state.stage is None:
            raise ValueError("No stage found in the latest suggestion's trainer state.")
        self._dataverse_suggestion = _dataverse_suggestion
        self._trainer_state = _dataverse_suggestion.trainer_state
        launcher.set_task_logic(_dataverse_suggestion.trainer_state.stage.task)

    assert self._trainer_state is not None
    if not self._trainer_state.is_on_curriculum:
        logging.warning("Deserialized TrainerState is NOT on curriculum.")
    self._sync_session_metadata(launcher)
    return self.trainer_state

push_new_suggestion ¶

push_new_suggestion(
    launcher: Launcher[TRig, TSession, TTaskLogic],
    trainer_state: TrainerState,
) -> None

Pushes a new suggestion to Dataverse for the current subject in the launcher. Args: launcher: The Launcher instance containing the current session and subject information. trainer_state: The TrainerState object to be pushed as a new suggestion.

Source code in src\clabe\pickers\dataverse.py

def push_new_suggestion(self, launcher: Launcher[TRig, TSession, TTaskLogic], trainer_state: TrainerState) -> None:
    """
    Pushes a new suggestion to Dataverse for the current subject in the launcher.
    Args:
        launcher: The Launcher instance containing the current session and subject information.
        trainer_state: The TrainerState object to be pushed as a new suggestion.
    """
    logger.info("Pushing new suggestion to Dataverse for subject %s", launcher.get_session(strict=True).subject)
    _append_suggestion(self._dataverse_client, launcher.get_session(strict=True).subject, trainer_state)

register_ui_helper ¶

register_ui_helper(ui_helper: UiHelper) -> Self

Registers a UI helper with the picker.

Associates a UI helper instance with this picker for user interactions.

Parameters:

Name	Type	Description	Default
`ui_helper`	`UiHelper`	The UI helper to register	required

Returns:

Name	Type	Description
`Self`	`Self`	The picker instance for method chaining

Raises:

Type	Description
`ValueError`	If a UI helper is already registered

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

def register_ui_helper(self, ui_helper: ui.UiHelper) -> Self:
    """
    Registers a UI helper with the picker.

    Associates a UI helper instance with this picker for user interactions.

    Args:
        ui_helper: The UI helper to register

    Returns:
        Self: The picker instance for method chaining

    Raises:
        ValueError: If a UI helper is already registered
    """
    if self._ui_helper is None:
        self._ui_helper = ui_helper
    else:
        raise ValueError("UI Helper is already registered")
    return self

initialize ¶

initialize(
    launcher: Launcher[TRig, TSession, TTaskLogic],
) -> None

Initializes the picker by creating required directories if needed.

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

def initialize(self, launcher: Launcher[TRig, TSession, TTaskLogic]) -> None:
    """
    Initializes the picker by creating required directories if needed.
    """
    self._launcher = launcher
    self._ui_helper = launcher.ui_helper
    self._ensure_directories(launcher)

pick_rig ¶

pick_rig(
    launcher: Launcher[TRig, TSession, TTaskLogic],
) -> 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.

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, launcher: Launcher[TRig, TSession, TTaskLogic]) -> 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.

    Returns:
        TRig: The selected rig configuration.

    Raises:
        ValueError: If no rig configuration files are found or an invalid choice is made.
    """
    rig = launcher.get_rig()
    if rig is not None:
        logger.info("Rig already set in launcher. Using existing rig.")
        return rig
    available_rigs = glob.glob(os.path.join(self.rig_dir, "*.json"))
    if len(available_rigs) == 0:
        logger.error("No rig config files found.")
        raise ValueError("No rig config files found.")
    elif len(available_rigs) == 1:
        logger.info("Found a single rig config file. Using %s.", {available_rigs[0]})
        rig = model_from_json_file(available_rigs[0], launcher.get_rig_model())
        launcher.set_rig(rig)
        return rig
    else:
        while True:
            try:
                path = self.ui_helper.prompt_pick_from_list(available_rigs, prompt="Choose a rig:")
                if not isinstance(path, str):
                    raise ValueError("Invalid choice.")
                rig = model_from_json_file(path, launcher.get_rig_model())
                logger.info("Using %s.", path)
                launcher.set_rig(rig)
                return rig
            except pydantic.ValidationError as e:
                logger.error("Failed to validate pydantic model. Try again. %s", e)
            except ValueError as e:
                logger.info("Invalid choice. Try again. %s", e)

pick_session ¶

pick_session(
    launcher: Launcher[TRig, TSession, TTaskLogic],
) -> 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.

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, launcher: Launcher[TRig, TSession, TTaskLogic]) -> 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.

    Returns:
        TSession: The created or selected session configuration.
    """
    if (session := launcher.get_session()) is not None:
        logger.info("Session already set in launcher. Using existing session.")
        return session

    experimenter = self.prompt_experimenter(strict=True)
    if launcher.subject is not None:
        logger.info("Subject provided via CLABE: %s", launcher.subject)
        subject = launcher.subject
    else:
        subject = self.choose_subject(self.subject_dir)
        launcher.subject = subject

        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 = launcher.get_session_model()(
        experiment="",  # Will be set later
        root_path=str(Path(launcher.settings.data_dir).resolve() / subject),
        subject=subject,
        notes=notes,
        experimenter=experimenter if experimenter is not None else [],
        commit_hash=launcher.repository.head.commit.hexsha,
        allow_dirty_repo=launcher.settings.debug_mode or launcher.settings.allow_dirty,
        skip_hardware_validation=launcher.settings.skip_hardware_validation,
        experiment_version="",  # Will be set later
    )
    launcher.set_session(session)
    return session

pick_task_logic ¶

pick_task_logic(
    launcher: Launcher[TRig, TSession, TTaskLogic],
) -> TTaskLogic

Prompts the user to select or create a task logic configuration.

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

Returns:

Name	Type	Description
`TTaskLogic`	`TTaskLogic`	The created or selected task logic configuration.

Raises:

Type	Description
`ValueError`	If no valid task logic file is found.

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

def pick_task_logic(self, launcher: Launcher[TRig, TSession, TTaskLogic]) -> TTaskLogic:
    """
    Prompts the user to select or create a task logic configuration.

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

    Returns:
        TTaskLogic: The created or selected task logic configuration.

    Raises:
        ValueError: If no valid task logic file is found.
    """
    if (task_logic := launcher.get_task_logic()) is not None:
        logger.info("Task logic already set in launcher. Using existing task logic.")
        self._sync_session_metadata(launcher)
        return task_logic

    # Else, we check inside the subject folder for an existing task file
    try:
        if launcher.subject is None:
            logger.error("No subject set in launcher. Cannot load task logic.")
            raise ValueError("No subject set in launcher.")
        f = self.subject_dir / launcher.subject / (ByAnimalFiles.TASK_LOGIC.value + ".json")
        logger.info("Attempting to load task logic from subject folder: %s", f)
        task_logic = model_from_json_file(f, launcher.get_task_logic_model())
    except (ValueError, FileNotFoundError, pydantic.ValidationError) as e:
        logger.warning("Failed to find a valid task logic file. %s", e)
    else:
        logger.info("Found a valid task logic file in subject folder!")
        _is_manual = not self.ui_helper.prompt_yes_no_question("Would you like to use this task logic?")
        if not _is_manual:
            if task_logic is not None:
                launcher.set_task_logic(task_logic)
                return task_logic
            else:
                logger.error("No valid task logic file found in subject folder.")
                raise ValueError("No valid task logic file found.")
        else:
            task_logic = None

    # If not found, we prompt the user to choose/enter a task logic file
    while task_logic is None:
        try:
            _path = Path(os.path.join(self.config_library_dir, self.task_logic_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="Choose a task logic:")
            if not isinstance(path, str):
                raise ValueError("Invalid choice.")
            if not os.path.isfile(path):
                raise FileNotFoundError(f"File not found: {path}")
            task_logic = model_from_json_file(path, launcher.get_task_logic_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_logic is None:
        logger.error("No task logic file found.")
        raise ValueError("No task logic file found.")

    launcher.set_task_logic(task_logic)
    self._sync_session_metadata(launcher)
    return task_logic

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")
        ```
    """
    subject = None
    while subject is None:
        subject = self.ui_helper.input("Enter subject name: ")
        if subject == "":
            subject = self.ui_helper.prompt_pick_from_list(
                [
                    os.path.basename(folder)
                    for folder in os.listdir(directory)
                    if os.path.isdir(os.path.join(directory, folder))
                ],
                prompt="Choose a subject:",
                allow_0_as_none=True,
            )
        else:
            return 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	`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

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

    Example:
        ```python
        # Prompt for experimenter with validation
        names = picker.prompt_experimenter(strict=True)
        print("Experimenters:", names)
        ```
    """
    experimenter: Optional[List[str]] = None
    while experimenter is None:
        _user_input = self.ui_helper.prompt_text("Experimenter name: ")
        experimenter = _user_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
    return experimenter

dump_model ¶

dump_model(
    launcher: Launcher[TRig, TSession, TTaskLogic],
    model: Union[
        AindBehaviorRigModel,
        AindBehaviorTaskLogicModel,
        TrainerState,
    ],
) -> Optional[Path]

Saves the provided model to the appropriate configuration file.

Parameters:

Name	Type	Description	Default
`launcher`	`Launcher[TRig, TSession, TTaskLogic]`	The launcher instance managing the experiment.	required
`model`	`Union[AindBehaviorRigModel, AindBehaviorTaskLogicModel, 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,
    launcher: Launcher[TRig, TSession, TTaskLogic],
    model: Union[AindBehaviorRigModel, AindBehaviorTaskLogicModel, TrainerState],
) -> Optional[Path]:
    """
    Saves the provided model to the appropriate configuration file.

    Args:
        launcher: The launcher instance managing the experiment.
        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, AindBehaviorRigModel):
        path = self.rig_dir / ("rig.json")
    elif isinstance(model, AindBehaviorTaskLogicModel):
        if launcher.subject is None:
            raise ValueError("No subject set in launcher. Cannot dump task logic.")
        path = Path(self.subject_dir) / launcher.subject / (ByAnimalFiles.TASK_LOGIC.value + ".json")
    elif isinstance(model, TrainerState):
        if launcher.subject is None:
            raise ValueError("No subject set in launcher. Cannot dump trainer state.")
        path = Path(self.subject_dir) / launcher.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

pickers.dataverse¶

DataverseSuggestion ¶

validate_trainer_state classmethod ¶

from_request_output classmethod ¶

from_trainer_state classmethod ¶

DataversePicker ¶

has_ui_helper property ¶

ui_helper property ¶

trainer_state property ¶

config_library_dir property ¶

rig_dir property ¶

subject_dir property ¶

task_logic_dir property ¶

pick_trainer_state ¶

push_new_suggestion ¶

register_ui_helper ¶

initialize ¶

pick_rig ¶

pick_session ¶

pick_task_logic ¶

choose_subject ¶

prompt_experimenter ¶

dump_model ¶

validate_trainer_state `classmethod` ¶

from_request_output `classmethod` ¶

from_trainer_state `classmethod` ¶

has_ui_helper `property` ¶

ui_helper `property` ¶

trainer_state `property` ¶

config_library_dir `property` ¶

rig_dir `property` ¶

subject_dir `property` ¶

task_logic_dir `property` ¶