Apppublic

veereshdev/prompt-injection-env

sourceHugging Facemitupdated 6mo agoView on Hugging Face

0likes

App README

title: Prompt Injection Defense Environment emoji: ??? colorFrom: red colorTo: purple sdk: docker pinned: false license: mit tags:

—openenv
—security
—prompt-injection ---

Prompt Injection Defense Environment

An OpenEnv-compatible reinforcement-learning environment where an AI agent learns to detect, classify, and neutralise prompt injection attacks embedded in documents — addressing OWASP Top 10 for LLM Applications, Risk #1.

Why this matters

Prompt injection is the most critical unsolved security problem in AI today. Every RAG system, AI email assistant, and AI web browser is currently vulnerable. No standardised training environment existed — until now.

Environment Description

The agent acts as a security layer reviewing documents before they reach a downstream LLM. It must identify whether a document contains an injection, classify the attack type, and return a sanitised version with the attack removed and all legitimate content preserved.

Action Space

python

# Single-document tasks (easy, medium)
Action(
    is_injection=True,
    attack_type="direct_injection",
    sanitised_content="...cleaned document...",
    confidence=0.95,
    explanation="Found explicit override instruction at end of document."
)

# Multi-document task (hard)
Action(
    document_analyses=[
        {"is_injection": False, "attack_type": "no_injection", "sanitised_content": "...", ...},
        {"is_injection": True,  "attack_type": "semantic_injection", "sanitised_content": "...", ...},
        {"is_injection": True,  "attack_type": "multi_hop_injection", "sanitised_content": "...", ...},
    ]
)

Observation Space

Field	Type	Description
`task_id`	str	Current task identifier
`difficulty`	str	easy / medium / hard
`instruction`	str	What the agent must do
`document`	str?	Single document (easy/medium tasks)
`documents`	list?	3-document corpus (hard task)
`attack_taxonomy`	list	Valid attack type labels
`last_score`	float?	Score from previous step
`last_feedback`	str?	Feedback from previous step

Attack Taxonomy

Type	Description
`direct_injection`	Explicit override instructions, admin commands
`encoded_injection`	Attack hidden in base64 or unicode escapes
`semantic_injection`	Subtle role-reframing without obvious keywords
`multi_hop_injection`	References other documents to amplify attack
`no_injection`	Clean document, no attack present

Tasks

Task ID	Difficulty	Description	Expected Baseline Score
`task_easy`	Easy	Detect direct injection with explicit keywords	~0.85
`task_medium`	Medium	Detect encoded injection (base64/unicode)	~0.60
`task_hard`	Hard	Analyse 3-doc corpus with semantic + multi-hop attacks	~0.35

Reward Function

Dimension	Weight	Description
Detection	0.30–0.40	Correct is_injection boolean
Classification	0.25–0.30	Correct attack_type from taxonomy
Sanitisation	0.25–0.40	Injection absent from output
Preservation	0.25–0.30	Legitimate content retained

Penalties: over-sanitisation (removing benign content) reduces score.

Setup & Usage

Local

bash

pip install -r requirements.txt

set HF_TOKEN=hf_your_token_here
set API_BASE_URL=https://router.huggingface.co/v1
set MODEL_NAME=meta-llama/Llama-3.1-8B-Instruct

python inference.py

Docker

bash

docker build -t prompt-injection-env .
docker run -e HF_TOKEN=hf_your_token -e MODEL_NAME=meta-llama/Llama-3.1-8B-Instruct prompt-injection-env

Project Structure

prompt-injection-env/
├── inference.py                    # Mandatory baseline inference script
├── Dockerfile
├── requirements.txt
├── openenv.yaml
├── README.md
└── prompt_injection_env/
    ├── __init__.py
    ├── env.py                      # Main environment (reset/step/state)
    ├── corpus.py                   # Attack corpus and document builders
    └── graders.py                  # Deterministic scoring logic

Baseline Scores

Task	Score
task_easy	~0.85
task_medium	~0.60
task_hard	~0.35
Average	~0.60