Metadata-Version: 2.4
Name: quaveo_guardrail
Version: 0.1.0
Summary: Composable prompt-injection guardrails for Python
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: google-auth>=2.0
Provides-Extra: armor
Requires-Dist: google-cloud-modelarmor>=0.7.1; extra == "armor"
Requires-Dist: presidio-analyzer>=2.2.364; extra == "armor"
Requires-Dist: presidio-anonymizer>=2.2.364; extra == "armor"
Provides-Extra: gemini
Requires-Dist: google-genai>=1.0; extra == "gemini"
Provides-Extra: deberta
Requires-Dist: torch>=2.0; extra == "deberta"
Requires-Dist: transformers>=4.0; extra == "deberta"
Provides-Extra: all
Requires-Dist: google-cloud-modelarmor>=0.7.1; extra == "all"
Requires-Dist: presidio-analyzer>=2.2.364; extra == "all"
Requires-Dist: presidio-anonymizer>=2.2.364; extra == "all"
Requires-Dist: google-genai>=1.0; extra == "all"
Requires-Dist: torch>=2.0; extra == "all"
Requires-Dist: transformers>=4.0; extra == "all"

# guardrails

A lightweight Python library that packages two LLM guardrail detectors — **Model Armor** (cloud API) and **DeBERTa** (ML classifier) — into a clean, importable interface.

Both implement the same `BaseGuardrail` interface, so they are interchangeable and can be layered.

---

## Detectors

### ModelArmorGuardrail
Connects to the real [Google Cloud Model Armor](https://cloud.google.com/security/products/model-armor) API. Scans user prompts and model responses against a configured template. Supports SDP (Sensitive Data Protection) de-identification — when enabled on the template, the result includes a redacted version of the text instead of just a block decision.

- **Strengths**: production-grade, template-configurable, handles both input and output, SDP redaction
- **Limitations**: requires GCP project, Model Armor template, and `google-cloud-modelarmor` SDK; cloud latency
- **Context scanning**: combined into prompt (context prepended to prompt in one API call)

### DebertaGuardrail
Uses [`protectai/deberta-v3-base-prompt-injection-v2`](https://huggingface.co/protectai/deberta-v3-base-prompt-injection-v2), a fine-tuned transformer classifier. Understands semantic variants of injection attacks, including those embedded in RAG-retrieved context.

- **Strengths**: catches novel phrasing, multilingual attacks, indirect/RAG injections
- **Limitations**: requires `torch` + `transformers`; ~500 MB model download on first use
- **Context scanning**: configurable via `scan_mode`

---

## Installation

Install the package from PyPI (or from a local checkout while developing):

```bash
pip install quaveo_guardrail
```

Install optional integrations only when you need them:

```bash
pip install "quaveo_guardrail[armor]"    # Google Cloud Model Armor + Presidio
pip install "quaveo_guardrail[gemini]"   # Gemini semantic check
pip install "quaveo_guardrail[deberta]"  # DeBERTa classifier
pip install "quaveo_guardrail[all]"      # all integrations
```

The import name remains `guardrails` for compatibility with existing code:

```python
from guardrails import DebertaGuardrail, DetectionResult, ModelArmorGuardrail, WafGuardrail
```

To install the current checkout in editable mode:

```bash
.venv/bin/pip install -e ".[all]"
```

---

## Creating a Model Armor template

Before using `ModelArmorGuardrail` you must create a template in your GCP project. The template name is then set as `MODEL_ARMOR_TEMPLATE` in your `.env` file or passed as the `template=` argument.

### Required IAM permissions

**Model Armor template management**

To create and manage Model Armor templates you need one of the following IAM roles on the GCP project:

| Role | ID | Use case |
|---|---|---|
| Model Armor Admin | `roles/modelarmor.admin` | Full access — create, update, delete templates |
| Model Armor Editor | `roles/modelarmor.editor` | Create and update templates (no delete) |

The specific permission gating template creation is `modelarmor.templates.create`. Project `Owner` and `Editor` primitive roles also carry this permission.

```bash
gcloud projects add-iam-policy-binding PROJECT_ID \
  --member="user:EMAIL" \
  --role="roles/modelarmor.admin"
```

**Sensitive Data Protection (advanced SDP mode only)**

Advanced SDP mode requires pre-existing Sensitive Data Protection inspect/de-identify templates. To create those you need:

| Role | ID | Use case |
|---|---|---|
| DLP Editor | `roles/dlp.editor` | Create and manage inspect and de-identify templates |
| DLP Inspect Templates Editor | `roles/dlp.inspectTemplatesEditor` | Create/edit inspect templates only |
| DLP De-identify Templates Editor | `roles/dlp.deidentifyTemplatesEditor` | Create/edit de-identify templates only |

If the SDP templates live in a **different project** from the Model Armor template, grant the Model Armor service agent access to that project:

```bash
# replace PROJECT_NUMBER with the number of the Model Armor project
gcloud projects add-iam-policy-binding SDP_PROJECT_ID \
  --member="serviceAccount:service-PROJECT_NUMBER@gcp-sa-modelarmor.iam.gserviceaccount.com" \
  --role="roles/dlp.user"

gcloud projects add-iam-policy-binding SDP_PROJECT_ID \
  --member="serviceAccount:service-PROJECT_NUMBER@gcp-sa-modelarmor.iam.gserviceaccount.com" \
  --role="roles/dlp.reader"
```

No extra grants are needed when the SDP templates reside in the **same project** as the Model Armor template.

### Enable the API

```bash
gcloud services enable modelarmor.googleapis.com --project=PROJECT_ID
```

### Create a template

**Option 1 — Google Cloud console**

1. Go to [Model Armor in the console](https://console.cloud.google.com/security/modelarmor).
2. Click **Create Template**.
3. Set **Template ID**, **Region** (cannot be changed later), and configure detections.
4. Click **Create**.

**Option 2 — REST API (curl)**

Model Armor uses regional endpoints. Set `LOCATION` to a supported region (e.g. `us`, `eu`, `us-central1`).

```bash
export PROJECT_ID=your-gcp-project
export LOCATION=us
export TEMPLATE_ID=my-guardrail-template

export TEMPLATE_CONFIG='{
  "filter_config": {
    "rai_settings": {
      "rai_filters": [
        {"filter_type": "HATE_SPEECH",      "confidence_level": "MEDIUM_AND_ABOVE"},
        {"filter_type": "HARASSMENT",       "confidence_level": "MEDIUM_AND_ABOVE"},
        {"filter_type": "DANGEROUS",        "confidence_level": "MEDIUM_AND_ABOVE"},
        {"filter_type": "SEXUALLY_EXPLICIT","confidence_level": "MEDIUM_AND_ABOVE"}
      ]
    },
    "pi_and_jailbreak_filter_settings": {
      "filter_enforcement": "ENABLED",
      "confidence_level": "HIGH"
    },
    "malicious_uri_filter_settings": {
      "filter_enforcement": "ENABLED"
    }
  },
  "template_metadata": {
    "enforcement_type": "INSPECT_AND_BLOCK",
    "multi_language_detection": {"enable_multi_language_detection": true}
  }
}'

curl -X POST \
  -d "$TEMPLATE_CONFIG" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  "https://modelarmor.${LOCATION}.rep.googleapis.com/v1/projects/${PROJECT_ID}/locations/${LOCATION}/templates?template_id=${TEMPLATE_ID}"
```

The template name returned (`projects/PROJECT_ID/locations/LOCATION/templates/TEMPLATE_ID`) is what you pass as `MODEL_ARMOR_TEMPLATE` in your `.env` file, or as the `template=` argument to `ModelArmorGuardrail`.

> **Tip**: Start with `"enforcement_type": "INSPECT_ONLY"` to observe block rates before switching to `INSPECT_AND_BLOCK`.

For full documentation see [Create and manage templates](https://docs.cloud.google.com/model-armor/manage-templates) and [Model Armor roles and permissions](https://docs.cloud.google.com/model-armor/access-control/roles-permissions).

---

## Quick start

```python
from guardrails import ModelArmorGuardrail, DebertaGuardrail, DetectionResult

# Real Google Cloud Model Armor (config from .env)
armor = ModelArmorGuardrail(
    template="my-template",      # name or full resource path
    project_id="my-gcp-project", # omit to auto-detect from ADC / env
    location="us",
    model="gemini-2.0-flash",    # optional: adds Gemini semantic step
    vertexai=True,
    vertex_location="us-central1",
)

# Scan user input
result: DetectionResult = armor.scan(prompt=user_input, context=rag_chunk)
print(result.flagged)         # True = blocked
print(result.reason)          # "ARMOR_BLOCKED" | "ALLOW"
print(result.sanitized_text)  # de-identified text if SDP is configured, else None

# Scan model output before returning to user
out_result = armor.scan_response(response=llm_output, user_prompt=user_input)
if out_result.flagged:
    raise PermissionError("model response blocked")
safe_text = out_result.sanitized_text or llm_output

# DeBERTa ML classifier (model downloads on first scan())
deberta = DebertaGuardrail()  
if not deberta.is_safe(user_input, context=rag_chunk):
    raise PermissionError("injection detected")
```

---

## Detection pipelines

### ModelArmorGuardrail — `scan()` (input path)

```
            User prompt + optional RAG context
                           │
                           ▼
┌─────────────────────────────────────────────────────────┐
│ Step 1 · Presidio injection pre-check (local, no I/O)   │
│                                                         │
│  Regex recognizer scans the prompt for patterns such    │
│  as "ignore previous instructions", "reveal … system    │
│  prompt", "developer mode", "override safety", or       │
│  "output … api_key".                                    │
│                                                         │
│  Language: English only. Non-English text skips this    │
│  step entirely.                                         │
│                                                         │
│  ✗ match → flagged=True, reason="PRESIDIO_INJECTION"    │
└───────────────────────────┬─────────────────────────────┘
                            │ no match
                            ▼
┌─────────────────────────────────────────────────────────┐
│ Step 2 · Google Cloud Model Armor API                   │
│                                                         │
│  context + "\n\n" + prompt are merged into one string   │
│  so RAG-injected context is evaluated together with     │
│  the user turn.                                         │
│                                                         │
│  The combined text is sent to                           │
│  sanitize_user_prompt() against the configured          │
│  template. The template can enable any combination of:  │
│    • Prompt injection & jailbreak filter                │
│    • Responsible AI (hate, harassment, dangerous, etc.) │
│    • Malicious URI detection                            │
│    • Sensitive Data Protection (SDP) de-identification  │
│                                                         │
│  ✗ MATCH_FOUND → flagged=True, reason="ARMOR_BLOCKED"   │
│  ✓ no match + SDP enabled → sanitized_text populated    │
└───────────────────────────┬─────────────────────────────┘
                            │ no match
                            ▼ (only when model= is set)
┌─────────────────────────────────────────────────────────┐
│ Step 3 · Gemini semantic classifier (optional)          │
│                                                         │
│  The same merged text is sent to a Gemini model with    │
│  a zero-shot system prompt that instructs it to reply   │
│  with exactly INJECTION or SAFE.                        │
│                                                         │
│  Catches multilingual attacks and indirect injections   │
│  that defeat regex and template-based filters.          │
│                                                         │
│  ✗ INJECTION → flagged=True, reason="GEMINI_INJECTION"  │
└───────────────────────────┬─────────────────────────────┘
                            │ SAFE
                            ▼
               flagged=False, reason="ALLOW"
```

### ModelArmorGuardrail — `scan_response()` (output path)

```
                    LLM response text
                            │
                            ▼
┌─────────────────────────────────────────────────────────┐
│ Step 1 · Google Cloud Model Armor API                   │
│                                                         │
│  Calls sanitize_model_response() with the response and  │
│  the original user prompt (for context).                │
│                                                         │
│  Template rules apply as configured (RAI filters,       │
│  malicious URIs, SDP de-identification, etc.).          │
│                                                         │
│  ✗ MATCH_FOUND → flagged=True, reason="ARMOR_BLOCKED"   │
│  ✓ no match + SDP → sanitized_text populated            │
└───────────────────────────┬─────────────────────────────┘
                            │ no match (use_presidio=True)
                            ▼
┌─────────────────────────────────────────────────────────┐
│ Step 2 · Presidio PII masking (local, no I/O)           │
│                                                         │
│  Runs the standard Presidio recognizer set (names,      │
│  emails, phone numbers, credit cards, …) plus a custom  │
│  12-digit Vietnamese national ID recognizer.            │
│                                                         │
│  Any detected entities are replaced with                │
│  [REDACTED_PII] or [REDACTED_VN_ID].                    │
│  The redacted text is stored in sanitized_text.         │
│  The result is never flagged here — PII masking is      │
│  informational, not a block decision.                   │
└───────────────────────────┬─────────────────────────────┘
                            │
                            ▼
               flagged=False, reason="ALLOW"
         sanitized_text = masked text (if any PII found)
```

---

### DebertaGuardrail — `scan()` (both modes)

The DeBERTa classifier is a fine-tuned transformer that outputs a probability score for the `INJECTION` class. A score above `threshold` (default `0.5`) is treated as a detection.

**`scan_mode="dual"` (default)**

```
prompt ──────────────────► _classify(prompt)  → (flagged_p, score_p)
                                                         │
context (if provided) ──► _classify(context) → (flagged_c, score_c)
                                                         │
                          flagged = flagged_p OR flagged_c
                          confidence = max(score_p, score_c)
```

Each field is tokenized and scored independently. This is the recommended mode for RAG pipelines because an injection hidden in a retrieved document is detected even if the user prompt itself is clean.

**`scan_mode="combined"`**

```
context + "\n\n" + prompt
           │
           ▼
      _classify(combined text)  →  (flagged, confidence)
```

The two fields are concatenated before the model sees them. Useful when the injection spans both the context and the prompt, and the classifier needs the full surrounding text to detect it.

**Inference detail**

```
text
  │  AutoTokenizer  (truncation=True, max_length=512 by default)
  ▼
token IDs (on device: cpu / cuda / mps — auto-detected)
  │  DeBERTa-v3-base (sequence classification head)
  ▼
logits  →  softmax  →  [P(SAFE), P(INJECTION)]
                                  │
                        > threshold? → flagged
```

The model is lazy-loaded on the first `scan()` call and cached for subsequent calls.

---

## API reference

### `DetectionResult`
```python
@dataclass
class DetectionResult:
    flagged: bool              # True = threat detected
    reason: str                # "PRESIDIO_INJECTION" | "ARMOR_BLOCKED" | "GEMINI_INJECTION" | "INJECTION" | "ALLOW"
    confidence: float          # 0.0–1.0  (armor is always 1.0 or 0.0)
    scanned_fields: list[str]  # e.g. ["prompt"], ["response"], ["prompt", "context"]
    sanitized_text: str | None # de-identified text from Model Armor SDP; None otherwise
```

### `BaseGuardrail`
All detectors implement:
```python
def scan(self, prompt: str, context: str = "") -> DetectionResult: ...
def is_safe(self, prompt: str, context: str = "") -> bool: ...
```

### `ModelArmorGuardrail(template, project_id, location, use_presidio, model, api_key, vertexai, vertex_location)`

Mirrors the full `agent.py` pipeline: Presidio injection pre-check → Model Armor API → optional Gemini semantic check → Presidio PII masking.
Requires `google-cloud-modelarmor`. When `use_presidio=True` also requires `presidio-analyzer`, `presidio-anonymizer`.
When `model` is set, also requires `google-genai`.

| Parameter | Type | Default | Description |
|---|---|---|---|
| `template` | `str \| None` | `None` | Template name or full resource path. Reads `MODEL_ARMOR_TEMPLATE` env if `None` |
| `project_id` | `str \| None` | `None` | GCP project. Auto-resolved from `GOOGLE_CLOUD_PROJECT` env or ADC |
| `location` | `str \| None` | `None` | Model Armor API region. Reads `MODEL_ARMOR_LOCATION` env, defaults to `"us"` |
| `use_presidio` | `bool` | `True` | Run Presidio injection pre-check on input and PII masking on output |
| `model` | `str \| BaseLlm \| None` | `None` | Gemini model name or `BaseLlm` instance. When set, adds a Gemini semantic check as step 3 of `scan()`. Disabled by default |
| `api_key` | `str \| None` | `None` | Gemini API key for Gemini API mode. Reads `GOOGLE_API_KEY` / `GEMINI_API_KEY` env if `None` and `vertexai=False` |
| `vertexai` | `bool \| None` | `None` | `True` → Vertex AI mode; `False` → Gemini API mode; `None` → auto-detect from `GOOGLE_GENAI_USE_VERTEXAI` env |
| `vertex_location` | `str \| None` | `None` | Vertex AI region for the Gemini client (e.g. `"us-central1"`, `"global"`). Reads `GOOGLE_CLOUD_LOCATION` env, defaults to `"global"`. Ignored when `vertexai=False` |

**`scan()` pipeline:**
1. Presidio custom recognizer checks prompt for injection patterns (`ignore ... instructions`, `reveal ... system`, `developer mode`, `override safety`, `output ... api_key`, etc.) — **English only** (`language="en"`); non-English injections bypass this step
2. If not caught, Model Armor `sanitize_user_prompt` on `context + prompt`
3. If `model` is set and steps 1–2 passed, a Gemini call classifies the text as `INJECTION` or `SAFE` — catches multilingual and indirect injections missed by the earlier steps

**`scan_response()` pipeline:**
1. Model Armor `sanitize_model_response`
2. Presidio PII masking on result — standard PII entities + Vietnamese ID (`\b\d{12}\b`); redacted text returned in `sanitized_text`

Additional method:
```python
def scan_response(self, response: str, user_prompt: str = "") -> DetectionResult: ...
```

```python
# all config read from .env
armor = ModelArmorGuardrail()

# explicit config, Presidio disabled
armor = ModelArmorGuardrail(template="my-template", location="us", use_presidio=False)

# with Gemini semantic check via Vertex AI
armor = ModelArmorGuardrail(
    template="my-template",
    model="gemini-2.0-flash",
    vertexai=True,
    vertex_location="us-central1",
)

# with Gemini semantic check via Gemini API key
armor = ModelArmorGuardrail(
    template="my-template",
    model="gemini-2.0-flash",
    vertexai=False,
    api_key="YOUR_GEMINI_API_KEY",
)
```

### `DebertaGuardrail(model_name, threshold, device, scan_mode, truncation, max_length)`

| Parameter | Type | Default | Description |
|---|---|---|---|
| `model_name` | `str` | `"protectai/deberta-v3-base-prompt-injection-v2"` | HuggingFace model ID |
| `threshold` | `float` | `0.5` | Injection confidence cutoff — raise to reduce false positives |
| `device` | `str \| None` | `None` | `"cpu"`, `"cuda"`, or `"mps"`. Auto-detected if `None` |
| `scan_mode` | `"dual" \| "combined"` | `"dual"` | See below |
| `truncation` | `bool` | `True` | Whether to truncate inputs that exceed `max_length` |
| `max_length` | `int` | `512` | Maximum token length passed to the tokenizer |

**`scan_mode` options:**
- `"dual"` — scans `prompt` and `context` separately; flags if **either** exceeds threshold. Best for RAG pipelines where injections may be hidden in retrieved documents.
- `"combined"` — concatenates `context + prompt` into a single string before classifying. May improve detection when injection spans both fields.

```python
deberta = DebertaGuardrail(threshold=0.6, scan_mode="dual")
deberta = DebertaGuardrail(threshold=0.6, scan_mode="combined", max_length=256)
```

---

## Layered defence

DeBERTa for semantic depth, Model Armor for cloud-grade scanning and output sanitization:

```python
deberta = DebertaGuardrail(threshold=0.6, scan_mode="dual")
armor   = ModelArmorGuardrail(template="my-template")

def check_input(prompt: str, context: str = "") -> str:
    if not deberta.is_safe(prompt, context):
        raise PermissionError("request blocked by guardrail")
    result = armor.scan(prompt, context)
    if result.flagged:
        raise PermissionError("request blocked by Model Armor")
    # return de-identified prompt if SDP redacted anything
    return result.sanitized_text or prompt

def check_output(response: str, user_prompt: str = "") -> str:
    result = armor.scan_response(response, user_prompt)
    if result.flagged:
        raise PermissionError("response blocked by Model Armor")
    return result.sanitized_text or response
```

---

## Environment configuration

Create a `.env` file in the repo root for Model Armor credentials. All values are read automatically by `ModelArmorGuardrail` on first use.

```dotenv
GOOGLE_CLOUD_PROJECT=your-gcp-project-id
MODEL_ARMOR_TEMPLATE=your-template-name
MODEL_ARMOR_LOCATION=us

# Gemini semantic check (required when model= is set)
GOOGLE_GENAI_USE_VERTEXAI=true        # use Vertex AI; omit or set false for Gemini API mode
GOOGLE_CLOUD_LOCATION=global          # Vertex AI region for the Gemini client
GOOGLE_API_KEY=your-gemini-api-key    # only needed for Gemini API mode
```

`ModelArmorGuardrail` resolves config in this priority order:
1. Constructor args (`template`, `project_id`, `location`, `api_key`, `vertex_location`, …)
2. Environment variables (`GOOGLE_CLOUD_PROJECT`, `MODEL_ARMOR_TEMPLATE`, `MODEL_ARMOR_LOCATION`, `GOOGLE_CLOUD_LOCATION`, `GOOGLE_API_KEY` / `GEMINI_API_KEY`, `GOOGLE_GENAI_USE_VERTEXAI`)
3. Application Default Credentials (ADC) for project discovery

---

## Evaluation scripts

Pre-built scripts in `scripts/` evaluate each guardrail against the Rule 1 (Prompt Injection) and Rule 2 (Sensitive Info Disclosure) test datasets. Each script loads `.env` from the repo root, runs the full test set, and writes results to `results/`.

| Script | Guardrail | Rule | Output |
|---|---|---|---|
| `scripts/run_rule1_model_armor_lib.py` | `ModelArmorGuardrail` | Rule 1 | `results-rule1-model-armor-lib.json` |
| `scripts/run_rule2_model_armor_lib.py` | `ModelArmorGuardrail` | Rule 2 | `results-rule2-model-armor-lib.json` |
| `scripts/run_rule1_deberta_lib.py` | `DebertaGuardrail` | Rule 1 | `results-rule1-deberta-lib.json` |
| `scripts/run_rule2_deberta_lib.py` | `DebertaGuardrail` | Rule 2 | `results-rule2-deberta-lib.json` |

```bash
# run all four
uv run scripts/run_rule1_model_armor_lib.py
uv run scripts/run_rule2_model_armor_lib.py
uv run scripts/run_rule1_deberta_lib.py
uv run scripts/run_rule2_deberta_lib.py
```

**Pass/fail criteria** (consistent across all scripts):
- *Attack cases* — guardrail must `flagged=True` → pass
- *Benign Control* — guardrail must `flagged=False` → pass
- *Boundary Test* (Rule 2 only) — excluded from accuracy score; binary classifiers cannot partially disclose (e.g. "last 4 digits only")

---

## Extending

Both classes implement `BaseGuardrail`. You can subclass to add custom logic:

```python
from guardrails._base import BaseGuardrail
from guardrails.result import DetectionResult

class MyGuardrail(BaseGuardrail):
    def scan(self, prompt: str, context: str = "") -> DetectionResult:
        ...
```
