Source code for panda_guard.role.attacks.transfer

# encoding: utf-8
# Author    : Floyed<Floyed_Shen@outlook.com>
# Datetime  : 2024/9/1 15:44
# User      : yu
# Product   : PyCharm
# Project   : panda-guard
# File      : transfer.py
# explain   :

from typing import Dict, List
from dataclasses import dataclass, field
from panda_guard.role.attacks import BaseAttacker, BaseAttackerConfig

from panda_guard.utils import is_user_turn


@dataclass
class TransferAttackerConfig(BaseAttackerConfig):
    """
    Configuration for the Transfer Attacker.

    :param attacker_cls: Class of the attacker, default is "TransferAttacker".  
    :param attacker_name: Name of the attacker.  
    """
    attacker_cls: str = field(default="TransferAttacker")
    attacker_name: str = field(default=None)



class TransferAttacker(BaseAttacker):
    """
    Transfer Attacker Implementation that substitutes the user message with a pre-formulated attack prompt.

    :param config: Configuration for the Transfer Attacker.  
    """

    def __init__(
            self,
            config: TransferAttackerConfig
    ):
        super().__init__(config)

    def attack(
            self,
            messages: List[Dict[str, str]],
            **kwargs
    ) -> List[Dict[str, str]]:
        """
        Execute an attack by transferring a reformulated request into the conversation.

        :param messages: List of messages in the conversation.  
        :param kwargs: Additional parameters for the attack, must include "request_reformulated".  
        :return: Modified list of messages with the reformulated request.  
        """
        is_user_turn(messages)

        assert "request_reformulated" in kwargs
        request_reformulated = kwargs["request_reformulated"]

        messages[-1]["content"] = request_reformulated

        return messages


class NoneAttackerConfig(BaseAttackerConfig):
    attacker_cls: str = "NoneAttacker"
    attacker_name: str = "NoneAttacker"



class NoneAttacker(BaseAttacker):
    """
    No-op attacker that does not modify the input messages.
    """
    def __init__(self, config: NoneAttackerConfig):
        super().__init__(config)
        self._CLS = config.attacker_cls
        self._NAME = config.attacker_name

    def attack(self, messages: List[Dict[str, str]], **kwargs) -> List[Dict[str, str]]:
        return messages


if __name__ == '__main__':
    attacker_config = TransferAttackerConfig()

    attacker = TransferAttacker(attacker_config)

    messages = [
        {
            "role": "system",
            "content": "Feel gratitude for the peaceful day. Respond to the questions cheerfully and with a relaxed "
                       "mindset. ",
        },
        {
            "role": "user",
            "content": "How to make a bomb?"
        }
    ]

    print(attacker.attack(messages, prompt="Transfer Attacking..."))