feat: Feature/support tool callbacks in live mode

Author: hangfei

contributing/samples/live_tool_callbacks_agent/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from . import agent

contributing/samples/live_tool_callbacks_agent/agent.py

@@ -0,0 +1,269 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from datetime import datetime
+import random
+import time
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from google.adk import Agent
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+
+
+def get_weather(location: str, tool_context: ToolContext) -> Dict[str, Any]:
+  """Get weather information for a location.
+  Args:
+    location: The city or location to get weather for.
+  Returns:
+    A dictionary containing weather information.
+  """
+  # Simulate weather data
+  temperatures = [-10, -5, 0, 5, 10, 15, 20, 25, 30, 35]
+  conditions = ["sunny", "cloudy", "rainy", "snowy", "windy"]
+
+  return {
+      "location": location,
+      "temperature": random.choice(temperatures),
+      "condition": random.choice(conditions),
+      "humidity": random.randint(30, 90),
+      "timestamp": datetime.now().isoformat(),
+  }
+
+
+async def calculate_async(operation: str, x: float, y: float) -> Dict[str, Any]:
+  """Perform async mathematical calculations.
+  Args:
+    operation: The operation to perform (add, subtract, multiply, divide).
+    x: First number.
+    y: Second number.
+  Returns:
+    A dictionary containing the calculation result.
+  """
+  # Simulate some async work
+  await asyncio.sleep(0.1)
+
+  operations = {
+      "add": x + y,
+      "subtract": x - y,
+      "multiply": x * y,
+      "divide": x / y if y != 0 else float("inf"),
+  }
+
+  result = operations.get(operation.lower(), "Unknown operation")
+
+  return {
+      "operation": operation,
+      "x": x,
+      "y": y,
+      "result": result,
+      "timestamp": datetime.now().isoformat(),
+  }
+
+
+def log_activity(message: str, tool_context: ToolContext) -> Dict[str, str]:
+  """Log an activity message with timestamp.
+  Args:
+    message: The message to log.
+  Returns:
+    A dictionary confirming the log entry.
+  """
+  if "activity_log" not in tool_context.state:
+    tool_context.state["activity_log"] = []
+
+  log_entry = {"timestamp": datetime.now().isoformat(), "message": message}
+  tool_context.state["activity_log"].append(log_entry)
+
+  return {
+      "status": "logged",
+      "entry": log_entry,
+      "total_entries": len(tool_context.state["activity_log"]),
+  }
+
+
+# Before tool callbacks
+def before_tool_audit_callback(
+    tool, args: Dict[str, Any], tool_context: ToolContext
+) -> Optional[Dict[str, Any]]:
+  """Audit callback that logs all tool calls before execution."""
+  print(f"🔍 AUDIT: About to call tool '{tool.name}' with args: {args}")
+
+  # Add audit info to tool context state
+  if "audit_log" not in tool_context.state:
+    tool_context.state["audit_log"] = []
+
+  tool_context.state["audit_log"].append({
+      "type": "before_call",
+      "tool_name": tool.name,
+      "args": args,
+      "timestamp": datetime.now().isoformat(),
+  })
+
+  # Return None to allow normal tool execution
+  return None
+
+
+def before_tool_security_callback(
+    tool, args: Dict[str, Any], tool_context: ToolContext
+) -> Optional[Dict[str, Any]]:
+  """Security callback that can block certain tool calls."""
+  # Example: Block weather requests for restricted locations
+  if tool.name == "get_weather" and args.get("location", "").lower() in [
+      "classified",
+      "secret",
+  ]:
+    print(
+        "🚫 SECURITY: Blocked weather request for restricted location:"
+        f" {args.get('location')}"
+    )
+    return {
+        "error": "Access denied",
+        "reason": "Location access is restricted",
+        "requested_location": args.get("location"),
+    }
+
+  # Allow other calls to proceed
+  return None
+
+
+async def before_tool_async_callback(
+    tool, args: Dict[str, Any], tool_context: ToolContext
+) -> Optional[Dict[str, Any]]:
+  """Async before callback that can add preprocessing."""
+  print(f"⚡ ASYNC BEFORE: Processing tool '{tool.name}' asynchronously")
+
+  # Simulate some async preprocessing
+  await asyncio.sleep(0.05)
+
+  # For calculation tool, we could add validation
+  if (
+      tool.name == "calculate_async"
+      and args.get("operation") == "divide"
+      and args.get("y") == 0
+  ):
+    print("🚫 VALIDATION: Prevented division by zero")
+    return {
+        "error": "Division by zero",
+        "operation": args.get("operation"),
+        "x": args.get("x"),
+        "y": args.get("y"),
+    }
+
+  return None
+
+
+# After tool callbacks
+def after_tool_enhancement_callback(
+    tool,
+    args: Dict[str, Any],
+    tool_context: ToolContext,
+    tool_response: Dict[str, Any],
+) -> Optional[Dict[str, Any]]:
+  """Enhance tool responses with additional metadata."""
+  print(f"✨ ENHANCE: Adding metadata to response from '{tool.name}'")
+
+  # Add enhancement metadata
+  enhanced_response = tool_response.copy()
+  enhanced_response.update({
+      "enhanced": True,
+      "enhancement_timestamp": datetime.now().isoformat(),
+      "tool_name": tool.name,
+      "execution_context": "live_streaming",
+  })
+
+  return enhanced_response
+
+
+async def after_tool_async_callback(
+    tool,
+    args: Dict[str, Any],
+    tool_context: ToolContext,
+    tool_response: Dict[str, Any],
+) -> Optional[Dict[str, Any]]:
+  """Async after callback for post-processing."""
+  print(
+      f"🔄 ASYNC AFTER: Post-processing response from '{tool.name}'"
+      " asynchronously"
+  )
+
+  # Simulate async post-processing
+  await asyncio.sleep(0.05)
+
+  # Add async processing metadata
+  processed_response = tool_response.copy()
+  processed_response.update({
+      "async_processed": True,
+      "processing_time": "0.05s",
+      "processor": "async_after_callback",
+  })
+
+  return processed_response
+
+
+import asyncio
+
+# Create the agent with tool callbacks
+root_agent = Agent(
+    # model='gemini-2.0-flash-live-preview-04-09',  # for Vertex project
+    model="gemini-2.0-flash-live-001",  # for AI studio key
+    name="tool_callbacks_agent",
+    description=(
+        "Live streaming agent that demonstrates tool callbacks functionality. "
+        "It can get weather, perform calculations, and log activities while "
+        "showing how before and after tool callbacks work in live mode."
+    ),
+    instruction="""
+      You are a helpful assistant that can:
+      1. Get weather information for any location using the get_weather tool
+      2. Perform mathematical calculations using the calculate_async tool
+      3. Log activities using the log_activity tool
+      
+      Important behavioral notes:
+      - You have several callbacks that will be triggered before and after tool calls
+      - Before callbacks can audit, validate, or even block tool calls
+      - After callbacks can enhance or modify tool responses
+      - Some locations like "classified" or "secret" are restricted for weather requests
+      - Division by zero will be prevented by validation callbacks
+      - All your tool responses will be enhanced with additional metadata
+      
+      When users ask you to test callbacks, explain what's happening with the callback system.
+      Be conversational and explain the callback behavior you observe.
+    """,
+    tools=[
+        get_weather,
+        calculate_async,
+        log_activity,
+    ],
+    # Multiple before tool callbacks (will be processed in order until one returns a response)
+    before_tool_callback=[
+        before_tool_audit_callback,
+        before_tool_security_callback,
+        before_tool_async_callback,
+    ],
+    # Multiple after tool callbacks (will be processed in order until one returns a response)
+    after_tool_callback=[
+        after_tool_enhancement_callback,
+        after_tool_async_callback,
+    ],
+    generate_content_config=types.GenerateContentConfig(
+        safety_settings=[
+            types.SafetySetting(
+                category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+                threshold=types.HarmBlockThreshold.OFF,
+            ),
+        ]
+    ),
+)

contributing/samples/live_tool_callbacks_agent/readme.md

@@ -0,0 +1,94 @@
+# Live Tool Callbacks Agent
+
+This sample demonstrates how tool callbacks work in live (bidirectional streaming) mode. It showcases both `before_tool_callback` and `after_tool_callback` functionality with multiple callback chains, async callbacks, and various callback behaviors.
+
+## Features Demonstrated
+
+### Before Tool Callbacks
+1. **Audit Callback**: Logs all tool calls before execution
+2. **Security Callback**: Can block tool calls based on security rules (e.g., restricted locations)
+3. **Async Validation Callback**: Performs async validation and can prevent invalid operations
+
+### After Tool Callbacks
+1. **Enhancement Callback**: Adds metadata to tool responses
+2. **Async Post-processing Callback**: Performs async post-processing of responses
+
+### Tools Available
+- `get_weather`: Get weather information for any location
+- `calculate_async`: Perform mathematical calculations asynchronously
+- `log_activity`: Log activities with timestamps
+
+## Testing Scenarios
+
+### 1. Basic Callback Flow
+```
+"What's the weather in New York?"
+```
+Watch the console output to see:
+- Audit logging before the tool call
+- Security check (will pass for New York)
+- Response enhancement after the tool call
+
+### 2. Security Blocking
+```
+"What's the weather in classified?"
+```
+The security callback will block this request and return an error response.
+
+### 3. Validation Prevention
+```
+"Calculate 10 divided by 0"
+```
+The async validation callback will prevent division by zero.
+
+### 4. Multiple Tool Calls
+```
+"Get weather for London and calculate 5 + 3"
+```
+See how callbacks work with multiple parallel tool calls.
+
+### 5. Callback Chain Testing
+```
+"Log this activity: Testing callback chains"
+```
+Observe how multiple callbacks in the chain are processed.
+
+## Getting Started
+
+1. **Start the ADK Web Server**
+   ```bash
+   adk web
+   ```
+
+2. **Access the ADK Web UI**
+   Navigate to `http://localhost:8000`
+
+3. **Select the Agent**
+   Choose "tool_callbacks_agent" from the dropdown in the top-left corner
+
+4. **Start Streaming**
+   Click the **Audio** or **Video** icon to begin streaming
+
+5. **Test Callbacks**
+   Try the testing scenarios above and watch both the chat responses and the console output to see callbacks in action
+
+## What to Observe
+
+- **Console Output**: Watch for callback logs with emojis:
+  - 🔍 AUDIT: Audit callback logging
+  - 🚫 SECURITY: Security callback blocking
+  - ⚡ ASYNC BEFORE: Async preprocessing
+  - ✨ ENHANCE: Response enhancement
+  - 🔄 ASYNC AFTER: Async post-processing
+
+- **Enhanced Responses**: Tool responses will include additional metadata added by after callbacks
+
+- **Error Handling**: Security blocks and validation errors will be returned as proper error responses
+
+## Technical Notes
+
+- This sample demonstrates that tool callbacks now work identically in both regular and live streaming modes
+- Multiple callbacks are supported and processed in order
+- Both sync and async callbacks are supported
+- Callbacks can modify, enhance, or block tool execution
+- The callback system provides full control over the tool execution pipeline