Generative AI Google Brain

gemini-responsible-ai-react-rag-attacks-mitigations-examples

3 min read

Gen AI and LLM Security - ReAct and RAG attacks & mitigations

Source notebook

Repo path: gemini/responsible-ai/react_rag_attacks_mitigations_examples.ipynb · Open on GitHub · intermediate

Demonstrates ReAct and RAG prompt-injection attacks with Gemini and simple mitigations.

Summary

This notebook teaches security risks in simplified ReAct agent and RAG implementations using Vertex AI Gemini. It builds toy weather and order tools, shows how compromised tool output can induce unsafe actions, then adds schema validation and out-of-band confirmation. It also demonstrates a PDF-based RAG flow, a manipulated PDF attack using hidden text, and OCR-based extraction as a mitigation.

Key code patterns

Initialize Gemini on Vertex AI

vertexai.init(project=project_id, location="us-central1")
model = GenerativeModel("gemini-2.0-flash")
config = GenerationConfig(
    temperature=0.0,
    max_output_tokens=2048,
    top_k=1,
    top_p=0.1,
    candidate_count=1,
)

Sets a low-temperature Gemini client for reproducible agent and RAG examples.

Simple ReAct loop

def chat(question: str) -> str:
    agent_scratchpad = ""
    for i in range(3):
        response = model.generate_content(
            prompt_template.format(input=question, agent_scratchpad=agent_scratchpad),
            generation_config=config,
        )
        response_last_lines = "\n".join(response.text.splitlines()[-3:])
        if "WAITING" in response_last_lines:
            action, action_input = extract_action(response_last_lines)
            observation = weather_city(action_input) if action == "weather_city" else order_store(action_input)
            agent_scratchpad += response.text + f"Observation: {observation}\n"
        else:
            agent_scratchpad += response.text
            break

Shows how tool observations are fed back into the model, creating an injection surface.

Validate tool output

def validate_weather(observation: str) -> str:
    pattern = r"(?i)(sunny|snowy|cloudy|rainy),\s+-?\d+\s+°C"
    matches = re.findall(pattern, observation)
    if matches:
        return observation
    print(">>> Error: Not proper weather tool output")
    return "Weather is unknown. Stop using the tool weather"

Filters compromised tool responses before they enter the agent scratchpad.

PDF text RAG

doc = pymupdf.open("Beyond41.pdf")
 
def search_snippets(query: str) -> str:
    text = ""
    for page in doc:
        text += page.get_text()
    return text
 
def chat_rag(question: str) -> str:
    documents = search_snippets(question)
    response = model.generate_content(
        prompt_template.format(input=question, documents=documents),
        generation_config=config,
    )
    return response.text

Demonstrates a minimal RAG pattern over PDF text and how document content can steer answers.

OCR mitigation for hidden text

pdf_file = "Beyond41mal.pdf"
 
def search_snippets(query: str) -> str:
    pages = convert_from_path(pdf_file)
    full_text = ""
    for page_num, page_image in enumerate(pages):
        text = pytesseract.image_to_string(page_image)
        full_text += f"{text}\n"
    return full_text

Uses OCR extraction when invisible PDF text may poison retrieved context.

Models & APIs used

  • Models: gemini-2.0-flash
  • APIs / services: Vertex AI, Cloud Storage
  • SDKs / libraries: google-cloud-aiplatform, google-cloud, pymupdf, pdf2image, pytesseract, vertexai

When to use this

Use this pattern to teach or test prompt-injection risks in agent tool loops and PDF-backed RAG pipelines.

Gotchas & caveats

  • The notebook states it is learning and demonstration material and not production code.
  • Requires Google Cloud project, region, authentication, and required API services enabled.
  • Uses Colab runtime restart after installing dependencies.
  • Installs system packages poppler-utils and tesseract-ocr for PDF conversion and OCR.
  • The ReAct parser uses simple regex and is not a production schema validator.
  • OCR can introduce recognition errors and requires more resources.
  • The notebook recommends ready agents and RAG libraries for real implementations.

Best practices

  • Use ready Agents and RAG libraries such as Agent Builder, LangChain Agents, Vertex AI Search, and LangChain RAG.
  • Use strict schema validation of tool input and output.
  • Use out-of-band user consent for dangerous operations.
  • Apply defense in depth by layering multiple filters.
  • Use OCR for documents if concerned about invisible text.
  • Use low temperature for reproducible results.

Graph View

Backlinks