Skip to content

Latest commit

 

History

History
93 lines (77 loc) · 4.3 KB

File metadata and controls

93 lines (77 loc) · 4.3 KB

Prompt-injection guard (rag-injection, Step 7.3)

rag-injection detects and removes adversarial instructions embedded in retrieved context before they reach the LLM, and builds the prompt so untrusted content can only ever appear as delimited data in the user turn.

Overview

A retrieved chunk is data, not commands — but an attacker who controls a source document can plant text that hijacks the model. This package is the LLM-adapter-level defense the red-team gate requires (≥ 95 % block on the corpus, no user_supplied chunk in a system-trust position). It depends only on rag-core + rag-observability.

Usage

from rag_injection import (
    PromptInjectionGuard,
    INJECTION_RESISTANT_SYSTEM_PROMPT,
    build_user_message,
)
from rag_core.spi.llm import LLMMessage

guard = PromptInjectionGuard()                      # heuristic detector default
result, safe_chunks = await guard.inspect(ctx, chunks)
# result.verdict -> clean | flagged | blocked ; safe_chunks drops the injected ones

messages = [
    LLMMessage(role="system", content=INJECTION_RESISTANT_SYSTEM_PROMPT),
    LLMMessage(role="user", content=build_user_message(query, safe_chunks)),
]

In the gateway it is wired automatically on every answer surface (/v1/query, /v1/chat/completions, MCP) and gated by cfg.injection.enabled (off by default). Config:

injection:
  enabled: true
  threshold: 0.5
  action: block          # or "annotate"
  block_untrusted: false
  max_scan_chars: 20000

API surface

Symbol Purpose
PromptInjectionGuard(*, detector=None, config=None) Inspect context + drop injection payloads.
async PromptInjectionGuard.inspect(ctx, chunks) -> (InjectionResult, list[Chunk]) The safe chunk list + the decision record.
InjectionConfig(threshold, action, per_tenant_thresholds, block_untrusted, max_scan_chars) Behaviour knobs.
InjectionResult(verdict, action, threshold, scanned_n, blocked_n, untrusted_n, categories) Telemetry-safe decision record.
InjectionAction annotate | block.
InjectionVerdict clean | flagged | blocked.
InjectionDetector (ABC) Pluggable detector: scan(text) -> list[InjectionMatch], score(text) -> float.
HeuristicInjectionDetector Dependency-free default — a curated regex library.
InjectionMatch(category, pattern_id, weight) One signal that fired (never the payload text).
INJECTION_RESISTANT_SYSTEM_PROMPT System prompt that treats retrieved text as untrusted data.
build_user_message(query, chunks) / format_untrusted_context(chunks) Fenced user-turn data block.

Internals

  • Detection grammar. Patterns anchor on an adversarial verb (ignore / disregard / reveal / bypass / …) and a model-referencing target (its instructions / role / safety) within a bounded, sentence-local gap, so real phrasings ("ignore your previous instructions", "disregard the above directions") match while a verb or target alone in benign prose ("follow the install instructions", "ignore case") does not. Scores aggregate by noisy-OR over the distinct patterns that fire (bounded in [0, 1]).
  • Trust isolation. The system prompt is static (carries no chunk content); retrieved chunks are rendered into a non-spoofable fenced block in the user turn, each tagged with its trust level. A user_supplied chunk can therefore never reach a system-trust position.
  • Degrade-open. A detector failure logs injection.scan_degraded and keeps all chunks — a broken detector must never block answering.
  • PII-free telemetry. injection.blocked (and the InjectionEvent schema) carry only category names + counts + the detector identity — never the matched payload text.

Extension points

  • Plug in a real classifier: implement InjectionDetector.scan (a fine-tuned model, Lakera, Rebuff) and pass it: PromptInjectionGuard(detector=MyDetector()).
  • Per-tenant strictness: InjectionConfig.per_tenant_thresholds lowers the threshold for a stricter tenant.
  • Observation mode: action="annotate" keeps every chunk and only logs the verdict — useful for measuring injection rates before enforcing.

See guides/red-team.md and ADR-0045.