Guardrail implementation

[devevignesh]

Hello,

We have implemented a guardrail system to preprocess user inputs before they reach our root agent. Our current guardrail is designed to intercept requests to the root agent and apply a set of checks using a dedicated "parallel guardrail agent," based on the workflow agents documentation.

Our primary challenge is identifying off-topic inputs, so we want to use an LLM to detect them.

Since the documentation does not provide examples for advanced guardrail use cases, I would like to know the best approach. If you could share any documentation or feedback, it would be appreciated.

# In agents
root_agent = LlmAgent(
    name="root_agent",
    model=settings.MODEL_GEMINI_2_0_FLASH,
    description="",
    instruction=prompt_old.ROOT_AGENT_INSTR,
    tools=[...],
    sub_agents=[...],
    output_key="final_response",
    before_agent_callback=_load_user_profile,
    before_model_callback=_before_model_callback,
    generate_content_config={
        "safety_settings": [
            {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_LOW_AND_ABOVE"},
            {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
            {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_LOW_AND_ABOVE"},
            {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
        ]
    }
)


# In sub_agents/guard_rails/tools.py
from ai_coach.sub_agents.guard_rails.runner import parallel_guardrail_runner
from google.adk.agents.llm_agent import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.agents.callback_context import CallbackContext
from google.genai import types

async def _before_model_callback(callback_context: CallbackContext,
                                 llm_request: LlmRequest) -> LlmResponse | None:
    user_text = None
    if llm_request.contents:
        last_content = llm_request.contents[-1]
        if last_content.role == "user" and hasattr(last_content.parts[0], 'text'):
            user_text = last_content.parts[0].text

    if user_text:
        user_content = types.Content(role='user', parts=[types.Part(text=user_text)])
        state = callback_context.state
        user_id = str(state.get("user_id", "unknown_user"))
        session_id = state.get("session_id", "unknown_session")
        final_guardrail_response = "false" # Default to allow

        async for event in parallel_guardrail_runner.run_async(
                user_id=user_id,
                session_id=session_id,
                new_message=user_content
            ):
            if event.is_final_response() and event.content and event.content.parts:
                for part in event.content.parts:
                    if hasattr(part, 'text') and part.text is not None:
                        final_guardrail_response = part.text
                        break
                break
        
        if final_guardrail_response == "true":
            # Guardrail triggered, block main LLM
            return LlmResponse(
                content=types.Content(
                    parts=[types.Part(text="Input blocked by guardrails.")],
                    role="model"
                )
            )
    # No guardrail triggered or no user text to check
    return None