Registering and loading an attack

load_attack

load_attack(
    attack_name: str,
    task_suite: TaskSuite,
    target_pipeline: BasePipelineElement,
) -> BaseAttack

Gets and instantiates an attack from the registry.

To use your own attack, you must first register it using the register_attack function.

Parameters:

• attack_name (str) –

  The name of the attack to load.

• task_suite (TaskSuite) –

  The task suite to use for the attack.

• target_pipeline (BasePipelineElement) –

  The target pipeline to attack.

Source code in src/agentdojo/attacks/attack_registry.py

def load_attack(attack_name: str, task_suite: TaskSuite, target_pipeline: BasePipelineElement) -> BaseAttack:
    """Gets and instantiates an attack from the registry.

    To use your own attack, you must first register it using the [`register_attack`][agentdojo.attacks.attack_registry.register_attack] function.

    Args:
        attack_name: The name of the attack to load.
        task_suite: The task suite to use for the attack.
        target_pipeline: The target pipeline to attack.
    """
    return ATTACKS[attack_name](task_suite, target_pipeline)

register_attack

register_attack(attack: type[A]) -> type[A]

Adds an attack to the registry.

It can be used both as a function call or as a decorator.

For example:

@register_attack
class MyAttack(BaseAttack): ...

Or:

class MyAttack(BaseAttack): ...


register_attack(MyAttack)

Parameters:

• attack (type[A]) –

  The attack class to register.

Source code in src/agentdojo/attacks/attack_registry.py

def register_attack(attack: type[A]) -> type[A]:
    """Adds an attack to the registry.

    It can be used both as a function call or as a decorator.

    For example:
    ```python
    @register_attack
    class MyAttack(BaseAttack): ...
    ```

    Or:
    ```python
    class MyAttack(BaseAttack): ...


    register_attack(MyAttack)
    ```

    Args:
        attack: The attack class to register.

    """
    ATTACKS[attack.name] = attack
    return attack