iraajayakumar/pii-redaction-environment

class PiiRedactionAction(BaseModel):
    redacted_text: str  # Document with PII redacted (required)

class PiiRedactionObservation(BaseModel):
    task_id: str                    # "easy", "medium", or "hard"
    original_text: str              # Original document
    redacted_text: str              # Agent's redacted version
    score: float                    # Grader score [0.0, 1.0]
    pii_caught: int                 # Number of PII instances successfully redacted
    pii_total: int                  # Total PII instances in document
    judge_result: Optional[Dict]    # For hard task: re-identification assessment
    feedback: str                   # Human-readable feedback

easy_score = (pii_caught / pii_total) × reduction_penalty
  where reduction_penalty penalizes over-aggressive redaction (>50% loss)

medium_score = (pii_caught / pii_total) × context_quality
  where context_quality is LLM assessment of semantic preservation

hard_score = defense_effectiveness
  = 0.5 × baseline_score + 0.5 × (successful_defenses / max_attempts)

episode_score = (easy_score + medium_score + hard_score) / 3
  clamped to [0.0, 1.0]

# Clone repository
git clone https://github.com/iraajayakumar/pii-redaction-env
cd pii_redaction_env

# Create virtual environment
python -m venv .venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate

# Install dependencies
pip install -r pii_redaction_env/server/requirements.txt
pip install openai

# Verify installation
openenv validate

# Required for LLM-based grading/inference
export API_BASE_URL="http://localhost:11434"      # Ollama endpoint
export MODEL_NAME="mistral"                        # Model to use
export HF_TOKEN="your-huggingface-token"          # Optional, for HuggingFace API fallback

# Start Ollama (in separate terminal)
ollama serve

# In another terminal, pull Mistral model
ollama pull mistral

# Run baseline inference
cd pii_redaction_env
python inference.py

[START] task=pii_redaction env=pii_redaction_env model=mistral
[STEP] step=1 action=redact_structured_pii(caught=10/10) reward=0.92 done=false error=null
[STEP] step=2 action=redact_contextual_pii(method=agent_llm+grader_validation) reward=0.71 done=false error=null
[STEP] step=3 action=defend_against_reidentification(attempts=2/3) reward=0.68 done=true error=null
[END] success=true steps=3 score=0.77 rewards=0.92,0.71,0.68

# Build Docker image
cd pii_redaction_env
docker build -t pii-redaction:latest .

# Run container
docker run -e API_BASE_URL="http://localhost:11434" \
           -e MODEL_NAME="mistral" \
           -p 8000:8000 \
           pii-redaction:latest

# Ensure you're logged in
huggingface-cli login

# Create Space on Hugging Face (via web UI)
# Then clone and push:
git clone https://huggingface.co/spaces/YOUR_USERNAME/pii-redaction-env
cd pii-redaction-env
cp -r ../pii_redaction_env/* .
git add .
git commit -m "Initial deployment"
git push

python test_quick_judges.py

python test_llm_judges_integrated.py

python test_all_tasks_comprehensive.py

pii_redaction_env/
├── inference.py                      # Baseline agent + evaluation loop
├── openenv.yaml                      # OpenEnv spec
├── Dockerfile                        # Docker container definition
├── README.md                         # This file
├── server/
│   ├── app.py                       # FastAPI server
│   ├── requirements.txt              # Python dependencies
│   ├── tasks/
│   │   ├── easy.py                  # Easy task generator
│   │   ├── medium.py                # Medium task generator
│   │   └── hard.py                  # Hard task generator
│   ├── graders/
│   │   ├── easy_grader.py           # Easy task grader (regex)
│   │   ├── medium_grader.py         # Medium task grader (NER + LLM)
│   │   └── hard_grader.py           # Hard task grader (LLM judge)
│   └── data/
│       └── documents.py             # Professional sample documents
└── tests/
    ├── test_quick_judges.py         # Fast grader tests
    └── test_llm_judges_integrated.py # Full integration tests

from pii_redaction_env import PiiRedactionEnv

env = await PiiRedactionEnv.from_docker_image("pii-redaction:latest")
observation = await env.reset()
# Returns: PiiRedactionObservation with task_id, original_text

from pii_redaction_env import PiiRedactionAction

action = PiiRedactionAction(redacted_text="...redacted document...")
result = await env.step(action)
# Returns: (observation, reward, done, info)

state = await env.state()
# Returns: Current environment state snapshot

@misc{pii_redaction_env,
  title={PII Redaction Environment: Training Agents for Privacy Compliance},
  author={Jayakumar, Iraa},
  year={2026},
  howpublished={\url{https://github.com/iraajayakumar/pii-redaction-env}}
}

# Push to your personal namespace (defaults to username/env-name from openenv.yaml)
openenv push

# Push to a specific repository
openenv push --repo-id my-org/my-env

# Push with a custom base image
openenv push --base-image ghcr.io/meta-pytorch/openenv-base:latest

# Push as a private space
openenv push --private

# Combine options
openenv push --repo-id my-org/my-env --base-image custom-base:latest --private

from pii_redaction_env import PiiRedactionEnv

# Connect to existing server
pii_redaction_envenv = PiiRedactionEnv(base_url="<ENV_HTTP_URL_HERE>")

# Use as normal
result = pii_redaction_envenv.reset()
result = pii_redaction_envenv.step(PiiRedactionAction(message="Hello!"))

from pii_redaction_env import PiiRedactionAction, PiiRedactionEnv

# Connect with context manager (auto-connects and closes)
with PiiRedactionEnv(base_url="http://localhost:8000") as env:
    result = env.reset()
    print(f"Reset: {result.observation.echoed_message}")
    # Multiple steps with low latency
    for msg in ["Hello", "World", "!"]:
        result = env.step(PiiRedactionAction(message=msg))
        print(f"Echoed: {result.observation.echoed_message}")

# In server/app.py - use factory mode for concurrent sessions
app = create_app(
    PiiRedactionEnvironment,  # Pass class, not instance
    PiiRedactionAction,
    PiiRedactionObservation,
    max_concurrent_envs=4,  # Allow 4 concurrent sessions
)

from pii_redaction_env import PiiRedactionAction, PiiRedactionEnv
from concurrent.futures import ThreadPoolExecutor

def run_episode(client_id: int):
    with PiiRedactionEnv(base_url="http://localhost:8000") as env:
        result = env.reset()
        for i in range(10):
            result = env.step(PiiRedactionAction(message=f"Client {client_id}, step {i}"))
        return client_id, result.observation.message_length

# Run 4 episodes concurrently
with ThreadPoolExecutor(max_workers=4) as executor:
    results = list(executor.map(run_episode, range(4)))

# From the server directory
python3 server/pii_redaction_env_environment.py

uvicorn server.app:app --reload

pii_redaction_env/
├── .dockerignore         # Docker build exclusions
├── __init__.py            # Module exports
├── README.md              # This file
├── openenv.yaml           # OpenEnv manifest
├── pyproject.toml         # Project metadata and dependencies
├── uv.lock                # Locked dependencies (generated)
├── client.py              # PiiRedactionEnv client
├── models.py              # Action and Observation models
└── server/
    ├── __init__.py        # Server module exports
    ├── pii_redaction_env_environment.py  # Core environment logic
    ├── app.py             # FastAPI application (HTTP + WebSocket endpoints)
    └── Dockerfile         # Container image definition