# Agent Pipeline and Defenses

A core element of AgentDojo is the agent pipeline. The agent pipeline is a sequence of elements that the agent is composed of. In this way, elements can be composed together. Each element needs to inherit from the [`BasePipelineElement`][agentdojo.agent_pipeline.BasePipelineElement]. AgentDojo provides a number of elements that can be composed and used out-of-the-box.

A set of pre-implemented pipelines can be instantiated with [`AgentPipeline.from_config`][agentdojo.agent_pipeline.AgentPipeline.from_config].

## Base elements

Some base elements are:

  - [`AgentPipeline`][agentdojo.agent_pipeline.AgentPipeline]: executes a sequence of elements.
  - [`InitQuery`][agentdojo.agent_pipeline.InitQuery]: initializes the pipeline execution with the user's query.
  - [`SystemMessage`][agentdojo.agent_pipeline.SystemMessage]: adds a system message to the messages list.

## Function execution elements

A core element of a pipeline, is the execution of tools by using a [`FunctionsRuntime`][agentdojo.functions_runtime.FunctionsRuntime]. We provide some base components to run the tools needed by the pipeline:

  - [`ToolsExecutor`][agentdojo.agent_pipeline.ToolsExecutor]: executes the tools requested in the last [`ChatAssistantMessage`][agentdojo.types.ChatAssistantMessage], by using the [`FunctionsRuntime`][agentdojo.functions_runtime.FunctionsRuntime] passed by the previous component.
  - [`ToolsExecutionLoop`][agentdojo.agent_pipeline.ToolsExecutionLoop]: loops a sequence of pipeline elements until the last [`ChatAssistantMessage`][agentdojo.types.ChatAssistantMessage] does not require running any tool.

## LLMs

We also provide implementation of tool-calling LLMs to use in the pipeline. All of these elements call the respective LLM with the tools available in the [`FunctionsRuntime`][agentdojo.functions_runtime.FunctionsRuntime], and return the mode output to the pipeline. The output is added to the list of [`ChatMessage`][agentdojo.types.ChatMessage]s passed along the pipeline. In particular, we provide the following LLMs:

  - [`OpenAILLM`][agentdojo.agent_pipeline.OpenAILLM]: gives access to GPT models.
  - [`AnthropicLLM`][agentdojo.agent_pipeline.AnthropicLLM]: gives access to Claude Opus, Sonnet, and Haiku.
  - [`GoogleLLM`][agentdojo.agent_pipeline.GoogleLLM]: gives access to Gemini 1.5 Pro and Flash.
  - [`CohereLLM`][agentdojo.agent_pipeline.CohereLLM]: gives access to Command-R and Command-R+.
  - [`PromptingLLM`][agentdojo.agent_pipeline.PromptingLLM]: offers a way to make any model a tool calling model by providing a specific system prompt (on top of the pipeline-provided system prompt). We use this to run Llama-3 70B via TogetherAI's API.

## Combining pipeline components

Pipeline components can be combined to make a pipeline. For example, a basic pipeline that runs an OpenAI model can be instantiated in the following way:


```python title="pipeline.py"
--8<-- "pipeline.py:init_pipeline"
```

1. Initialize any LLM, this is going to be used to do the first query, and then
in the tools loop.
2. You can combine together components to create a tool execution loop. The loop is going to run as long as the [`ToolsExecutor`][agentdojo.agent_pipeline.ToolsExecutor] finds some tools to be run in the last message in the history.
3. You can also place defense components in the [`ToolsExecutionLoop`][agentdojo.agent_pipeline.ToolsExecutionLoop]!
4. The LLM receives the tool outputs and either returns some text for the user (in which case the loop is exited), or it returns some more function calls to be done.
5. Finally, all components can be combined together with [`AgentPipeline`][agentdojo.agent_pipeline.AgentPipeline]

## Running a pipeline

An [`AgentPipeline`][agentdojo.agent_pipeline.AgentPipeline] can be run by using the `query` method. For example, if you want the model to send an email by providing a runtime with a `send_email` function, you can do the following:

```python title="pipeline.py"
--8<-- "pipeline.py:run_pipeline"
```

## Creating a fully customized pipeline from scratch

You don't need to use the pre-implemented pipelines, or even an `AgentPipeline`. This is intentional, to not limit how a defense would work. The only thing that matters is that the pipeline components inherit from [`BasePipelineElement`][agentdojo.agent_pipeline.BasePipelineElement].

!!! tip
    The pipeline does not even have to be made of multiple elements. You could have a one-component pipeline that does everything.

We hope that this design is generic enough to not limit the creation of new defenses. But let us know if you find any limitations!

## Registering and running a pipeline

We are still working on providing a better way to instantiate pipelines and allow users to create and register their own pipelines as easily as it is possible with attacks and task suites. We plan to add this feature in the next couple of weeks.

Meanwhile, you can create your own pipelines and benchmark them by using the benchmarking functions provided in the [`agentdojo.benchmark`](../api/benchmark.md) module.
