Build adversarial embedding defense

Introduction

When you ship a RAG system to production, attackers will probe the retrieval layer long before they bother with prompt injection — embedding spaces are mathematically exploitable, and a single crafted query can pull back documents the LLM was never supposed to see. Teams that defend only the generation step learn this the hard way when a stuffed or topic-targeted query leaks confidential chunks into a chat response and the incident shows up in a compliance review. By the end of this lesson you'll be able to layer three concrete defenses — structural query validation, intent classification, and retrieval diversity enforcement — so adversarial embedding attacks fail at the retrieval boundary instead of the LLM output.

Key Terminology

• Query stuffing — an attack that submits an extremely long or repetitive query so the resulting embedding becomes a noise vector that matches unpredictably; the first defense layer you build here rejects it on length and repetition.
• Nearest-neighbor attack — a sequence of syntactically benign queries each crafted to sit semantically close to one confidential chunk; matters here because no single query looks suspicious, so you defend by diversifying results, not by inspecting individual queries.
• Query intent classification — using a fast LLM to label a query as benign, exploratory, suspicious, or adversarial before it reaches the vector store; this is the semantic layer that catches attacks rule-based validation cannot see.
• Retrieval diversity enforcement — a post-retrieval filter that caps how many results can come from a single source, blocking systematic extraction across repeated queries.

Concepts

Structural query validation

The first defense rejects queries on shape, not meaning. A maximum token length (around 500 tokens for OpenAI embeddings) prevents stuffed queries from collapsing into noise vectors, and a repetition-ratio check catches the common attack of repeating a phrase to amplify its weight in the embedding. Encoded-character scans block payloads that try to smuggle special tokens past the embedder. This layer is cheap, deterministic, and runs before any model call — it should reject every obviously malformed query so the downstream layers only see plausible-looking input (see Code Walkthrough).

Semantic intent classification

Structural checks cannot catch a query like "summarize the Q3 board memo on the pending acquisition" — it is well-formed and short, but its intent is to retrieve specific confidential content. A small, fast model (Gemini Flash here) classifies each query into one of four intents — benign, exploratory, suspicious, adversarial — and the last two are blocked before retrieval. The classifier returns confidence and reasoning so security audits can review borderline calls (see Code Walkthrough).

Retrieval result diversity

The third layer assumes attacks will get through and limits their yield. By capping the share of results coming from any one source (e.g. no source contributes more than 40% of the top-k), a nearest-neighbor attacker who issues many narrow queries cannot reassemble a single restricted document, because each query returns at most a fraction of it. Diversity enforcement runs on the result list, not the query, so it composes cleanly with the first two layers.

Code Walkthrough

The two snippets below implement the three concepts above. The first combines structural validation and intent classification into a single pre-retrieval gate; the second enforces diversity on whatever the vector store returns.

Code snippetpython
1import json
2import os
3import re
4from dataclasses import dataclass, field
5
6from google import genai
7
8@dataclass
9class ValidationResult:
10    is_valid: bool
11    query: str
12    violations: list[str] = field(default_factory=list)
13
14@dataclass
15class IntentClassification:
16    intent: str
17    confidence: float
18    reasoning: str
19    is_blocked: bool
20
21class QueryValidator:
22    """Structural pre-retrieval check: length, repetition, encoded payloads."""
23
24    def __init__(self, max_tokens: int = 500, max_repetition_ratio: float = 0.5):
25        self.max_tokens = max_tokens
26        self.max_repetition_ratio = max_repetition_ratio
27
28    def validate(self, query: str) -> ValidationResult:
29        violations = []
30        tokens = query.split()
31        if len(tokens) > self.max_tokens:
32            violations.append(f"exceeds {self.max_tokens}-token cap ({len(tokens)})")
33        if tokens:
34            repetition = 1 - (len(set(tokens)) / len(tokens))
35            if repetition > self.max_repetition_ratio:
36                violations.append(f"repetition ratio {repetition:.2f}")
37        if re.search(r"%[0-9a-fA-F]{2}|\\x[0-9a-fA-F]{2}", query):
38            violations.append("encoded-character payload")
39        return ValidationResult(is_valid=not violations, query=query, violations=violations)
40
41class QueryIntentClassifier:
42    """Semantic pre-retrieval check via Gemini Flash."""
43
44    PROMPT = (
45        "Classify the query into exactly one of: benign, exploratory, "
46        "suspicious, adversarial.\n\nQuery: {query}\n\n"
47        'Respond as JSON: {{"intent": "...", "confidence": 0.0-1.0, "reasoning": "..."}}'
48    )
49
50    def __init__(self):
51        self.client = genai.Client(
52            api_key="student-token",
53            http_options={"api_endpoint": os.environ.get("GEMINI_PROXY_URL")},
54        )
55
56    def classify(self, query: str) -> IntentClassification:
57        response = self.client.models.generate_content(
58            model="gemini-2.0-flash",
59            contents=self.PROMPT.format(query=query),
60        )
61        parsed = json.loads(response.text)
62        intent = parsed.get("intent", "benign")
63        return IntentClassification(
64            intent=intent,
65            confidence=float(parsed.get("confidence", 0.5)),
66            reasoning=parsed.get("reasoning", ""),
67            is_blocked=intent in ("suspicious", "adversarial"),
68        )

• ValidationResult and IntentClassification are the structured outputs each layer returns so the calling pipeline can log audit detail and short-circuit on a block decision.
• QueryValidator.validate runs three deterministic checks (length, repetition, encoded payloads) and returns the union of violations — cheap, runs first, no model calls.
• QueryIntentClassifier.classify sends the query to Gemini Flash with a four-label prompt and blocks on suspicious/adversarial; it runs only on queries that passed structural validation.

Code snippetpython
1from collections import Counter
2from dataclasses import dataclass
3
4@dataclass
5class DiversityConfig:
6    max_source_percentage: float = 0.4
7    min_unique_sources: int = 2
8
9class DiversityEnforcer:
10    """Post-retrieval cap: no single source dominates the result list."""
11
12    def __init__(self, config: DiversityConfig = None):
13        self.config = config or DiversityConfig()
14
15    def enforce(self, results: list[dict], source_field: str = "source") -> list[dict]:
16        if len(results) <= 1:
17            return results
18        max_per_source = max(1, int(len(results) * self.config.max_source_percentage))
19        counts: Counter = Counter()
20        kept, deferred = [], []
21        for r in results:
22            source = r.get(source_field, "unknown")
23            if counts[source] < max_per_source:
24                kept.append(r)
25                counts[source] += 1
26            else:
27                deferred.append(r)
28        remaining = len(results) - len(kept)
29        kept.extend(deferred[:remaining])
30        return kept

• DiversityConfig parameterizes the per-source cap (default 40%) so different collections can tune it without code changes.
• enforce walks results in similarity order, admits each one only while its source is under the cap, defers the rest, and backfills from the deferred list to keep the total result count stable.

You'll know it works when a stuffed query is rejected at QueryValidator with a violation list, a topic-targeted query is blocked by QueryIntentClassifier with reasoning logged, and a benign query that happens to match many chunks from one source returns a result list where no source exceeds the configured percentage.

Do's and Don'ts

Building on the three-layer pipeline you just wired up, the following rules turn the defenses into operational habits — what to instrument, what to tune, and what to avoid when this lands in front of real traffic.

Do's

1. ✓Do run validation, classification, and diversity as three separate stages — each layer catches attacks the others can't, and isolating them keeps audit logs interpretable when an incident review asks which gate caught what.
2. ✓Do log every block decision with the violating query and reasoning — security review needs the evidence trail, and false positives can only be tuned if you can replay the rejected queries.
3. ✓Do tune the diversity cap against your corpus — 40% is a starting point; a collection with few sources needs a higher cap, and one with sensitive single-document topics needs a lower one.

Don'ts

1. ✗Don't rely on length limits alone — structural checks miss semantically-crafted nearest-neighbor attacks, which is exactly what intent classification and diversity enforcement are for.
2. ✗Don't classify intent with the same large model that generates answers — use a small fast model (e.g. Gemini Flash) so the classifier adds milliseconds, not seconds, to every query.
3. ✗Don't enforce diversity before validation and classification — running it last means it only ever sees queries that already passed the cheaper gates, and it composes correctly with whatever the vector store returned.