Samyak17Jain/prompt-injection-env-final

Prompt Injection Detection Environment

OpenEnv-compatible RL environment for detecting and neutralizing prompt injection attacks across support bots, code assistants, document summarizers, and analyst workflows.

Built for the Meta x Scaler Hackathon 2026 Round 1 submission.

Problem Domain

Prompt injection is a practical AI security problem: a malicious or hidden instruction tries to override a model's intended behavior. This environment gives an agent structured observations and expects a structured security decision in return.

The project covers three task families:

• —easy: direct role-override attacks
• —medium: indirect or embedded prompt injection inside otherwise legitimate content
• —hard: multi-turn persona hijack attempts that require conversation awareness

The dataset contains 30 tasks total: 10 easy, 10 medium, and 10 hard, with safe examples in every tier to measure false positives.

Task Families

The submission is organized into three concrete human-aligned task families, each with deterministic graders and progressively harder attack patterns:

• —task_1_easy: direct role-override detection in realistic support, coding, document, and analyst workflows
• —task_2_medium: indirect prompt injection hidden inside otherwise legitimate business content such as notes, attachments, or embedded instructions
• —task_3_hard: multi-turn persona hijack attempts that rely on conversation history, authority framing, or roleplay setup

Each family simulates a task humans actually perform when evaluating AI system inputs: triaging user requests, reviewing content for malicious instructions, and safely rewriting inputs before handing them to downstream assistants.

Observation Space

reset() returns an ObservationModel with the following fields:

Action Space

step() expects an ActionModel JSON payload:

Validation is strict:

• —extra keys are rejected
• —safe actions must use attack_type="none" or null
• —safe actions must not include injection_span
• —injection actions must provide a non-none attack type

Reward Design

All rewards are normalized to 0.0 to 1.0 and returned by the environment through StepResponse.reward.

• —easy: label 0.40 + attack type 0.30 + sanitization 0.30
• —medium: label 0.30 + span F1 0.40 + explanation 0.25 + sanitization 0.30, normalized to stay in [0,1]
• —hard: label 0.25 + trajectory awareness proxy 0.35 + explanation 0.25 + safe deflection 0.15

grader.py includes deterministic scoring for:

• —label correctness
• —attack type matching
• —token-level span F1
• —hard-task trajectory proxy scoring
• —explanation quality heuristics
• —sanitization safety

While episodes are single-step for submission simplicity, the reward function is still shaped rather than binary: it gives partial credit for correct labels, attack-type classification, exact-span matching, explanation quality, and safe sanitization instead of only rewarding perfect end states.

OpenEnv Compatibility

The environment implements the full submission-facing OpenEnv interface:

• —typed Pydantic contracts in models.py for observation, action, reward breakdown, and step responses
• —reset() for initial observations
• —step(action) returning observation, reward, done, and info
• —state() for current episode state
• —openenv.yaml metadata describing tasks, spaces, and reward range

The repo includes validate_submission.py for local structural checks. If an official openenv validate command is provided by the hackathon tooling, it should be run against this same environment before final submission.

API Endpoints

The FastAPI server in server.py exposes:

• —GET / -> {"status": "ok"}
• —GET /health -> health plus redacted env-var readiness
• —POST /reset -> start a new episode, optionally with task_id
• —POST /step -> validate an action and return StepResponse
• —GET /state -> current episode state
• —GET /tasks -> list all task ids with difficulty and app context

Malformed POST /step payloads return HTTP 422 with structured validation details.

Project Files

• —environment.py: task loading, dataset validation, episode lifecycle
• —grader.py: scoring logic and local grader tests
• —models.py: frozen Pydantic contracts
• —server.py: OpenEnv-compatible API
• —inference.py: baseline inference and optional offline LLM judge
• —tasks.json: curated task set
• —openenv.yaml: environment metadata for submission
• —validate_submission.py: one-command local submission checks

Local Setup

1. 1.Create and activate a Python 3.11 virtual environment.
2. 2.Install dependencies:

pip install -r requirements.txt

1. 1.Copy .env.example to .env and provide values for:

• —API_BASE_URL
• —MODEL_NAME
• —HF_TOKEN
• —optional LLM_API_KEY or OPENAI_API_KEY fallback for inference.py

The inference layer is provider-agnostic by design. It uses an OpenAI-compatible client interface, so the same code can run against providers such as OpenAI or Groq by changing environment variables rather than changing application logic. This lets the environment stay submission-compatible while supporting cost-performance tradeoffs during local testing and baseline evaluation.

For evaluator compatibility, inference.py accepts the mandatory HF_TOKEN variable directly as its model API key. It also accepts LLM_API_KEY or OPENAI_API_KEY as local-development fallbacks, but judges can run the project using only the mandatory variables listed above.

1. 1.Start the API server:

uvicorn server:app --host 0.0.0.0 --port 7860 --reload

Docker

Build and run locally:

docker build -t prompt-injection-env .
docker run -p 7860:7860 --env-file .env prompt-injection-env

Internal Robustness Evaluation

To improve the system without changing the official submission-facing tasks.json, this repo includes a private stress-test set in eval_tasks.json.

• —Run internal evaluation with judge:
• —python3 internal_eval.py --tasks-file eval_tasks.json
• —Run internal evaluation without judge:
• —python3 internal_eval.py --tasks-file eval_tasks.json --no-judge

internal_eval.py evaluates tricky safe cases, indirect injections, and multi-turn persona-hijack scenarios using the same action-generation path as inference.py, while keeping the official 30-task submission file untouched.

Final Readiness Checklist

• —python3 grader.py
• —python3 inference.py --server http://127.0.0.1:7860 --all-tasks
• —python3 internal_eval.py --tasks-file eval_tasks.json
• —docker build -t prompt-injection-env .
• —verify GET /, GET /health, GET /tasks, and POST /reset
• —confirm HF Space is on the same commit as GitHub main
• —rotate any shared secrets before submission

Validation

Run the deterministic grader tests:

python grader.py

Run the full local submission check:

python validate_submission.py

This validates:

• —dataset size and balance
• —safe-example coverage per difficulty
• —schema and endpoint behavior
• —reward bounds
• —malformed request handling
• —openenv.yaml structure
• —README submission coverage

Inference

inference.py runs a baseline agent against the environment and can optionally add an offline explanation judge for analysis.

The baseline client is intentionally built around an OpenAI-compatible API surface instead of locking the project to a single vendor. In practice, this means judges or contributors can point API_BASE_URL, MODEL_NAME, and HF_TOKEN to a compatible provider of their choice.

For local judging, contributors can also use OPENAI_API_KEY or another OpenAI-compatible provider key such as Groq via LLM_API_KEY. No code changes are required to switch providers.

Examples:

python inference.py --server http://127.0.0.1:7860 --task-id task_2_medium_003 --no-judge
python inference.py --server http://127.0.0.1:7860 --all-tasks --tasks-file tasks.json

Baseline Scores

Baseline scores were generated using the default OpenAI-compatible setup in .env.example with:

• —API_BASE_URL=https://api.groq.com/openai/v1
• —MODEL_NAME=llama-3.3-70b-versatile
• —python inference.py --server http://127.0.0.1:7860 --all-tasks --no-judge

Observed baseline on the current task set:

• —total tasks: 30
• —successful tasks: 30
• —average deterministic reward: about 0.8472
• —minimum reward observed: 0.03
• —maximum reward observed: 1.0

These values are reproducible for the same provider, model, and task set, and may vary if judges intentionally swap to a different OpenAI-compatible model.

Sample Run

Example local flow:

curl http://127.0.0.1:7860/
curl -X POST http://127.0.0.1:7860/reset -H "Content-Type: application/json" -d "{\"task_id\":\"task_1_easy_001\"}"
python validate_submission.py

Expected outcomes:

• —root endpoint returns {"status":"ok"}
• —reset returns a valid ObservationModel
• —validation script prints Submission validation passed.

Submission Notes

• —tasks.json is validated on environment startup, including duplicate ids and difficulty balance
• —rewards are clamped to 0.0-1.0
• —action payloads use strict schema validation with no extra fields
• —safe scenarios are present in every difficulty tier
• —the repo includes a Dockerfile, openenv.yaml, inference client, and local validator