# Task Suite and Tasks

Defining your own task suite and tasks enables you to create new benchmarks to extend ours, or to create a new benchmark from scratch.

You can find an example toy suite in our repo, in the `examples` [directory](https://github.com/ethz-spylab/agentdojo/tree/main/examples/counter_benchmark). This is the suite that we cover in this guide.

## Task suite

The suite operates on a specific environment type, and must be the same environment type as the one used in the
tasks. Each suite also has access to a set of tools that can be used by the agent to solve the tasks.

A new suite can be created in the following way:

```python title="counter_benchmark/suites/counter/task_suite.py"
--8<-- "counter_benchmark/suites/counter/task_suite.py:suite_creation"
```

1. [TaskSuite][agentdojo.task_suite.task_suite.TaskSuite] is a generic of the environment type. In this case, the `CounterEvironment`.
2. This is the name of the suite.
3. These are the tools that are available to the suite.
4. This is the path to the suite data. This should be a directory that contains the following files: `environment.yaml` (The data of the environment). `injection_vectors.yaml`: The injection vectors in the environment. See [below](#environment-files) for more information on these files.

### Environment files

The path `"data/suites/counter"` provided above is the path to the suite data. The path must contain the following files:

- `environment.yaml`: The data of the environment.
- `injection_vectors.yaml`: The injection vectors in the environment.

The `environment.yaml` file is used to define the initial state of the environment and should be parseable as a Pydantic Model using the environment type. For example, the `environment.yaml` file for the `CounterEnvironment` could be:

```yaml title="counter_benchmark/data/suites/counter/environment.yaml"
--8<-- "counter_benchmark/data/suites/counter/environment.yaml"
```

1. The initial counter value.
2. The description of the counter. The `{injection_counter_0}` is a placeholder for the injection vector. This is where the injection attack will be inserted while the model is benchmarked. If the benchmark is not using injection attacks, the placeholder will be replaced with the default value defined in `injection_vectors.yaml`.

The `injection_vectors.yaml` file, instead, is used to define the injection vectors in the environment. This should be parseable as a nested dictionary, where the keys are the names of the injection vectors and the values are a dict with the description of the injection vector and the a default value to use when no injection is placed. For example, the `injection_vectors.yaml` file for the `CounterEnvironment` could be:

```yaml title="counter_benchmark/data/suites/counter/injection_vectors.yaml"
--8<-- "counter_benchmark/data/suites/counter/injection_vectors.yaml"
```

### Registering tasks with the suite

After the suite is created, it can be used to register [`UserTask`][agentdojo.base_tasks.BaseUserTask]s and [`InjectionTask`][agentdojo.base_tasks.BaseInjectionTask]s. This can be done by adding the decorator [`@task_suite.register_user_task`][agentdojo.task_suite.task_suite.TaskSuite.register_user_task] or [`@task_suite.register_injection_task`][agentdojo.task_suite.task_suite.TaskSuite.register_injection_task] to the task class.

For example, to register a user task:

```python
@task_suite.register_user_task
class MyUserTask(BaseUserTask[CounterEnvironment]): ...
```

And to register an injection task:

```python
@task_suite.register_injection_task
class MyInjectionTask(BaseInjectionTask[CounterEnvironment]): ...
```

[Below](#tasks), we explain how to create user and injection tasks.

### Registering the suite

Once the suite is created, you can register and run the benchmark by using our script. You can do that with the following code:

```python title="benchmark.py"
--8<-- "counter_benchmark/benchmark.py"
```

Then you can use our `check-suites` scripts to check whether the ground truths solve the tasks, and if all tasks are injectable. Assuming you are in the `examples` directory, you can run

```bash
python -m agentdojo.scripts.check_suites \
    --benchmark-version counter_benchmark \
    -ml counter_benchmark.benchmark  # (1)!
```

1. Loads the module where the benchmark is registered, to make it available to the script.

!!! tip
    If you are creating a utility-only benchmark, then you should use the `--no-check-injectable` flag. This will prevent the script from checking whether the tasks are injectable (which of course will be the case if your benchmark has no prompt injections at all!)

Once the suite is correctly verified, you can run the benchmark with the following command:

```bash
python -m agentdojo.scripts.benchmark \
    --model gpt-4o-2024-05-13 \
    -ml counter_benchmark.benchmark \   # (1)!
    --benchmark-version counter_benchmark
```

1. Loads the module where the benchmark is registered, to make it available to the script.

## Tasks

### User tasks

User tasks are tasks that the agent can solve by interacting with the environment. They are defined by the user and can be registered to a task suite.

A user task should inherit from [`BaseUserTask`][agentdojo.base_tasks.BaseUserTask] and be generic of the environment type. The user task should define a [`PROMPT`][agentdojo.base_tasks.BaseUserTask.PROMPT] class variable that describes the task in a way that is directly understandable by the agent. Furthermore, the class should implement the [`utility`][agentdojo.base_tasks.BaseUserTask.utility] and [`ground_truth`][agentdojo.base_tasks.BaseUserTask.ground_truth] methods. `utility` is used to see whether the task was executed, and `ground_truth` is used to run the [`GroundTruthPipeline`][agentdojo.agent_pipeline.GroundTruthPipeline].

Finally, a user task may define a [`GROUND_TRUTH_OUTPUT`][agentdojo.base_tasks.BaseUserTask.GROUND_TRUTH_OUTPUT] variable that is going to be used by the [`GroundTruthPipeline`][agentdojo.agent_pipeline.GroundTruthPipeline] as the model output. This is useful if

A user task can be created in the following way:

```python title="counter_benchmark/suites/counter/user_tasks.py"
--8<-- "counter_benchmark/suites/counter/user_tasks.py:user_task_0"
```

1. The user task class name should be `UserTask<N>` where `N` is a number that identifies the task.

A user task may also define the [utility_from_traces][agentdojo.base_tasks.BaseUserTask.utility_from_traces] method to check the security of the task from the function calls traces. This is especially useful when the task does not leave a trace in the environment. Consider as an example the following task:

```python title="counter_benchmark/suites/counter/user_tasks.py"
--8<-- "counter_benchmark/suites/counter/user_tasks.py:user_task_1"
```

1. The counters should match before and after the task is executed as the same amount is incremented and decremented.
2. We can use `traces` to check the function calls that were made. In this case, we check that the first function call is an increment of 5 and the second is a decrement of 5.

### Injection tasks

Injection tasks are tasks that an adversary wants the agent to execute. An injection task should inherit from [`BaseInjectionTask`][agentdojo.base_tasks.BaseInjectionTask] and be generic of the environment type. The user task should define a [`GOAL`][agentdojo.base_tasks.BaseInjectionTask.GOAL] class variable that describes the task in a way that is directly understandable by the agent. Furthermore, the class should implement the [`security`][agentdojo.base_tasks.BaseInjectionTask.security] and [`ground_truth`][agentdojo.base_tasks.BaseInjectionTask.ground_truth] methods. `security` is used to see whether the task was executed, and `ground_truth` is used to run the [`GroundTruthPipeline`][agentdojo.agent_pipeline.GroundTruthPipeline].

Furthermore, injection tasks should define, for each function call in the ground truth, some placeholder arguments that the attacker does not beforehand. This is useful for attackers that know the names and arguments of the tools, and can build their prompt injection by inserting a sequence of tool calls that the model should do.

Consider the following example:

```python title="counter_benchmark/suites/counter/injection_tasks.py"
--8<-- "counter_benchmark/suites/counter/injection_tasks.py:injection_task_0"
```

1. The attacker does not know the value of the counter, so we use a placeholder `$value` that will the agent will (hopefully!) replace with the value returned by the previous function call.
