Gen AI & LLM Security for developers

Source notebook

Repo path: gemini/responsible-ai/gemini_prompt_attacks_mitigation_examples.ipynb · Open on GitHub · advanced

Shows prompt injection attacks against Gemini and layered mitigations with DLP, NL API, safety filters, and embeddings.

Summary

This notebook teaches Gen AI and LLM security patterns for developers by demonstrating data leaks, jailbreaking, hallucinations, payload splitting, virtualization, obfuscation, multimodal attacks, and model poisoning. It builds a vulnerable Gemini travel assistant, then adds mitigations using Cloud DLP, Cloud Natural Language classification and sentiment checks, VirusTotal URL checks, LLM validation, DARE prompting, strict token validation, embeddings with ScaNN, and Gemini safety filters.

Key code patterns

Gemini model with system instructions

MODEL_ID = "gemini-2.0-flash"
config = GenerationConfig(temperature=0.0, max_output_tokens=2048, top_k=1, top_p=0.1)
model = GenerativeModel(
    model_name=MODEL_ID,
    system_instruction=["You are friendly travel expert...", "The tech support password is Bacon&Egg1234"],
)

Creates the baseline assistant and intentionally exposes why secrets and PII must not be placed in prompts.

DLP input and output guard

def generate_response_with_dlp(input: str) -> str:
    if not valid_dlp_text(input):
        return "Please do not provide PII as input."
    output = model.generate_content(input, generation_config=config).text
    if not valid_dlp_text(output):
        return "Sorry, I cannot provide the answer."
    return output

Checks both user input and Gemini output before returning potentially sensitive content.

Cloud Natural Language category blocklist

BLOCKED_CATEGORIES = ["/Sensitive Subjects", "/Business & Industrial/Advertising & Marketing"]
CONFIDENCE_THRESHOLD = 0.1
response = client.classify_text(document=language_v2.Document(content=text_content, type_=language_v2.Document.Type.PLAIN_TEXT, language_code="en"))
for category in response.categories:
    if blocked_category in category.name and category.confidence > CONFIDENCE_THRESHOLD:
        return False

Uses topic classification to reject inputs outside the intended application category.

Strict input validation token

token = str(random.getrandbits(128))
prompt_template = f"""
You should respond with TRUE if the question violates any of these guidelines and with "{token}" if it is safe to answer.
Question: {input}
"""
response = model_guardrail.generate_content(prompt_template, generation_config=config)
return token in response.text

Requires the guardrail model to return an unpredictable token for safe prompts, reducing prompt override risk.

Embedding similarity screen

embedding_model = TextEmbeddingModel.from_pretrained("google/text-embedding-005")
df["embedding"] = df["textContent"].apply(lambda x: get_embedding(x))
searcher = scann.scann_ops_pybind.builder(normalized_dataset, 10, "dot_product").tree(num_leaves=record_count, num_leaves_to_search=record_count, training_sample_size=record_count).score_ah(2).reorder(100).build()

Builds a small dangerous-prompt vector index to detect semantically similar attacks.

Models & APIs used

  • Models: gemini-2.0-flash, google/text-embedding-005
  • APIs / services: Vertex AI, Cloud Data Loss Prevention, Cloud Natural Language API, VirusTotal API
  • SDKs / libraries: google-cloud-aiplatform, google-cloud-dlp, google-cloud-language, vertexai, scann, virustotal-python, pandas, numpy, colorama

When to use this

Use this pattern when building a Gemini application that needs practical prompt-attack demonstrations and layered preflight, postflight, and semantic guardrails.

Gotchas & caveats

  • The notebook states it is learning and demonstration material and is not production code.
  • Vertex AI API must be enabled and vertexai.init requires a Google Cloud project and location.
  • DLP and Cloud Natural Language API require gcloud auth application-default login and a quota project.
  • Cloud DLP can miss transformed or encoded sensitive data such as inserted characters between digits.
  • VirusTotal checks require a user-provided API key from getpass.
  • The poisoned model endpoint is a placeholder and must be replaced with a tuned model endpoint.
  • The guardrail example turns safety filters to BLOCK_NONE for demonstration and says not to do this in production.
  • The embedding dataset is small and not comprehensive for evaluating models.

Best practices

  • Do not store sensitive information in the prompt.
  • Use low temperature for reproducible results in security demonstrations.
  • Check both input and output with DLP before sending to or returning from Gemini.
  • Treat links, binaries, and files from users as untrusted and validate them.
  • Use Responsible AI safety filters and understand how to configure safety attributes.
  • Use embeddings to find similar and dangerous prompts.
  • Use application mission checks with DARE prompting before answering.