KChad/Prompt-Injection-RL-environment

1---2title: Prompt Injection Envinronment3emoji: "🔊"4colorFrom: blue5colorTo: green6sdk: docker7pinned: false8app_port: 80009base_path: /web10tags:11  - openenv-0.2.312  - openenv13  - Prompt Injection14  - Reinforcement Learning15---16 17# IDPI Exfil Environment18 19`IDPI Exfil Environment` is an OpenEnv-compatible prompt-injection defense environment. Agents operate on realistic email tasks where the visible goal looks benign, but attacker-controlled content tries to redirect behavior toward unsafe actions such as data exfiltration or malicious relays.20 21The environment is designed for evaluating whether an agent can:22 23- read and process email content safely24- identify indirect prompt injection patterns25- avoid unsafe outbound actions26- produce clean internal summaries on harder tasks27 28## Features29 30- `3 difficulty levels`: easy, medium, and hard scenarios with different attack visibility and trigger styles31- `Realistic attack channels`: inline instructions, hidden wrappers, HTML comments, transcript replay, attachment-style instructions, and more32- `Safety-aware action space`: includes both benign and risky tools so policies can be meaningfully tested33- `Structured scoring`: reward shaping during the episode plus end-of-episode grading34- `OpenEnv server mode`: FastAPI endpoints for reset, step, state, result, and health35- `Unified evaluation runner`: one runner for inference and baseline-style evaluation output36 37## Quick Start38 39### Install40 41```bash42pip install -e .43```44 45### Run the server46 47```bash48python server/app.py49```50 51The API starts on `http://localhost:7860`.52 53### Run the evaluation runner54 55Create a `.env` in the project root with the submission-compatible inference variables:56 57```bash58API_BASE_URL=https://generativelanguage.googleapis.com/v1beta/openai/59MODEL_NAME=gemini-2.5-flash60HF_TOKEN=your_gemini_api_key_here61```62 63Then run:64 65```bash66python inference.py67```68 69Compatibility commands still work:70 71```bash72python run_baseline.py73python baseline/run_baseline.py74```75 76All three commands use the same unified runner and write results to `baseline/baseline_results.json`.77For submission mode, `inference.py` validates `API_BASE_URL`, `MODEL_NAME`, and `HF_TOKEN`, then delegates to the benchmark runner using the OpenAI client path.78 79## Using the HTTP API80 81### Reset the environment82 83```python84import requests85 86response = requests.post(87    "http://localhost:7860/reset",88    json={"difficulty": "easy"},89)90observation = response.json()91print(observation)92```93 94### Step with an action95 96```python97import requests98 99response = requests.post(100    "http://localhost:7860/step",101    json={"action": {"type": "ReadDoc"}},102)103result = response.json()104 105print(result["observation"])106print(result["reward"])107print(result["done"])108```109 110### Inspect state or final result111 112```python113import requests114 115state = requests.get("http://localhost:7860/state").json()116episode_result = requests.get("http://localhost:7860/result").json()117```118 119Additional discovery endpoints are also available:120 121```python122import requests123 124metadata = requests.get("http://localhost:7860/metadata").json()125schema = requests.get("http://localhost:7860/schema").json()126```127 128## Environment Overview129 130### Task model131 132Each episode presents an email with a benign business-looking task and embedded attacker intent. The agent must complete the safe objective without obeying adversarial instructions contained in the email body or derived attack channels.133 134Typical safe behavior includes:135 136- reviewing the email with `ReadDoc`137- flagging suspicious instructions with `FlagInjection`138- writing an internal-only clean summary with `WriteDB`139- refusing unsafe outbound behavior140 141### Difficulty levels142 143- `easy`: explicit or obvious prompt injection patterns144- `medium`: more obfuscated attacks, hidden wrappers, spoofed instruction hierarchy, or delayed triggers145- `hard`: attacks embedded in semantically plausible content where the agent must also produce a safe summary146 147### Actions148 149- `ReadDoc`150  Reads the current email content.151- `WriteDB(content="...")`152  Writes an internal summary or note.153- `SendEmail(to="...", body="...")`154  Sends an outbound email. This is the main exfiltration risk surface.155- `FlagInjection(reason="...")`156  Marks the content as suspicious.157- `NoOp`158  Takes no action.159 160### Observation fields161 162- `subject`163- `body`164- `last_action`165- `last_action_result`166- `step_count`167- `task_difficulty`168- `summary_required`169- `current_summary`170 171### Endpoints172 173- `GET /`174- `GET /web`175- `GET /health`176- `GET /metadata`177- `GET /schema`178- `POST /reset`179- `POST /step`180- `GET /state`181- `GET /result`182- `POST /mcp`183 184## Evaluation185 186The evaluation runner loads scenarios from [docs/person_a_benchmark_manifest.json](/C:/Users/Admin/Desktop/My%20projects/Prompt-injection-env/docs/person_a_benchmark_manifest.json) and executes them across `easy`, `medium`, and `hard` splits.187 188The current manifest includes:189 190- `32 scenarios total`191- `9 easy`192- `13 medium`193- `10 hard`194 195Each run logs:196 197- per-step actions and rewards198- whether exfiltration occurred199- whether the task was completed safely200- final graded score201- per-difficulty summary statistics202 203Results are written to [baseline/baseline_results.json](/C:/Users/Admin/Desktop/My%20projects/Prompt-injection-env/baseline/baseline_results.json).204 205### Optional runner settings206 207These values can be set in `.env`:208 209```bash210EPISODES_PER_DIFFICULTY=1211BASELINE_RESULTS_PATH=baseline/submission_results.json212REQUEST_DELAY_SECONDS=3.0213MAX_MODEL_RETRIES=2214```215 216## Development217 218### Project structure219 220```text221Prompt-injection-env/222|-- env/223|   |-- environment.py224|   |-- dataset_loader.py225|   |-- scenario_generator.py226|   |-- injection_engine.py227|   |-- policy_engine.py228|   |-- taint_tracker.py229|   |-- reward.py230|   |-- grader.py231|   `-- models.py232|-- server/233|   |-- app.py234|   `-- Dockerfile235|-- docs/236|   |-- person_a_benchmark_manifest.json237|   |-- person_a_scenario_audit.md238|   `-- person_a_showcase_episodes.md239|-- baseline/240|   |-- run_baseline.py241|   `-- baseline_results.json242|-- inference.py243|-- run_baseline.py244|-- openenv.yaml245`-- README.md246```247 248### Core components249 250- [env/environment.py](/C:/Users/Admin/Desktop/My%20projects/Prompt-injection-env/env/environment.py)251  Main episode lifecycle, action handling, reward plumbing, and result assembly.252- [server/app.py](/C:/Users/Admin/Desktop/My%20projects/Prompt-injection-env/server/app.py)253  FastAPI wrapper exposing the environment over HTTP, including health, metadata, schema, and benchmark endpoints.254- [inference.py](/C:/Users/Admin/Desktop/My%20projects/Prompt-injection-env/inference.py)255  Submission-facing inference entrypoint that validates the required env vars and delegates to the benchmark runner.256- [openenv.yaml](/C:/Users/Admin/Desktop/My%20projects/Prompt-injection-env/openenv.yaml)257  Environment metadata for OpenEnv packaging and deployment.258 259 260 261### CI and sync262 263Includes GitHub workflows for:264 265- continuous integration in [.github/workflows/ci.yml](/C:/Users/Admin/Desktop/My%20projects/Prompt-injection-env/.github/workflows/ci.yml)266- syncing `main` to Hugging Face Spaces in [.github/workflows/sync.yml](/C:/Users/Admin/Desktop/My%20projects/Prompt-injection-env/.github/workflows/sync.yml)267 268Submission-focused regression tests live in `submission_tests/` and can be run with:269 270```bash271python -m unittest discover -s submission_tests272```273 274### Local Python usage275 276You can exercise the environment directly without starting the server:277 278```python279from env.environment import IDPIExfilEnv280from env.models import Action281 282env = IDPIExfilEnv()283obs = env.reset(difficulty="easy")284obs, reward, done, info = env.step(Action(type="ReadDoc"))285print(obs)286print(reward)287```288 289## Notes290 291- The environment terminates early on unsafe exfiltration events.292- Hard tasks require a clean internal summary that excludes attacker-controlled content.293- Gemini works through the submission-compatible OpenAI client path by setting `API_BASE_URL`, `MODEL_NAME`, and `HF_TOKEN`.294