search-query-safety-guard

Search Query Safety Guard

A DistilBERT-base-uncased guardrail that classifies an AI agent's search query as SAFE or DANGEROUS before it reaches a web search/scrape tool. It is designed to sit as a proxy between an agent's search tool and the real search service.

• 0 → SAFE — allow the search
• 1 → DANGEROUS — block and return an error to the agent

Intended use

Wrap a search/scrape tool so every query is classified first. Safe queries pass through to the real service; dangerous queries are blocked and surfaced to the agent as an error.

The model is trained to flag queries that seek to obtain, produce, or facilitate clearly harmful or illegal outcomes, across categories including:

• terrorism and violent extremism
• sexual content, including material involving minors
• controlled substances — procurement, manufacture, distribution
• weapons and explosives — unlawful manufacture or acquisition
• violence and physical harm
• self-harm
• hate, harassment, and targeting of individuals or groups
• kidnapping and human trafficking
• fraud, counterfeiting, and money laundering
• arson
• theft or sale of credentials and secrets — e.g. stolen authentication tokens, leaked API keys, session cookies, private keys

It also distinguishes these from legitimate, allow-listed intent: historical and policy education, public-health research, reporting a crime, and seeking help or resources.

Cyber policy (intent, not keywords)

This model also backs a cyber-defense agent, so the security boundary is drawn by intent rather than by topic:

• SAFE — researching how an attack technique works, or how to detect, prevent, or train against it (e.g. explaining ransomware behavior and incident-response steps, or how credential stuffing works and how rate limiting mitigates it).
• DANGEROUS — procuring, selling, harvesting, or exfiltrating credentials, secrets, or PII, or targeting a specific person's or organization's accounts.

The rule of thumb: a query about understanding or defending against a threat is allowed; a query about obtaining or selling the means to carry one out is blocked.

Usage

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

model = AutoModelForSequenceClassification.from_pretrained("urbanspr1nter/search-query-safety-guard")
tokenizer = AutoTokenizer.from_pretrained("urbanspr1nter/search-query-safety-guard")
model.eval()

def predict(query: str) -> dict:
    enc = tokenizer(query, return_tensors="pt", max_length=256, truncation=True, padding=True)
    with torch.no_grad():
        probs = torch.softmax(model(**enc).logits, dim=-1)
    pred = int(torch.argmax(probs, dim=-1).item())
    return {
        "query": query,
        "label": model.config.id2label[pred],   # "SAFE" or "DANGEROUS"
        "safe": pred == 0,
        "confidence": float(probs[0, pred].item()),
    }

predict("how to cook pasta")          # -> SAFE
predict("how does a buffer overflow work and how is it prevented")   # -> SAFE
# a query seeking to buy or harvest credentials -> DANGEROUS

Evaluation

Trained and evaluated with two held-out sets, both verified disjoint from the training data:

Recall on dangerous queries is 1.000 on both held-out sets — no dangerous query reached the web. The second set uses fresh wording of the same patterns and exists specifically to guard against the model memorizing the first.

Training

• Base: distilbert-base-uncased, 2 labels, max 256 tokens
• Optimizer: AdamW, lr 2e-5, batch 8, early stopping on F1
• Data: 2,313 examples (~61% SAFE / ~39% DANGEROUS), grown from 1,973 with a leakage guard that rejected any addition ≥82% similar to an eval query
• The training dataset and held-out evaluation sets are not published; only the model is distributed.

Limitations

• English only
• Evaluates a single query — no conversation history
• A small dataset cannot cover every edge case; deploy with a confidence threshold and log the grey zone