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.