Guardrail Classifier Agent
Source notebook
Repo path:
gemini/agents/genai-experience-concierge/agent-design-patterns/guardrail-classifier.ipynb· Open on GitHub · advanced
Builds a LangGraph guardrail classifier agent for Cymbal retail chat using Gemini on Vertex AI.
Summary
The notebook teaches how to add an LLM-based guardrail classifier before response generation in an agentic chat workflow. It builds a LangGraph with guardrails, chat, and post-processing nodes, using Gemini via the Google GenAI SDK on Vertex AI. It also generates valid and invalid test cases, runs them through the graph, and visualizes guardrail classification results with a confusion matrix.
Key code patterns
Vertex AI Gemini client
client = genai.Client(
vertexai=True,
project=agent_config.project,
location=agent_config.region,
)Configures Google GenAI SDK calls to use Vertex AI with an explicit project and region.
Structured guardrail output
response = await client.aio.models.generate_content(
model=agent_config.guardrail_model_name,
contents=contents,
config=genai_types.GenerateContentConfig(
response_mime_type="application/json",
response_schema=RequestClassification,
),
)
classification = RequestClassification.model_validate_json(response.text.strip())Forces the classifier to return parseable JSON matching a Pydantic schema.
Guardrail routing
next_node = CHAT_NODE_NAME
if current_turn["classification"].blocked:
stream_writer({"text": current_turn["response"]})
next_node = POST_PROCESS_NODE_NAME
return lg_types.Command(update=GraphSession(current_turn=current_turn), goto=next_node)Blocks invalid inputs before chat generation and routes allowed inputs to the chat node.
LangGraph compile
state_graph = graph.StateGraph(state_schema=GraphSession)
state_graph.add_node(GUARDRAILS_NODE_NAME, ainvoke_guardrails)
state_graph.add_node(CHAT_NODE_NAME, ainvoke_chat)
state_graph.add_node(POST_PROCESS_NODE_NAME, ainvoke_post_process)
state_graph.set_entry_point(GUARDRAILS_NODE_NAME)
compiled_graph = state_graph.compile(memory_checkpoint.MemorySaver())Defines a stateful agent workflow with memory-backed checkpoints.
Classifier evaluation
cf_matrix = pd.crosstab(
example_df["blocked_actual"],
example_df["blocked_label"],
)
sns.heatmap(cf_matrix, annot=True)Compares generated labels with classifier outputs to inspect guardrail performance.
Models & APIs used
- Models: gemini-3.5-flash
- APIs / services: Vertex AI
- SDKs / libraries:
google-genai,langgraph,langgraph-checkpoint-sqlite,pydantic,pandas,seaborn
When to use this
Use this pattern when an agent must classify and block out-of-scope or adversarial user inputs before generating a response.
Gotchas & caveats
- Requires PROJECT_ID or GOOGLE_CLOUD_PROJECT to be set.
- Uses REGION set to us-central1.
- Colab users must authenticate with google.colab.auth.authenticate_user().
- The notebook says restarting the runtime is required after installing dependencies.
- Sequential guardrail classification increases latency but can reduce cost by skipping answer generation for blocked inputs.
- The chat response is intentionally not grounded because the demo focuses on the guardrail classifier.
- The code initializes genai_client with PROJECT, while notebook parameters define PROJECT_ID.
Best practices
- Use a dedicated guardrails node before the chat node to decide whether to answer or block.
- Use temperature=0 and seed=0 for guardrail classification consistency.
- Return structured JSON with response_schema and validate it with Pydantic.
- Provide a safe guardrail_response for blocked requests.
- Store conversation history in state so the classifier and chat node can use prior turns.
- Stream blocked fallback text and chat response text through get_stream_writer().
- Evaluate the classifier with valid and invalid test cases and inspect mismatches.
Related
- Concepts: Agents & ADK · Function Calling & Tools · Evaluation
- Entities: Vertex AI · Google GenAI SDK · LangGraph · Gemini
- Area: Gemini Notebooks
- Best practices: Agents & ADK - Best Practices · Function Calling & Tools - Best Practices · Evaluation - Best Practices