Claret Ibeawuchi

Posted on Oct 29

Callback the Police: Enforcing Business Rules in AI Agents 👮‍♂️

#programming #ai #agents #python

The Traffic Cop Analogy
What Are Callbacks in AI Agents?
The before_tool_callback: Your First Line of Defense
Real-World Pattern: Admin Payment Exemption
Real-World Pattern: Tool Interception and Routing
Real-World Pattern: Auto-Injecting Context
Real-World Pattern: Detecting Admin Message Leaks
Common Mistakes and How to Avoid Them
When to Use Which Callback
Key Takeaways
Resources

The Traffic Cop Analogy

Imagine a busy intersection without traffic lights or stop signs, just drivers making their own decisions. Most would probably stop and check. Some might slow down. A few might just go, assuming others will yield.

That's your AI agent without callbacks.

Now add a traffic cop who:

Blocks unsafe moves before they happen
Redirects traffic when routes are closed
Enforces rules consistently, regardless of driver intent

That's your AI agent with callbacks.

The difference? One relies on hoping drivers (the LLM) make good decisions. The other enforces good decisions with authority.

What Are Callbacks in AI Agents?

Callbacks are functions that run at specific points in your agent's lifecycle, giving you programmatic control over agent behavior.

Think of them as hooks where you can:

✅ Intercept actions before they execute
✅ Modify parameters or route to different functions
✅ Block unauthorized operations
✅ Inject additional context
✅ Validate outputs before users see them

The Callback Lifecycle

User: "Create an order for me"
         ↓
    [LLM reasoning...]
         ↓
    Decides to call: process_payment(amount=250)
         ↓
    ┌────────────────────────────────────┐
    │  before_tool_callback()            │  ← YOU ARE HERE
    │  "Should this be allowed?"         │    (Pre-execution)
    └────────────────────────────────────┘
         ↓
    [Tool executes: process_payment()]
         ↓
    [LLM generates response...]
         ↓
    ┌────────────────────────────────────┐
    │  after_model_callback()            │  ← YOU ARE HERE
    │  "Is this response correct?"       │    (Post-generation)
    └────────────────────────────────────┘
         ↓
    User sees: "Payment of $250 processed successfully"

Callbacks sit between the LLM's decisions and their execution/output, giving you the final say.

The before_tool_callback: Your First Line of Defense

The before_tool_callback runs before every tool call the LLM tries to make. This is where you enforce business rules.

Signature and Components

def before_tool_callback(
    tool: BaseTool,           # The tool about to be called
    args: dict,               # Arguments the LLM provided
    tool_context: ToolContext # Session state, auth token, etc.
) -> Optional[dict]:
    """
    Returns:
    - None: Proceed with normal tool execution
    - dict: Skip execution and return this result instead
    """
    pass

What You Can Do

def before_tool_callback(tool, args, tool_context):
    tool_name = tool.func.__name__

    # 1. BLOCK unauthorized actions
    if tool_name == "delete_user" and not is_admin(tool_context):
        return {
            "error": True,
            "message": "Unauthorized: Admin access required"
        }

    # 2. ROUTE to different implementations
    if tool_name == "get_price":
        if is_premium_tier(tool_context):
            return calculate_premium_price(**args)
        else:
            return calculate_standard_price(**args)

    # 3. MODIFY parameters before execution
    if tool_name == "search_database":
        # Add tenant filter for multi-tenant systems
        args["tenant_id"] = tool_context.state.get("tenant_id")

    # 4. INJECT missing context
    if tool_name == "create_order" and not args.get("user_id"):
        args["user_id"] = tool_context.state.get("current_user_id")

    return None  # Proceed with (possibly modified) execution

The callback can:

Return a result → Skip tool execution entirely, use your result
Modify args in place → Tool executes with modified parameters
Return None → Tool executes normally

Real-World Pattern: Admin Payment Exemption

Problem: Our agent was asking admin users to pay for services. Admins should create orders for free on behalf of customers.

Why it happened: The LLM doesn't understand business roles, it just follows the general user flow.

Solution: Intercept payment tool calls and block them for admin users.

Implementation

def before_tool_callback(tool, args, tool_context):
    tool_name = tool.func.__name__

    # Get or cache admin status
    is_admin = tool_context.state.get("session:is_admin")

    if is_admin is None:
        # First time - check auth token and cache result
        auth_token = tool_context.state.get("auth_token")
        is_admin = check_admin_status(auth_token)
        tool_context.state["session:is_admin"] = is_admin  # Cache it

    # ========================================================================
    # ADMIN PAYMENT EXEMPTION
    # Block payment tools for admin users - admins create orders for free
    # ========================================================================
    if tool_name in ["process_payment", "charge_card"]:
        if is_admin:
            logger.info(f"🔐 ADMIN USER: Blocking {tool_name} - no payment required")

            # Mark session as payment exempt
            tool_context.state["session:payment_status"] = "admin_exempt"

            # Return success WITHOUT executing payment
            return {
                "success": True,
                "message": "✅ Admin user - no payment required. You can proceed immediately.",
                "admin_exemption": True,
                "payment_status": "admin_exempt"
            }

    return None  # Non-payment tools execute normally

What Happens

Before callbacks:

Admin: "Create an order"
Agent: "I'll need to process your payment first..."
Admin: /frustration/ "I'm an admin!"

After callbacks:

Admin: "Create an order"
[Callback intercepts payment tool call]
[Returns success without executing payment]
Agent: "✅ Admin user - no payment required. Creating your order now..."
Admin: /happy/

Key Insights

Admin status is cached in session state, we only check once per session
Multiple payment tool names are covered (any payment-related tool)
Logs include context (🔐 emoji for security actions)
Session state updated to track exemption status
Clean response that explains why (for debugging) without confusing the user

Performance Optimization: Caching Expensive Checks

def before_tool_callback(tool, args, tool_context):
    # Check cache first
    is_admin = tool_context.state.get("session:is_admin")

    if is_admin is None:
        # First time - check auth and cache result
        auth_token = tool_context.state.get("auth_token")
        is_admin = check_admin_status(auth_token)  # Expensive: JWT decode + lookup

        # Cache for entire session
        tool_context.state["session:is_admin"] = is_admin
        logger.info(f"🔐 Cached admin status: {is_admin}")

    # Use cached value
    if tool_name in PAYMENT_TOOLS and is_admin:
        return {"success": True, "admin_exempt": True}

    return None

Performance impact:

Before: Auth check on every tool call (~50-100ms each × 10 calls = 500-1000ms)
After: Auth check once per session (~50ms total)

What to cache:

✅ User roles/permissions
✅ Tenant/organization IDs
✅ Feature flags
❌ Dynamic data (current balance, order status)

Real-World Pattern: Tool Interception and Routing

Problem: The LLM sometimes calls the wrong pricing tool:

Calls regular pricing for admin users → Wrong (too expensive)
Calls admin pricing for regular users → Security violation

Why it happened: The LLM chooses tools based on context and phrasing, which is inconsistent.

Solution: Intercept pricing tool calls and route to the correct implementation.

Implementation

def before_tool_callback(tool, args, tool_context):
    tool_name = tool.func.__name__

    # Get user type from session
    is_admin = tool_context.state.get("session:is_admin", False)

    # ========================================================================
    # DETERMINISTIC PRICING ROUTING
    # If LLM calls wrong pricing tool, redirect to correct one
    # ========================================================================

    if tool_name == "get_standard_price" and is_admin:
        # Admin called standard pricing → Route to premium
        logger.warning(
            f"🔄 LLM called standard pricing for ADMIN user - "
            f"intercepting and routing to premium pricing"
        )

        # Import and call correct function
        from app.tools.pricing import get_premium_price

        # Pass through all args + context
        args["tool_context"] = tool_context
        return get_premium_price(**args)

    elif tool_name == "get_premium_price" and not is_admin:
        # Regular user called admin pricing → Security violation, route to standard
        logger.warning(
            f"🔄 LLM called premium pricing for REGULAR user - "
            f"intercepting and routing to standard pricing"
        )

        from app.tools.pricing import get_standard_price
        args["tool_context"] = tool_context
        return get_standard_price(**args)

    return None

What Happens

Regular User: "How much for this service?"
LLM: [Decides to call get_premium_price() - WRONG TOOL]
      ↓
[Callback intercepts]
      ↓
[Routes to get_standard_price() instead]
      ↓
User sees correct standard pricing ✅

Logs show:
"🔄 LLM called premium pricing for REGULAR user - intercepting and routing"

Benefits

Zero reliance on LLM choosing the right tool
Security enforced in code, not prompts
Transparent logging of routing decisions
No user-visible errors - routing happens silently
Consistent behavior regardless of how user phrases request

Real-World Pattern: Auto-Injecting Context

Problem: After getting a quote, users would say "proceed with order" and the agent would re-ask for all the parameters, even though all that information was already collected.

Why it happened: The LLM didn't pass previous context to the order creation tool.

Solution: Auto-inject stored quote details when order tools are called.

Implementation

def before_tool_callback(tool, args, tool_context):
    tool_name = tool.func.__name__

    # ========================================================================
    # AUTO-INJECT QUOTE DETAILS INTO ORDER
    # If order tool is called without params, fill from stored quote
    # ========================================================================

    if tool_name == "create_order":
        # Get previously stored quote from session
        latest_quote = tool_context.state.get("session:latest_quote")

        if latest_quote:
            # Auto-inject missing required parameters
            for key in ["param_a", "param_b", "param_c", "total_price"]:
                if not args.get(key) and latest_quote.get(key):
                    args[key] = latest_quote[key]
                    logger.info(f"✅ Auto-injected {key}: {args[key]}")

    return None  # Proceed with augmented args

Storing Quote Data

When a quote is generated, store it:

def calculate_quote(params, tool_context):
    # Calculate quote
    price = api.get_price(params)

    # Store complete quote in session state
    quote_data = {
        "param_a": params.get("param_a"),
        "param_b": params.get("param_b"),
        "param_c": params.get("param_c"),
        "total_price": price,
        "timestamp": time.time()
    }

    tool_context.state["session:latest_quote"] = quote_data

    return {"price": price, ...}

User Experience

Before callbacks:

User: "Quote for service from A to B"
Agent: "That'll be $1,200"
User: "Book it"
Agent: "I'll need parameter A"
User: "I just told you!"
Agent: "I need the exact value..."

After callbacks:

User: "Quote for service from A to B"
Agent: "That'll be $1,200"
[Quote stored in session state]
User: "Book it"
[Callback injects all required params from session]
Agent: "Great! Creating your order now..."

Logs show:

✅ Auto-injected param_a: value_a
✅ Auto-injected param_b: value_b
✅ Auto-injected param_c: value_c
✅ Auto-injected price: $1200.00

Key Insights

Session state bridges turns - Data persists across messages
LLM doesn't need to remember - We do it for them
Only inject if missing - Don't override if LLM provided it
Log every injection - Helps debugging and monitoring
Structured storage - Use consistent keys for retrieval

Common Mistakes and How to Avoid Them

❌ Mistake 1: Not Returning Anything

def before_tool_callback(tool, args, tool_context):
    if should_block_tool(tool):
        logger.error("Blocking tool")
        # BUG: Forgot to return! Tool will still execute

Fix:

def before_tool_callback(tool, args, tool_context):
    if should_block_tool(tool):
        logger.error("Blocking tool")
        return {"error": True, "message": "Unauthorized"}  # ✅

❌ Mistake 2: Modifying Args Without In-Place Updates

def before_tool_callback(tool, args, tool_context):
    # BUG: Creates new dict, doesn't modify original
    args = {**args, "tenant_id": "abc123"}
    return None

Fix:

def before_tool_callback(tool, args, tool_context):
    # ✅ Modifies dict in-place
    args["tenant_id"] = "abc123"
    return None

❌ Mistake 3: Forgetting Tool Context Parameter

# Original tool signature
def my_tool(param1, param2):
    pass

# In callback - BUG: Doesn't pass tool_context
return my_other_tool(param1=args["param1"])

Fix:

# If tool needs context, include it
return my_other_tool(
    param1=args["param1"],
    tool_context=tool_context  # ✅
)

❌ Mistake 4: Missing Tool Name Variants

def before_tool_callback(tool, args, tool_context):
    # BUG: Only catches one tool name, misses variants
    if tool.func.__name__ == "charge_payment":
        if is_admin(tool_context):
            return {"success": True, "admin_exempt": True}

Problem: Your system might have multiple payment-related tools:

charge_payment
initiate_payment
process_card_payment

Missing one means the callback doesn't catch it!

Fix:

# Define at module level for easy maintenance
PAYMENT_TOOLS = [
    "charge_payment",
    "initiate_payment", 
    "process_card_payment",
    "charge_card"
]

def before_tool_callback(tool, args, tool_context):
    tool_name = tool.func.__name__

    # ✅ Catches all variants
    if tool_name in PAYMENT_TOOLS:
        if is_admin(tool_context):
            return {"success": True, "admin_exempt": True}

    return None

Maintenance tip: When adding new tools, update the list once. All callbacks automatically cover it.

❌ Mistake 5: Catching All Tools Too Broadly

def before_tool_callback(tool, args, tool_context):
    # BUG: Every tool will get user_id added
    args["user_id"] = get_user_id(tool_context)
    return None

Fix:

def before_tool_callback(tool, args, tool_context):
    # ✅ Only inject for tools that need it
    if tool.func.__name__ in ["create_order", "update_profile"]:
        args["user_id"] = get_user_id(tool_context)
    return None

❌ Mistake 6: No Logging

def before_tool_callback(tool, args, tool_context):
    if is_admin(tool_context):
        return admin_version(**args)
    return None

Fix:

def before_tool_callback(tool, args, tool_context):
    if is_admin(tool_context):
        logger.info(f"🔄 Routing {tool.func.__name__} to admin version")
        return admin_version(**args)
    return None

Logging is critical for debugging "why did this happen?" in production.

Real-World Pattern: Detecting Admin Message Leaks

Even with perfect tool routing and parameter injection, the LLM can still leak admin-only information in its natural language responses.

The Problem

Admin sees:
"Using admin access for premium pricing... Your cost: $1,200"

Non-admin sees:
"Using admin access for premium pricing... Total: $1,500" ❌

The tool correctly returned different prices, but the LLM copy-pasted admin messaging into non-admin responses!

The Solution: Response Content Validation

Use after_model_callback to detect and clean role-specific messaging:

def after_model_callback(callback_context, llm_response):
    """Detect and remove admin-only phrases from non-admin responses"""

    # Get user role from session
    is_admin = callback_context.state.get("session:is_admin", False)

    if is_admin:
        return None  # Admin can see admin messages

    # Extract response text
    response_text = llm_response.content.parts[0].text

    # Define admin-only phrases
    admin_only_phrases = [
        "using admin access",
        "premium pricing",
        "admin access for",
        "your premium cost"
    ]

    # Check if any admin phrase leaked
    has_admin_phrase = any(
        phrase.lower() in response_text.lower() 
        for phrase in admin_only_phrases
    )

    if has_admin_phrase:
        logger.warning(
            f"🚨 Admin message leaked to non-admin user! "
            f"Cleaning response..."
        )

        # Get correct price from session
        display_price = callback_context.state.get("session:last_display_price")

        # Override with clean, role-appropriate message
        clean_response = f"Total shipping cost: ${display_price:.2f} 🚚"

        from google.genai import types
        return LlmResponse(
            content=types.Content(
                role="model",
                parts=[types.Part(text=clean_response)]
            )
        )

    return None  # Response is clean

Why This Happens

LLMs learn patterns from context. If they recently saw:

Admin messages in earlier conversation
Admin-formatted responses in training data
Tool responses that mention "premium" or "admin"

They might echo these patterns even for non-admin users.

When to Use This Pattern

✅ Use when:

Different user roles see different information
LLM has access to role-specific tool responses
Sensitive information must never leak

❌ Don't use when:

All users see the same information
No role-based differentiation needed

Lesson learned: Callbacks protect against both wrong tool execution AND wrong LLM messaging. Defense in depth applies to language, not just code.

When to Use Which Callback

Google ADK provides several callback types. Here's when to use each:

`before_tool_callback` - Pre-Execution Enforcement

Use when:

✅ Blocking unauthorized actions
✅ Routing to different implementations
✅ Injecting authentication context
✅ Auto-filling parameters from session
✅ Validating parameters before execution

Example use cases:

Admin payment exemption
Multi-tenant data isolation
Rate limiting
Parameter validation

`after_model_callback` - Post-Generation Validation

Use when:

✅ Validating LLM output accuracy
✅ Sanitizing responses (removing sensitive data)
✅ Formatting output consistently
✅ Correcting known LLM mistakes

Example use cases:

Price validation (next article!)
PII detection and removal
Response format enforcement
Output translation

Comparison

Callback Type	When It Runs	What to Use It For
`before_tool_callback`	Before tool execution	Business rule enforcement, routing, injection
`after_tool_callback`	After tool execution	Result transformation, logging
`after_model_callback`	After LLM generates response	Output validation, correction

Most common pattern: Use before_tool_callback for enforcement and after_model_callback for validation.

Key Takeaways

Callbacks are enforcers - They don't ask, they tell
Use them for business logic - Don't rely on LLM prompt-following
Log everything - You'll need it for debugging
Cache expensive checks - Store results in session state
Return early for blocks - Don't execute what shouldn't happen
Modify args in-place - When you want execution to proceed with changes
Test edge cases - Callbacks run on EVERY tool call

The Enforcement Hierarchy

Strongest:   Code enforcement (callbacks)
             ↑
Medium:      Tool filtering by role
             ↑
Weakest:     LLM instructions

Put your critical business logic in callbacks, not instructions.

Resources

Previous: "How Reliable Are Your AI Agents?"

Next: "The Ground Truth Principle: Session State and Output Validation"

What business rules are you enforcing with callbacks? Share your patterns!

Table of Contents

The Traffic Cop Analogy

What Are Callbacks in AI Agents?

The Callback Lifecycle

The before_tool_callback: Your First Line of Defense

Signature and Components

What You Can Do

Real-World Pattern: Admin Payment Exemption

Implementation

What Happens

Key Insights

Performance Optimization: Caching Expensive Checks

Real-World Pattern: Tool Interception and Routing

Implementation

What Happens

Benefits

Real-World Pattern: Auto-Injecting Context

Implementation

Storing Quote Data

User Experience

Key Insights

Common Mistakes and How to Avoid Them

❌ Mistake 1: Not Returning Anything

❌ Mistake 2: Modifying Args Without In-Place Updates

❌ Mistake 3: Forgetting Tool Context Parameter

❌ Mistake 4: Missing Tool Name Variants

❌ Mistake 5: Catching All Tools Too Broadly

❌ Mistake 6: No Logging

Real-World Pattern: Detecting Admin Message Leaks

The Problem

The Solution: Response Content Validation

Why This Happens

When to Use This Pattern

When to Use Which Callback

before_tool_callback - Pre-Execution Enforcement

after_model_callback - Post-Generation Validation

Comparison

Key Takeaways

The Enforcement Hierarchy

Resources

`before_tool_callback` - Pre-Execution Enforcement

`after_model_callback` - Post-Generation Validation