Tool-based Data Segregation: a Prevention Strategy against Prompt Injection in AI Agents

Contents:

logs.7z - compacted log files.

executions.7z - compacted responses from the models.

code.7z - the java code used in this experiment.

models.txt - text file containing the models used in this experiment.

prompt-injection-results.ods - sheets with our results.

# Project Structure

├── executions/│   └── execution_number/│       ├── REFERENCE.txt│       │   └── Reference responses│       ├── INJECTION-1.txt│       │   └── Prompt injection responses without tools│       ├── INJECTION-2.txt│       │   └── Prompt injection responses with tools│       ├── INJECTION-3.txt│       │   └── Prompt injection responses with tools and system messages│       └── INJECTION-4.txt│           └── Prompt injection responses with tools and tool instructions│├── logs/│   └── execution_number/│       └── results_for_a_given_model/│           └── openai.log│               └── Execution logs for all five scenarios│├── promptinjection-tools/│   └── Java project used in this experiment│└── models.txt    └── List of models executed by the experiment

# Running the Experiment

From the promptinjection-tools directory, run:

mvn clean package

mvn exec:java -D"exec.mainClass"="com.promptinjection.tooldefense.MainExperiment" -Dexec.args="../models.txt number_of_executions"

Where:

- ../models.txt is the file containing the list of models to execute.- number_of_executions is the number of executions to perform.

# Extending Support to Other Models

To add support for additional models, first insert the model name in ../models.txt providers or model families, extend the method getModelFor in the class:

```textcom.promptinjection.tooldefense.ChatModelSelector```

Example:

public static ChatModel getModelfor(String modelName) {

boolean isGptModel =
modelName.contains("gpt");

boolean isLlamaModel =
modelName.contains("llama");

boolean isYourModelName =
modelName.contains("YOUR_MODEL_NAME");

if (isGptModel) {

buildGptModelsMap(OpenAiChatModelName.values());

return buildOpenAiChatmodel(
gptModels.get(modelName.toLowerCase()));

} else if (isLlamaModel) {

return buildOllamaChatModel(modelName);

} else if (isYourModelName) {

// implement your model instantiation logic

} else {

throw new ModelNotFoundException(
"Could not resolve model for " + modelName);
}
}