def _set_globals_based_on_role()

in pyrit/orchestrator/multi_turn/multi_turn_orchestrator.py [0:0]

• 18 lines of code
• 9 McCabe index (conditional complexity)

    def _set_globals_based_on_role(self, last_message: PromptRequestPiece):
        """Sets the global variables of self._last_prepended_user_message and self._last_prepended_assistant_message
        based on the role of the last message in the prepended conversation.
        """
        # There is specific handling per orchestrator depending on the last message
        if last_message.role == "user":
            self._last_prepended_user_message = last_message.converted_value
        elif last_message.role == "assistant":
            # Get scores for the last assistant message based off of the original id
            self._last_prepended_assistant_message_scores = self._memory.get_scores_by_prompt_ids(
                prompt_request_response_ids=[str(last_message.original_prompt_id)]
            )

            # Do not set last user message if there are no scores for the last assistant message
            if not self._last_prepended_assistant_message_scores:
                return

            # Check assumption that there will be a user message preceding the assistant message
            if (
                len(self._prepended_conversation) > 1
                and self._prepended_conversation[-2].request_pieces[0].role == "user"
            ):
                self._last_prepended_user_message = self._prepended_conversation[-2].get_value()
            else:
                raise ValueError(
                    "There must be a user message preceding the assistant message in prepended conversations."
                )