Integrate with LangChain

Use Model Armor with LangChain to enhance the security of your AI applications by proactively screening LLM prompts and responses. To help prevent malicious input, verify content safety, and enforce security policies consistently across your AI applications, add Model Armor into your LangChain workflows.

Before you begin

Before implementing the code, complete the following Google Cloud setup steps:

• Select or create a project and enable billing.
• Enable the Model Armor API in your project.
• Set up authentication using gcloud or by setting the GOOGLE_APPLICATION_CREDENTIALS environment variable.
• Make sure that your service account has the Model Armor User (roles/modelarmor.user) role. To manage templates, you need the Model Armor Admin(roles/modelarmor.admin) role
• Create templates in the Google Cloud console.

Runnables

To use Model Armor with LangChain, you need two primary runnables that fit into the standard LangChain interface.

• ModelArmorSanitizePromptRunnable: This component screens user input (prompts) before they reach the LLM. It detects malicious input or policy violations.

• ModelArmorSanitizeResponseRunnable: This component screens the output generated by the LLM before it is returned to the user. Any modifications that the user makes to the response after the primary security check are not filtered and will reach you in their raw state.

Note: These runnables are limited to text screening. If a prompt includes a document, the system scans only the extracted text.

Install the required Python packages

Install the following Python packages:

• langchain-google-community version 3.0.4 or later
• langchain-google-genai

pip install langchain-google-community langchain-google-genai

Usage patterns

You can use Model Armor with LangChain either by basic chaining or by using middleware.

Basic chaining

The most common way to use Model Armor is by chaining the prompt sanitizer, the model, and the response sanitizer together using the LangChain Expression Language (LCEL) pipe operator.

from langchain_google_community.model_armor import (
    ModelArmorSanitizePromptRunnable,
    ModelArmorSanitizeResponseRunnable,
)
from langchain_google_genai import ChatGoogleGenerativeAI

# Initialize Model Armor runnables
sanitize_prompt = ModelArmorSanitizePromptRunnable(
    project="PROJECT_ID",
    location="LOCATION",
    template_id="TEMPLATE_ID",
    fail_open="SANITIZE_PROMPT_ERROR_HANDLING"
)

sanitize_response = ModelArmorSanitizeResponseRunnable(
    project="PROJECT_ID",
    location="LOCATION",
    template_id="TEMPLATE_ID",
    fail_open="SANITIZE_RESPONSE_ERROR_HANDLING"
)

# Initialize the model
llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", <var>...</var>)

# Build the chain
chain = sanitize_prompt | llm | sanitize_response

# Invoke
try:
    result = chain.invoke("<var>YOUR_USER_PROMPT_HERE</var>")
    print(result)
except ValueError as e:
    print(f"Content blocked: {e}")

Replace the following:

• PROJECT_ID: your Google Cloud project ID.
• LOCATION: the location of the template.
• TEMPLATE_ID: the ID of the Model Armor template created in Before you begin.
• SANITIZE_PROMPT_ERROR_HANDLING: a boolean that controls error handling for ModelArmorSanitizePromptRunnable; True or False.

• SANITIZE_RESPONSE_ERROR_HANDLING: a boolean that controls error handling for ModelArmorSanitizeResponseRunnable; True or False.

  When handling errors, Model Armor behaves as follows:

  • True: logs a warning but lets content pass even if risks are detected.
  • False: raises a ValueError and blocks execution when unsafe content is detected.

Middleware for agents

For LangChain agents that use tools, you can use ModelArmorMiddleware. This middleware applies the relevant sanitizer runnable at different points of interaction that the LangChain agent goes through when using tools.

from langchain.agents import create_agent
from langchain.tools import tool
from langchain_google_genai import ChatGoogleGenerativeAI

from langchain_google_community.model_armor import (
    ModelArmorMiddleware,
    ModelArmorSanitizePromptRunnable,
    ModelArmorSanitizeResponseRunnable,
)

# Create Model Armor sanitizers (same configuration as earlier sections)
prompt_sanitizer = ModelArmorSanitizePromptRunnable(
    project="PROJECT_ID",
    location="LOCATION",
    template_id="TEMPLATE_ID",
    fail_open="SANITIZE_PROMPT_ERROR_HANDLING",
)

response_sanitizer = ModelArmorSanitizeResponseRunnable(
    project="PROJECT_ID",
    location="LOCATION",
    template_id="TEMPLATE_ID",
    fail_open="SANITIZE_RESPONSE_ERROR_HANDLING",
)

# Wrap them in the official Model Armor middleware
model_armor_middleware = ModelArmorMiddleware(
    prompt_sanitizer=prompt_sanitizer,
    response_sanitizer=response_sanitizer,
)

# --- Create agent with Model Armor middleware ---
agent = create_agent(
    model="SOME_GEMINI_VERSION",
    tools=[],
    middleware=[model_armor_middleware],
)

Replace the following:

• PROJECT_ID: your Google Cloud project ID.
• LOCATION: the location of the template.
• TEMPLATE_ID: the ID of the Model Armor template created in Before you begin.
• SANITIZE_PROMPT_ERROR_HANDLING: a boolean that controls error handling for ModelArmorSanitizePromptRunnable; True or False.

• SANITIZE_RESPONSE_ERROR_HANDLING: a boolean that controls error handling for ModelArmorSanitizeResponseRunnable; True or False.

  When handling errors, Model Armor behaves as follows:

  • True: logs a warning but lets content pass even if risks are detected.
  • False: raises a ValueError and blocks execution when unsafe content is detected.

This setup makes sure that even when the agent is executing multi-step reasoning or tool usage, inputs and outputs are screened.

Document screening

You can use Model Armor runnables to screen uploaded documents before processing them. This is essential for workflows like summarization, where the source document might contain hidden malicious content. You can load the file content and pass it through the ModelArmorSanitizePromptRunnable before sending it to the LLM.

Custom event handling

For observability, the runnables dispatch an on_model_armor_finding event when unsafe content is detected. You can attach a custom callback handler to log specific details (such as findings and content) or trigger alerts without crashing the application flow.

from langchain_core.callbacks import BaseCallbackHandler

class SecurityHandler(BaseCallbackHandler):
    def on_custom_event(self, name, data, **kwargs):
        if name == "on_model_armor_finding":
            print(f"Alert: {data['findings']} detected in {data['template_id']}")

# Attach to config
chain.invoke(user_input, config={"callbacks": [SecurityHandler()]})