How to wear Model Armor 2: Integrating with ADK and LangChain

approx. 2900 words
14 min
minherz

Quick recap of part 1

The first post about Model Armor, explored the fundamentals of Google Cloud’s managed security service for Generative AI applications, which provides a model-agnostic defense layer to sanitize both prompts and model responses.

It covered the two primary patterns for integrating Model Armor into your stack:

  • Direct Invocation: Using the Model Armor SDK or API for granular control over pre-call and post-call sanitization.
  • Built-in Integration: Configuring services like Vertex AI, GKE, and Gemini Enterprise to automatically enforce security policies through “Floor Settings” and user-defined templates without explicit API calls in your application logic.

And it walked through the practical configuration of these integrations using gcloud CLI and Terraform, establishing a secure baseline for your GenAI pipelines. In this post I shift my focus to examine how direct invocation works in practice. I will review the methods of interpreting sanitize API responses and incorporating the API calls in two agent frameworks: LangChain ‒ probably the most widespread framework today for implementing agentic workflows and the Agent Development Kit (ADK) which I personally prefer for its simplicity.

Interpreting Sanitize API responses

Before looking into best practices for incorporating the Sanitize API calls into your application using agent frameworks, let’s review what you do once you made the call and received a response. I will use the Model Armor client library in Python for code snippets. However, it should be easy to translate it to other languages. The sanitize API calls return SanitizeUserPromptResponse or SanitizeModelResponseResponse depending which API you call. These are simply the wrappers around the instance of SanitizationResult. Let’s look into several decisions you need to make once you receive and unwrap the response.

How to check returned status

If Model Armor does not detect any threat, no action is needed. You have to inspect the result to confirm that execution was successful and the filters of the Model Armor template did not detect anything.

from google.cloud.modelarmor_v1 import (
    FilterMatchState,
    InvocationResult,
    SanitizationResult,
)

# result: modelarmor_v1.SanitizationResult
if (result.invocation_result == InvocationResult.SUCCESS and
    result.filter_match_state == FilterMatchState.MATCH_FOUND):
    #... process filter results

If the filter_match_state value is FilterMatchState.NO_MATCH_FOUND, your job is done. No further processing is needed.

Error handling

There is a case when result.invocation_result is set to InvocationResult.FAILURE or InvocationResult.PARTIAL. This signifies that there were problems during execution of the sanitized action. If the invocation partially failed, you should use sanitization_metadata in the result to decide what action to take.

from google.cloud.modelarmor_v1 import (
    InvocationResult,
    SanitizationResult,
)

# result: modelarmor_v1.SanitizationResult
if result.invocation_result == InvocationResult.FAILURE:
    reason = result.sanitization_metadata
    return f"FAILURE: ({reason.error_code}) {reason.error_message}"
if result.invocation_result == InvocationResult.PARTIAL:
    reason = result.sanitization_metadata
    if not reason.ignore_partial_invocation_failures:
        return f"FAILURE: ({reason.error_code}) {reason.error_message}"
#... continue processing

Process findings

You confirmed that invocation succeeded and you have FilterMatchState.MATCH_FOUND. It is time to find what caused the result and decide what to do next. Model Armor templates define four detection filters:

  • Prompt injection and jailbreak detection
  • Malicious URL detection
  • Sensitive Data Protection
  • Ethical AI for hate speech, harassment, dangerous and sexually explicit content

Once there is a match found, you use the filter_results map to identify which filters triggered the match. The map contains a set of predefined keys with information about the filter results.

Key Filter object
csam CsamFilterResult
malicious_uris MaliciousUriFilterResult
rai RaiFilterResult
pi_and_jailbreak PiAndJailbreakFilterResult
sdp SdpFilterResult

Before processing filter results you need to check two fields: execution_state and match_state. These fields are defined for all types of filter results except for SdpFilterResult. You would expect to have the execution state set to FilterExecutionState.EXECUTION_SUCCESS unless you process the partially successful invocation. You need to process each filter result whose match state is set to FilterMatchState.MATCH_FOUND. Note that there can be more than one match in the filter_results map.
For the Sensitive Data Protection (SDP) filter you will need to inspect the deidentify_result field and update the user prompt (or model response) using the deidentified version if such version is found. The values in this filter’s result depend on how the SDP filter is defined in the Model Armor template that was used. The following code snippet checks Malicious URIs and SDP filters and either returns a predefined response or the deidentified model response.

from google.cloud.modelarmor_v1 import (
    FilterExecutionState,
    FilterMatchState,
    InvocationResult,
    MaliciousUriFilterResult,
    SanitizationResult,
    SdpFilterResult,
)

# result: modelarmor_v1.SanitizationResult
if (result.invocation_result == InvocationResult.SUCCESS and
    result.filter_match_state == FilterMatchState.MATCH_FOUND):
    sdp_result:SdpFilterResult = result.filter_results.get("sdp", None)
    if sdp_result and sdp_result.deidentify_result:
       deidentify_result = sdp_result.deidentify_result
       if (deidentify_result.execution_state == FilterExecutionState.EXECUTION_SUCCESS and
           deidentify_result.match_state == FilterMatchState.MATCH_FOUND):
           return deidentify_result.data.text
    muri_result:MaliciousUriFilterResult = result.filter_results.get("malicious_uris", None)
    if (muri_result and
        muri_result.execution_state == FilterExecutionState.EXECUTION_SUCCESS and
        muri_result.match_state == FilterMatchState.MATCH_FOUND):
        return "The response contains malicious URIs"

Best practices: Google ADK

ADK functions through an event-based architecture where every component interaction is treated as a discrete event. Callbacks provide a mechanism to intercept these events, allowing you to execute custom logic during the agent’s lifecycle. You can implement these either by defining a callback function or by creating a plugin. My article, Master ADK Callbacks: DOs and DON’Ts, provides guidance on determining the right approach for your specific needs. The following sections will demonstrate how to leverage these techniques to invoke the Model Armor API.

Which events to hook to

ADK has many events. The ones that you want to hook to in order to call Model Armor API are before_model_callback and after_model_callback. Why only these two events? There is, for example, before_tool_callback that is triggered before ADK calls a tool based on the “function call” response from a model. There are two reasons:

  • These event hooks are a natural match to Model Armor APIs SanitizeUserPrompt() and SanitizeModelResponse() making it good fit to the event workflow and guarantee that the agent does not pass to its model insecure prompt or uses a response that might yield sensitive information or inappropriate content.
  • In the majority of use cases using Model Armor for tool calls is redundant since the results are returned to the model and the model’s final response is filtered in the after_model_callback trigger.

There is an edge case when the model can decide to pass the sensitive data returned by one tool to another tool. In this case, you would want to intercept the call and obfuscate the sensitive data. In that case you would need to use Sensitive Data Protection (SDP) service that is specialized on masking and de-identifying sensitive data instead of Model Armor.

Callback or plugin

ADK provides two methods to hook up to these events: callbacks and plugins. Following the DOs and DON’Ts guidelines, if you want to protect a distinct agent with Model Armor, you should call Model Armor API using the callbacks. Hook up your implementation in the agent’s definition:

from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.llm_agent import Agent
from google.adk.models import LlmResponse, LlmRequest

def after_model_hook(
   callback_context: CallbackContext, llm_response: LlmResponse
) -> Optional[LlmResponse]:
    # ... callback implementation here

def before_model_hook(
   callback_context: CallbackContext, llm_request: LlmRequest
) -> Optional[LlmResponse]:
    # ... callback implementation here

agent = Agent(
    name="example_agent",
    model="gemini-3-flash-preview",
#   description="A helpful assistant for user questions.",
#   instruction="Answer user questions to the best of your knowledge",
    after_model_callback=after_model_hook,
    before_model_callback=before_model_hook,
)

If you want to expand the protection to all agents that run, you can do it using ADK Runner. You can define plugins using the App class constructor. However, this is less practical if you use runtime resources which might not be available at the time the App class instance is created. You can look at the safety plugins example at Github for the sample implementation of the plugin that calls Model Armor API.

Enforcing sanitized result

You can use the methods shown In Interpreting sanitize API responses to follow the callback guidelines for enforcing the results of the sanitize API calls. There are three options:

  1. Continue"as is". This option should be selected when the sanitize API call was successful and has no matches
  2. Continue after altering the parameters. This option should be selected when the sanitize API call was successful and has matches that allow it to continue.
  3. Returning the controllable response replacing the model generated response. You should choose this option when the API call has errors or the found matches prevent using the user prompt or model response.

How do you do it? For the first option, just return None:

from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.llm_agent import Agent
from google.adk.models import LlmResponse, LlmRequest
from google.cloud.modelarmor_v1 import (
    FilterMatchState,
    InvocationResult,
)

def before_model_hook(
   callback_context: CallbackContext, llm_request: LlmRequest
) -> Optional[LlmResponse]:
    #... implement resp = modelarmor_v1.sanitize_user_prompt()
    result = resp.sanitization_result
    if (result.invocation_result == InvocationResult.SUCCESS and
        result.filter_match_state == FilterMatchState.NO_MATCH_FOUND):
        return None
    #...continue processing

For the second option you also return None, but you’ll need to alter the input argument. You should do it for both before_ and after_ callbacks! The following code snippet alters the user prompt based on the SDP filter deidentified result:


from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.llm_agent import Agent
from google.adk.models import LlmResponse, LlmRequest
from google.cloud.modelarmor_v1 import (
    FilterExecutionState,
    FilterMatchState,
    InvocationResult,
)
from google.genai import types as genai_types

def before_model_hook(
   callback_context: CallbackContext, llm_request: LlmRequest
) -> Optional[LlmResponse]:
    #... implement resp = modelarmor_v1.sanitize_user_prompt()
    result = resp.sanitization_result
    if (result.invocation_result == InvocationResult.SUCCESS and
        result.filter_match_state == FilterMatchState.MATCH_FOUND):
        sdp_result = result.filter_results.get("sdp", None)
        if sdp_result and sdp_result.deidentify_result:
            deidentify_result = sdp_result.deidentify_result
            if (deidentify_result.execution_state == FilterExecutionState.EXECUTION_SUCCESS and
                deidentify_result.match_state == FilterMatchState.MATCH_FOUND):
                llm_request.contents[-1] = genai_types.Content(
                       role="user",
                       parts=[types.Part.from_text(text=deidentify_result.data.text)],
                )
                return None
    #...continue processing

Finally, for the third option you return your own response instead. In the following code snippet I will use a case that handles failed invocation:

from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.llm_agent import Agent
from google.adk.models import LlmResponse, LlmRequest
from google.cloud.modelarmor_v1 import (
    InvocationResult,
)

def before_model_hook(
   callback_context: CallbackContext, llm_request: LlmRequest
) -> Optional[LlmResponse]:
    #... implement resp = modelarmor_v1.sanitize_user_prompt()
    result = resp.sanitization_result
    if result.invocation_result == InvocationResult.FAILURE:
        error_message = f"FAILURE: ({result.sanitization_metadata.error_code}) {result.sanitization_metadata.error_message}"
        return LlmResponse(
                    content=types.Content(
                        role="model", parts=[types.Part.from_text(text=error_message)]
                    )
                )
    #...continue processing

About Model Armor integration with Agent Platform

Agent Platform (former Vertex AI) integration of Model Armor allows customization of the Model Armor template to be invoked during a content generation call to Gemini. In order to do it, you need to modify your agent initialization to provide the template ID. The implementation uses the model_armor_config field of the GenerateContentConfig type. You can use other fields of the type to customize temperature, the size of the model response (in tokens) and other content generation parameters.

from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.llm_agent import Agent
from google.adk.models import LlmResponse, LlmRequest
from google.genai import types as genai_types

my_generation_config = genai_types.GenerateContentConfig(
    model_armor_config=genai_types.ModelArmorConfig(
        prompt_template_name=f"projects/PROJECT_ID/locations/REGION_ID/templates/PROMPT_TEMPLATE_ID",
        response_template_name=f"projects/PROJECT_ID/locations/REGION_ID/templates/RESPONSE_TEMPLATE_ID",
    ),
)

agent = Agent(
    name="example_agent",
    model="gemini-3-flash-preview",
    generate_content_config=my_generation_config,
#   description="A helpful assistant for user questions.",
#   instruction="Answer user questions to the best of your knowledge",
)

Best Practices: LangChain

LangChain does not have a notion of plugins. It supports callbacks through CallbackHandlers for observation of the chain execution. Since we need the ability to change input parameters or to cancel following mode invocation completely, this method does not work. Instead we can use the following two methods. I will quickly review the pros and cons of each.

Wrap up integration model call as a single chain

You can use the @chain attribute for a wrapper method that implements the model call with two sanitize calls to Model Armor API.

from google.cloud.modelarmor_v1 import (
    DataItem,
    FilterExecutionState,
    FilterMatchState,
    InvocationResult,
    ModelArmorClient,
    SanitizeModelResponseRequest,
    SanitizeUserPromptRequest,
)
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.runnables import chain

model = ChatGoogleGenerativeAI(model="gemini-3-flash-preview")
ma_client = ModelArmorClient()

@chain
def secure_model_chain(user_prompt):
    prompt_template_name = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/templates/{PROMPT_TEMPLATE_ID}"
    response_template_name = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/templates/{RESPONSE_TEMPLATE_ID}"
    
    # 1: Sanitize User Input
    request = SanitizeUserPromptRequest(
        name=prompt_template_name,
        user_prompt_data=DataItem(text=user_prompt.content),
    )
    try:
        response = ma_client.sanitize_user_prompt(request=request)
        result = response.sanitization_result
    except Exception as e:
        return AIMessage(content="[SYSTEM ERROR]: Model Armor check failed. Request aborted.")
    if result.invocation_result == InvocationResult.FAILURE:
        error_message = f"[SYSTEM ERROR]: ({result.sanitization_metadata.error_code}) {result.sanitization_metadata.error_message}"
        return AIMessage(content=error_message)
    if result.invocation_result == InvocationResult.SUCCESS:
        if result.filter_match_state == FilterMatchState.MATCH_FOUND:
            sdp_result = result.filter_results.get("sdp", None)
            if sdp_result and sdp_result.deidentify_result:
                deidentify_result = sdp_result.deidentify_result
                if (deidentify_result.execution_state == FilterExecutionState.EXECUTION_SUCCESS and
                    deidentify_result.match_state == FilterMatchState.MATCH_FOUND):
                    user_prompt.content = deidentify_result.data.text
    
    # 2: Call Model
    model_response = model.invoke(user_prompt)

    # 3: Sanitize Model Response
    request =SanitizeModelResponseRequest(
        name=response_template_name,
        user_prompt=user_prompt,
        model_response_data=DataItem(text=model_response.content),
    )
    try:
        response = ma_client.sanitize_model_response(request=request)
        result = response.sanitization_result
    except Exception as e:
        return AIMessage(content="[SYSTEM ERROR]: Model Armor check failed. Request aborted.")
    if result.invocation_result == InvocationResult.FAILURE:
        error_message = f"[SYSTEM ERROR]: ({result.sanitization_metadata.error_code}) {result.sanitization_metadata.error_message}"
        return AIMessage(content=error_message)
    if result.invocation_result == InvocationResult.SUCCESS:
        if result.filter_match_state == FilterMatchState.MATCH_FOUND:
            sdp_result = result.filter_results.get("sdp", None)
            if sdp_result and sdp_result.deidentify_result:
                deidentify_result = sdp_result.deidentify_result
                if (deidentify_result.execution_state == FilterExecutionState.EXECUTION_SUCCESS and
                    deidentify_result.match_state == FilterMatchState.MATCH_FOUND):
                    return [HumanMessage(content=deidentify_result.data.text)]
    # example does not check other filters
    return [HumanMessage(content=model_response.content)]

Model initialization and chain invocation are done in the standard way. By wrapping the logic in @chain, the secure_model_chain object becomes a Runnable, , LangChain tracks this as a distinct step in its internal tracing (e.g., in LangSmith).

Use RunnableBranch to “interrupt” a complex chain

The drawback of the previous method is that Model Armor API calls and model invocation become a one single chain. It might not be good for observability scenarios and when you want to separate between different logical steps (chains). To break a complex chain into granular, trackable steps while retaining the ability to “interrupt” or “short-circuit” the flow, the authentic LangChain approach is to compose multiple named Runnables and use a RunnableBranch.

from google.cloud.modelarmor_v1 import (
    DataItem,
    FilterExecutionState,
    FilterMatchState,
    InvocationResult,
    ModelArmorClient,
    SanitizeUserPromptRequest,
)
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.runnables import RunnableBranch, chain

model = ChatGoogleGenerativeAI(model="gemini-3-flash-preview")
ma_client = ModelArmorClient()
prompt_template_name = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/templates/{PROMPT_TEMPLATE_ID}"
response_template_name = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/templates/{RESPONSE_TEMPLATE_ID}"

@chain
def before_model_hook(user_prompt):
    request = SanitizeUserPromptRequest(
        name=prompt_template_name,
        user_prompt_data=DataItem(text=user_prompt.content),
    )
    try:
        response = ma_client.sanitize_user_prompt(request=request)
        result = response.sanitization_result
    except Exception as e:
        return {"status": "error", "error_message": "[SYSTEM ERROR]: Model Armor check failed. Request aborted."}
    if result.invocation_result == InvocationResult.FAILURE:
        return {
           "status": "error",
           "error_message": f"[SYSTEM ERROR]: ({result.sanitization_metadata.error_code}) {result.sanitization_metadata.error_message}"
        }
    if result.invocation_result == InvocationResult.SUCCESS:
        if result.filter_match_state == FilterMatchState.MATCH_FOUND:
            sdp_result = result.filter_results.get("sdp", None)
            if sdp_result and sdp_result.deidentify_result:
                deidentify_result = sdp_result.deidentify_result
                if (deidentify_result.execution_state == FilterExecutionState.EXECUTION_SUCCESS and
                    deidentify_result.match_state == FilterMatchState.MATCH_FOUND):
                    user_prompt.content = deidentify_result.data.text
    return {"status": "safe", "payload": user_prompt}

@chain
def after_model_hook(ai_response):
    request =SanitizeModelResponseRequest(
        name=response_template_name,
        model_response_data=DataItem(text=ai_response.content),
    )
    try:
        response = ma_client.sanitize_model_response(request=request)
        result = response.sanitization_result
    except Exception as e:
        return AIMessage(content="[SYSTEM ERROR]: Model Armor check failed. Request aborted.")
    if result.invocation_result == InvocationResult.FAILURE:
        error_message = f"[SYSTEM ERROR]: ({result.sanitization_metadata.error_code}) {result.sanitization_metadata.error_message}"
        return AIMessage(content=error_message)
    if result.invocation_result == InvocationResult.SUCCESS:
        if result.filter_match_state == FilterMatchState.MATCH_FOUND:
            sdp_result = result.filter_results.get("sdp", None)
            if sdp_result and sdp_result.deidentify_result:
                deidentify_result = sdp_result.deidentify_result
                if (deidentify_result.execution_state == FilterExecutionState.EXECUTION_SUCCESS and
                    deidentify_result.match_state == FilterMatchState.MATCH_FOUND):
                    return [AIMessage(content=deidentify_result.data.text)]
    # example does not check other filters
    return ai_response

Set up the chain and the branching logic using the following command. Then you can use the standard way to invoke the chain.

conditional_execution = RunnableBranch(
    (
        # CONDITION: If safe, run the model AND the after_hook
        lambda x: x["status"] == "safe", 
        (lambda x: x["payload"]) | model | after_model_hook 
    ),
    (
        # CONDITION: If explicitly blocked
        lambda x: x["status"] == "blocked", 
        lambda x: AIMessage(content=x["error_message"])
    ),
    # DEFAULT: Handle system/API errors
    lambda x: AIMessage(content="[SYSTEM ERROR]: Security verification unavailable.")
)

# The logic flows: Hook -> Branch (Decide) -> [Model -> After] OR [Block]
full_secure_chain = before_model_hook | conditional_execution

LangChain Community solution

There are many lines of code to invoke Model Armor. You can use the Model Armor implementation posted in the LangChain community instead of coding it yourself. I look into the source code and it seems good. The main difference between the community implementation and my examples is that the community implementation does not alter the input if the sensitive data within was sanitized. It interrupts the chain by throwing a ValueError exception which seems to be less graceful compared to using the RunnableBranch.

Summary: Model Armor is Framework Agnostic

I decided not to show the LangGraph implementation because essentially, besides using the nodes to call Model Armor API and interpret the response and Conditional Edges to “short-circuit” the flow the integration looks the same as with ADK and LangChain.
These examples demonstrate that the integration of Model Armor is essentially framework agnostic. You use the same logic to invoke Model Armor API and interpret its results regardless which framework you use for your agents. The key is to wrap agent’s calls to the model with Model Armor sanitize_* calls. Then you use constructs of the framework to intercept and, if needed, interrupt the flow.
Model Armor direct invocation at application-level provides flexibility. However, this flexibility comes with a price tag. If you decide to use this approach, consider the potential friction points that can become blockers:

  • Introduces complexity to the application code that is unrelated to the business workflow. Lifts on-boarding requirements for developers and maintainers.
  • Increases security risks by adding dependencies to the runtimes thus expanding the attack surface.
  • Increase maintenance costs due to the need to cover Model Armor invocation into testing framework. Also increase the volume of the application’s observability data.
  • Hinders centralized enforcement of the security policies and may completely obstruct it. Tasks like modifying Model Armor templates (e.g. ensuring the use of the template closest to the agent’s runtime deployed location) or altering response interpretation coded in the application cannot be effectively changed without modifying application code.

To address these friction points you can review alternative integration solutions provided by Agent Platform, Agent gateway or 3P products that allow you to route inference calls to Model Armor before they reach the model and also sanitize the model response outside of the application code.

Remember that Security is a First-Class Citizen. It must be considered throughout development, and developers should remain vigilant regarding security risks, even when primary security measures are implemented outside of the application code.