Prompt Injection Guard¶

AegisAI includes a low-latency Prompt Guard that inspects prompts before they reach AI agents.

Architecture¶

Client → Guard Middleware → Guard Service
                              ├─ Rule detector (regex / banlist)
                              ├─ ML detector (TF-IDF + Isolation Forest)
                              └─ OPA Rego (policies/templates/guard.rego)
                         → Redis cache (TTL 300s)
                         → Audit Logger (SHA-256) on anomalies

Target latency: < 100ms (rules + cached ML; OPA timeout 2s with local fallback).

Scoring aggregates weighted detector outputs. Strong rule or OPA signals can hard-block regardless of the soft ML score. Aggregation uses max(weighted, rule, opa × 0.95) against GUARD_THRESHOLD_BLOCK / GUARD_THRESHOLD_WARN.

Threats Detected¶

Category	Examples
Prompt injection	`ignore previous instructions`, system prompt override
Jailbreak	`jailbreak`, `bypass security`, DAN / developer mode
Data leakage	`extract all data`, dump passwords / API keys
Token smuggling	Encoded / obfuscated instruction payloads

API¶

`GET /api/v1/guard/health`¶

Public health + current config (no auth).

`POST /api/v1/guard/validate`¶

Requires tenant X-API-Key.

{
  "agent_id": "uuid-or-name",
  "prompt": "User message…",
  "context": { "intent": "chat" }
}

Response:

{
  "threat_level": "blocked",
  "score": 0.86,
  "reasons": ["Matched pattern: system_override"],
  "blocked": true,
  "detected_patterns": ["Matched pattern: system_override"],
  "processing_time_ms": 12.4,
  "audit_id": "…",
  "cached": false
}

threat_level is one of: safe, suspicious, blocked.

`GET /api/v1/guard/stats`¶

In-process counters: blocked / suspicious / safe, average latency, top patterns. Also mirrored on dashboard via GET /api/v1/stats (guard_blocked, guard_avg_ms, guard_top_patterns).

Middleware¶

When ENABLE_GUARD=true, GuardMiddleware intercepts:

POST /api/v1/agents/{id}/chat
POST /api/v1/agents/{id}/action
POST /api/v1/agents/{id}/completions
POST /api/v1/chat

Blocked prompts return 403 with Guard details. Suspicious prompts are logged and allowed with X-Guard-* response headers. Detector failures are fail-open (request proceeds) so availability is preserved.

Configuration¶

Env var	Default	Description
`ENABLE_GUARD`	`true`	Master switch
`GUARD_ENABLE_ML`	`true`	Enable ML detector
`GUARD_ENABLE_RULES`	`true`	Enable regex / banlist
`GUARD_ENABLE_OPA`	`true`	Query OPA `guard.main`
`GUARD_ML_THRESHOLD`	`0.7`	ML soft threshold
`GUARD_RULE_THRESHOLD`	`0.8`	Hard block on strong rule hits
`GUARD_THRESHOLD_BLOCK`	`0.7`	Aggregate score → block
`GUARD_THRESHOLD_WARN`	`0.4`	Aggregate score → suspicious
`GUARD_RULE_WEIGHT`	`0.3`	Aggregation weight
`GUARD_ML_WEIGHT`	`0.5`	Aggregation weight
`GUARD_OPA_WEIGHT`	`0.2`	Aggregation weight
`GUARD_CACHE_TTL_SECONDS`	`300`	Redis cache TTL
`GUARD_MODEL_PATH`	`./models/guard/guard_model.pkl`	ML artifact
`GUARD_MODEL_BACKEND`	`hybrid`	`sklearn` / `onnx` / `hybrid`
`GUARD_MIN_TRAINING_SAMPLES`	`100`	Minimum samples for API training
`GUARD_EXTERNAL_ENABLED`	`false`	External API for uncertain scores

Training the ML model¶

CLI (offline)¶

poetry run python scripts/generate_prompt_dataset.py
poetry run python scripts/train_guard_model.py \
  --dataset data/prompt_dataset.csv \
  --model models/guard/guard_model.pkl
poetry run python scripts/benchmark_guard_model.py

Default model: TF-IDF + LogisticRegression (supervised). Isolation Forest is used as fallback without labels.

API (per-tenant, from audit + feedback)¶

POST /api/v1/guard/detector/train?days=30
GET /api/v1/guard/detector/status
GET /api/v1/guard/detector/versions
POST /api/v1/guard/detector/rollback/{version}
DELETE /api/v1/guard/detector/cancel

Dashboard UI: Prompt Guard ML page (/guard-ml).

ONNX (optional)¶

pip install onnxruntime transformers torch
poetry run python scripts/export_guard_onnx.py --dataset data/prompt_dataset.csv

Set GUARD_MODEL_BACKEND=onnx or hybrid.

Without a saved model, the detector bootstraps from a small seed corpus at startup.

SDK¶

from aegisai.sdk import AegisAIClient

async with AegisAIClient(api_key="…", agent_name="CustomerAgent") as client:
    await client.register()
    if await client.guard.is_safe(prompt):
        # call your LLM
        ...
    else:
        result = await client.guard.validate(prompt)
        # result["reasons"], result["score"]

OPA policy¶

Rego module: policies/templates/guard.rego (package guard.main).

Loaded automatically with other templates when OPA_ENABLED=true. If OPA is down, Guard uses a local regex fallback (fail-open for availability, still signals known attacks).