Skip to content

Repository files navigation

injection-guard

Python 3.10+ License: Apache 2.0 Async Ensemble OpenAI Anthropic Gemini HuggingFace Model Armor

Two-stage prompt injection detection system with an ensemble classifier architecture.

Stage 1 runs fast local models (~100-200ms) — Model Armor and DeBERTa for PI/JB signals, Safeguard for configurable safety policy signals — and enriches the SignalVector with their results.

Stage 2 runs frontier API classifiers (OpenAI, Anthropic, Gemini) on every prompt, receiving enriched signals from Stage 1 as additional context for high-accuracy classification.

Built with async-first Python 3.10+. A 6-stage preprocessor pipeline (Unicode normalization, encoding detection, structural analysis, token boundary detection, GLiNER NER entity extraction, and regex pattern matching) extracts a SignalVector from each prompt before classification. Stage 1 never short-circuits — it provides signal enrichment only. The cascade router runs Stage 2 classifiers tier-by-tier with quorum-based aggregation.

Architecture

block-beta
  columns 6

  block:input:1
    columns 1
    A["Prompt"]
  end

  block:preprocess:1
    columns 1
    style preprocess fill:#2d3748,color:#fff,stroke:#4a5568
    B["Preprocessor"]
    B1["Unicode"]
    B2["Encoding"]
    B3["Structural"]
    B4["Token"]
    B5["GLiNER NER"]
    B6["Regex"]
    BP["Signal Extraction"]
  end

  block:stage1:1
    columns 1
    style stage1 fill:#065f46,color:#fff,stroke:#047857
    S1["Stage 1"]
    S1A["Model Armor (PI/JB)"]
    S1B["DeBERTa (PI/JB)"]
    S1C["Safeguard (Policy)"]
    S1D["~100-200ms"]
    S1P["SignalVector Enrichment"]
  end

  block:stage2:1
    columns 1
    style stage2 fill:#2c5282,color:#fff,stroke:#2b6cb0
    S2["Stage 2"]
    S2A["OpenAI"]
    S2B["Anthropic"]
    S2C["Gemini"]
    S2P["Frontier Ensemble"]
  end

  block:decide:1
    columns 1
    style decide fill:#553c9a,color:#fff,stroke:#6b46c1
    D["Aggregator"]
    D1["Weighted Average"]
    D2["Threshold Engine"]
    DP["Score Aggregation"]
  end

  block:output:1
    columns 1
    style output fill:#2f855a,color:#fff,stroke:#38a169
    E1["ALLOW"]
    E2["FLAG"]
    E3["BLOCK"]
    EP["Decision"]
  end

  A --> B
  B --> S1
  S1 --> S2
  S2 --> D
  D --> E1
Loading

How It Works

flowchart TD
  subgraph Preprocessor["Preprocessor (6 stages)"]
    direction LR
    P1["Unicode"] --> P2["Encoding"] --> P3["Structural"] --> P4["Token"] --> P5["GLiNER NER"] --> P6["Regex"]
  end

  Preprocessor --> SV["SignalVector + risk_prior"]
  SV --> Stage1

  subgraph Stage1["Stage 1: Signal Enrichment (~100-200ms)"]
    direction LR
    MA["Model Armor\n(PI/JB signal)"] ~~~ HF["DeBERTa\n(PI/JB signal)"] ~~~ SG["Safeguard\n(Safety policy)"]
  end

  Stage1 -->|"enrich SignalVector\n(no early exit)"| Stage2

  subgraph Stage2["Stage 2: Frontier Ensemble (every prompt)"]
    direction LR
    S2A["OpenAI\n(+ reasoning)"] ~~~ S2B["Anthropic"] ~~~ S2C["Gemini"]
  end

  Stage2 --> DEC

  subgraph DEC["Decision"]
    direction LR
    AG["Weighted\nAggregation"] --> TH["Threshold\nCheck"] --> OUT["ALLOW | FLAG | BLOCK"]
  end

  style Preprocessor fill:#2d3748,color:#e2e8f0
  style Stage1 fill:#065f46,color:#d1fae5
  style Stage2 fill:#2c5282,color:#bee3f8
  style DEC fill:#2f855a,color:#c6f6d5
Loading

Recommended Ensemble Strategy

The architecture uses a tiered approach optimized from eval results on the Qualifire benchmark:

  1. Model Armor (PI/JB signal) — Google Cloud Model Armor evaluates prompts (~180ms). Its verdict and confidence are added to the SignalVector. Optional — requires GCP.
  2. DeBERTa (PI/JB signal) — Fine-tunable DeBERTa model (~100ms on GPU) provides fast local PI/JB classification. Score, label, and confidence are added to the SignalVector. Customers can fine-tune on their domain data.
  3. Safeguard (safety policy signal) — gpt-oss-safeguard evaluates content against configurable safety policies (violence, hate speech, fraud, etc.). Contributes category-level reasoning to the SignalVector. Not a PI/JB classifier.
  4. SignalVector enrichment — All Stage 1 results are added to the SignalVector and passed as context to Stage 2 classifiers. Frontier models see what Stage 1 already detected.
  5. Frontier ensemble (every prompt) — The cascade/parallel router fires frontier API classifiers (Anthropic, OpenAI with reasoning, Gemini) and waits for quorum. These receive enriched signals and provide 80-84% accuracy with nuanced scoring. No early exit — every prompt goes through Stage 2 to avoid false positive bypasses.
  6. Weighted aggregation — The aggregator combines all scores using learned weights, then applies threshold engine for ALLOW/FLAG/BLOCK.

Stage 1 never short-circuits — it enriches context for Stage 2 but does not make blocking decisions on its own. This avoids false positives from fast local models bypassing the frontier ensemble.

Preprocessor Pipeline

Six stages extract signals from the raw prompt before classification:

Stage Detects Key Signals
1. Unicode Homoglyphs, zero-width chars, BiDi overrides homoglyph_count, zero_width_count, script_mixing
2. Encoding Base64, hex, URL-encoding, nested encoding encodings_found, encoding_density, nested_encoding
3. Structural Chat delimiters, XML/HTML tags, instruction boundaries chat_delimiters_found, separator_density
4. Token Split-keyword attacks, prompt stuffing reconstructed_keywords, repetition_ratio
5. GLiNER NER Injection-specific semantic entities entity_count, max_entity_confidence
6. Regex Known injection patterns (12 built-in) match_count, matched_patterns

Signals feed into a risk_prior (0.0-1.0) that can block early or escalate routing. They're also formatted into natural language and appended to LLM classifier prompts as evidence.

See docs/ner-signals.md for details on how GLiNER NER works and how signals augment classifiers.

Classifiers

Stage 1: Signal Enrichment

Stage 1 provides fast local signals that enrich the SignalVector for Stage 2. Stage 1 never blocks or short-circuits — every prompt proceeds to Stage 2 to avoid false positive bypasses.

Component Type Role
Model Armor GCP API PI/JB signal. Verdict + confidence added to SignalVector. Optional — requires GCP.
HF DeBERTa Local GPU PI/JB signal (~100ms). Score + label added to SignalVector. Fine-tunable on customer data.
Safeguard Local GPU Safety policy signal. Evaluates configurable policies (P1-P6: violence, hate speech, self-harm, sexual content, dangerous activities, fraud). Category codes + reasoning added to SignalVector.

Model Armor and DeBERTa provide PI/JB-specific signals. Safeguard is a safety policy signal provider (not a PI/JB classifier) — it evaluates content against configurable safety policies and contributes category-level reasoning. Custom policies (spam, compliance, domain-specific) can be passed via the policy parameter. See docs/safeguard-policy.md for built-in and custom policy examples. See Google Cloud Model Armor docs for Model Armor template configuration.

Stage 2: Frontier Classifiers

Classifier Type Weight Category Accuracy Approach
Anthropic API 2.0 api 83.5% Claude with few-shot classification prompt
OpenAI API 1.5 api 82.0% GPT-5 with reasoning tokens (high effort)
Gemini API 1.5 api 80.5% Gemini via google-genai with few-shot prompt
Local LLM Local 1.5 local Any Ollama/vLLM model with classification prompt
ONNX Local 1.0 local ONNX Runtime inference

All classifiers implement the BaseClassifier protocol and receive the SignalVector from the preprocessor. Stage 2 API classifiers use a shared few-shot classification prompt enriched with Stage 1 signals. HF DeBERTa models (Stage 1) are served via litguard and can be fine-tuned on customer data. Safeguard (Stage 1) uses a configurable safety policy as system prompt — see docs/safeguard-policy.md for built-in and custom policy examples.

Routing

Two strategies control how classifiers are invoked:

Parallel Router — fires all classifiers concurrently, returns when a category quorum is met:

router:
  type: parallel
  timeout_ms: 10000
  category_quorum:
    local: 1   # at least 1 local model must respond
    api: 2     # at least 2 API models must respond

Cascade Router (recommended) — runs classifiers tier-by-tier (fast → medium → slow), exits early on high confidence. This is the recommended strategy for the tiered pre-filter architecture:

router:
  type: cascade
  timeout_ms: 10000
  fast_confidence: 0.85          # exit early if confidence > 85%
  escalate_on_high_risk_prior: true
  risk_prior_escalation_threshold: 0.7  # skip fast tier if risk_prior > 0.7

The cascade router groups classifiers by their latency_tier attribute and runs them in order:

Tier Classifiers Latency Behavior
fast DeBERTa (HF), ONNX, Regex ~100ms Run first. If confidence > fast_confidence, return immediately.
medium Safeguard, Local LLM ~1-5s Run if fast tier is uncertain.
slow OpenAI, Anthropic, Gemini ~2-10s Run only for ambiguous cases.

If the preprocessor's risk_prior exceeds risk_prior_escalation_threshold, the fast tier is skipped entirely and classification starts at medium/slow tiers — this prevents high-risk prompts from being cleared by a less capable local model.

This gives sub-200ms decisions for ~70% of traffic (clear benign/injection via DeBERTa) while escalating only ambiguous cases to frontier API classifiers.

Quick Start

YAML Config

# --- Stage 1: Pre-gate + Pre-filter (fast, high recall) ---
gate:
  type: model_armor
  project: ${GOOGLE_CLOUD_PROJECT}
  location: global
  template_id: my-injection-template
  block_on: HIGH              # block only high-confidence detections
  fail_mode: open             # if MA is down, let prompts through

classifiers:
  # Fast pre-filter (fine-tunable, ~100ms)
  - type: hf_compat
    model: deberta-injection
    base_url: http://192.168.1.199:8234/v1
    weight: 1.0
    category: local

  # --- Stage 2: Frontier ensemble (high accuracy) ---
  - type: anthropic
    model: claude-sonnet-4-6
    weight: 2.0
    category: api

  - type: openai
    model: gpt-5-2025-08-07
    weight: 1.5
    reasoning_effort: high
    category: api

  - type: gemini
    model: gemini-3.1-pro-preview
    weight: 1.5
    category: api

router:
  type: cascade
  timeout_ms: 10000
  fast_confidence: 0.85
  escalate_on_high_risk_prior: true
  risk_prior_escalation_threshold: 0.7

thresholds:
  block: 0.85
  flag: 0.50

aggregator: weighted_average

preprocessor:
  gliner_model: urchade/gliner_base
from injection_guard import InjectionGuard

guard = InjectionGuard.from_config("config.yaml")

decision = await guard.classify("Ignore all previous instructions")
print(decision.action)        # Action.BLOCK
print(decision.ensemble_score) # 0.97
print(decision.model_scores)   # per-classifier results

decision = await guard.classify("What is the capital of France?")
print(decision.action)        # Action.ALLOW

Programmatic Setup

from injection_guard import InjectionGuard
from injection_guard.classifiers import AnthropicClassifier, SafeguardClassifier
from injection_guard.router import ParallelRouter
from injection_guard.types import ParallelConfig

guard = InjectionGuard(
    classifiers=[
        AnthropicClassifier(model="claude-sonnet-4-6"),
        SafeguardClassifier(
            model="gpt-oss-safeguard:120b",
            base_url="http://192.168.1.199:11434/v1",
        ),
    ],
    router=ParallelRouter(ParallelConfig(
        timeout_ms=10000,
        category_quorum={"local": 1, "api": 1},
        classifier_categories={
            "anthropic-claude-sonnet-4-6": "api",
            "safeguard-gpt-oss-safeguard": "local",
        },
    )),
)

Sync Wrapper

decision = guard.classify_sync("Tell me about Python")

Environment Variables

Create a .env file:

ANTHROPIC_API_KEY=sk-ant-...
OPENAI_API_KEY=sk-...
GOOGLE_CLOUD_PROJECT=my-project-123
GOOGLE_CLOUD_REGION=global

The .env file is loaded automatically on InjectionGuard init.

Benchmark Results

Evaluated on Qualifire prompt-injections-benchmark (200 balanced samples: 100 injection, 100 benign). Full results in docs/eval-results.md.

Stage 2: Frontier Classifiers

Model Accuracy Precision Recall F1
Anthropic claude-opus-4.6 83.5% 0.860 0.800 0.829
Anthropic claude-sonnet-4.6 82.5% 0.788 0.890 0.836
OpenAI gpt-5 (high reasoning) 82.0% 0.758 0.940 0.839
Gemini 3.1-pro-preview 80.5% 0.740 0.940 0.828
Gemini 3-flash-preview 80.0% 0.724 0.970 0.829
OpenAI gpt-5-mini (medium reasoning) 79.0% 0.769 0.830 0.798

Stage 1: Pre-gate + Pre-filter (local/fast)

Model Accuracy Precision Recall F1 Latency
Model Armor (MA Low) 74.5% 0.889 0.560 0.687 ~800ms
Model Armor (MA High) 58.5% 0.947 0.180 0.303 ~180ms
protectai/deberta (open-weight) 69.5% 0.714 0.650 0.681 ~100ms
deepset/deberta (open-weight) 65.0% 0.589 0.990 0.739 ~100ms

Run benchmarks yourself:

# Quick model benchmarks (10-sample)
pytest tests/integration/test_model_benchmarks.py -v -s

# Full Qualifire eval (200-sample, requires API keys + HF token)
pytest tests/integration/test_eval_classifiers.py -v -s -k "test_openai_gpt_5_high"

Testing

Setup

pip install -e ".[dev]"

# For benchmark datasets (Qualifire, ToxicChat)
pip install -e ".[benchmark]"

Unit Tests

No API keys or external services needed. All classifiers, routers, and external calls are mocked.

# Run all unit tests (353 tests)
pytest tests/unit/ -v

# Run by module
pytest tests/unit/test_preprocessor/ -v       # Preprocessor pipeline (6 stages)
pytest tests/unit/test_classifiers/ -v         # All classifier mocks
pytest tests/unit/test_router/ -v              # Cascade + parallel router
pytest tests/unit/test_aggregator/ -v          # Weighted, voting, meta
pytest tests/unit/test_guard.py -v             # Guard orchestrator
pytest tests/unit/test_pipeline.py -v          # Full pipeline (mocked classifiers)
pytest tests/unit/test_engine.py -v            # Threshold engine

Integration Tests

Integration tests hit real APIs and services. Tests auto-detect available prerequisites and skip cleanly when services aren't reachable. A Rich-formatted prerequisite table prints at the start of each run.

Note: Always pass -s to see Rich-formatted output (prerequisite checks, dataset summaries, test result tables). Pytest captures stdout by default.

pytest tests/integration/test_full_ensemble.py -v -s

Prerequisites:

Requirement Env Variable / Service Needed For
OpenAI API key OPENAI_API_KEY Frontier ensemble tests
Anthropic API key ANTHROPIC_API_KEY Frontier ensemble tests
GCP project GCP_PROJECT_ID Model Armor tests
HuggingFace token HF_TOKEN Qualifire dataset (gated)
Safeguard (Ollama) SAFEGUARD_BASE_URL (default 192.168.1.199:11434) Safety policy tests (ToxicChat)
litguard (DeBERTa) 192.168.1.199:8234 DeBERTa signal tests
datasets package pip install datasets All dataset-based tests
# Full ensemble test (Qualifire + ToxicChat, requires API keys)
pytest tests/integration/test_full_ensemble.py -v -s

# Qualifire benchmark (RegexPrefilter only, requires HF_TOKEN)
pytest tests/integration/test_benchmark_qualifire.py -v -s

# Individual classifier evals (requires respective API keys)
pytest tests/integration/test_eval_classifiers.py -v -s -k "test_anthropic_sonnet"
pytest tests/integration/test_eval_classifiers.py -v -s -k "test_openai_gpt_5_high"

# Model Armor eval (requires GCP_PROJECT_ID)
pytest tests/integration/test_eval_model_armor.py -v -s

# Quick model benchmarks (10-sample)
pytest tests/integration/test_model_benchmarks.py -v -s

# Live classifier smoke tests
pytest tests/integration/test_live_classifiers.py -v -s

Eval Datasets

The shared dataset loader at injection_guard.eval.dataset normalizes HuggingFace datasets into a common TestSample format:

from injection_guard.eval.dataset import load_qualifire, load_toxicchat, load_mixed

# Qualifire — PI/JB detection (gated, requires HF_TOKEN)
samples = load_qualifire(n=100, seed=42, balanced=True)

# ToxicChat — safety/toxicity detection (public)
samples = load_toxicchat(n=100, seed=42, balanced=True)

# Mixed — combined from both sources
samples = load_mixed(n_per_source=50, seed=42)

Each TestSample has prompt, label ("injection" / "benign" / "toxic" / "safe"), source, and is_attack property.

Test Structure

tests/
  unit/                              # All mocked, no external deps
    test_preprocessor/               # 6-stage signal extraction
    test_classifiers/                # OpenAI, Anthropic, Gemini, Safeguard, etc.
    test_router/                     # Cascade + parallel routing
    test_aggregator/                 # Score aggregation strategies
    test_guard.py                    # Guard orchestrator
    test_pipeline.py                 # Full pipeline (mocked)
    test_engine.py                   # Threshold engine
    test_config.py                   # YAML config loading
    test_eval/                       # Eval runner, report, batch adapters
  integration/                       # Real APIs, auto-skips on missing prereqs
    test_full_ensemble.py            # Full Stage 1 + Stage 2 with real datasets
    test_benchmark_qualifire.py      # Qualifire benchmark (RegexPrefilter)
    test_eval_classifiers.py         # Per-classifier eval (Batch API)
    test_eval_model_armor.py         # Model Armor eval
    test_model_benchmarks.py         # Quick multi-model benchmarks
    test_live_classifiers.py         # Smoke tests against live services

Documentation

  • Deployment Guide — production deployment patterns, FastAPI/LangChain integration, scaling, and monitoring
  • Eval Results — full benchmark results across all classifiers and models
  • Fine-Tuning Strategy — how to fine-tune DeBERTa models to improve detection metrics
  • Domain Fine-Tuning — domain-specific tuning for healthcare, finance, legal, and other verticals
  • NER Signals & Preprocessor — how GLiNER NER works and how signals augment classifiers
  • Safeguard Safety Policies — gpt-oss-safeguard as Stage 1 safety policy signal provider, custom policy authoring, and deployment
  • litguard Spec — LitServe-based model serving platform for HuggingFace models

Project Structure

src/injection_guard/
    types.py              # All shared types (single source of truth)
    guard.py              # Main orchestrator
    config.py             # YAML config loader & factory
    engine.py             # Threshold decision engine
    cli.py                # CLI entry point
    reporting.py          # Rich-powered reporting output
    preprocessor/
        pipeline.py       # 6-stage pipeline orchestration
        unicode.py        # Stage 1: Unicode normalization
        encoding.py       # Stage 2: Encoding detection
        structural.py     # Stage 3: Structural analysis
        token.py          # Stage 4: Token boundary detection
        gliner.py         # Stage 5: GLiNER entity detection
        regex.py          # Stage 6: Regex pattern matching
    classifiers/
        prompts.py        # Shared few-shot prompt & signal formatting
        openai.py         # OpenAI API classifier
        anthropic.py      # Anthropic API classifier
        gemini.py         # Google Gemini via Vertex AI
        safeguard.py      # gpt-oss-safeguard safety policy signal provider
        local_llm.py      # Ollama, vLLM, OpenAI-compatible
        onnx.py           # Local ONNX model
        regex.py          # Legacy regex prefilter
    router/
        cascade.py        # Tier-by-tier with early exit
        parallel.py       # Concurrent with category quorum
    aggregator/
        weighted.py       # Weighted average
        voting.py         # Majority voting
        meta.py           # Meta-classifier stacking
    gate/
        model_armor.py    # Google Cloud Model Armor (optional)
    eval/
        runner.py         # Dataset loading & evaluation
        report.py         # Metrics & threshold recommendation

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages