From f52950deac8c35c8423a68172d1b456a1bd19540 Mon Sep 17 00:00:00 2001 From: sitalakshmi Date: Wed, 9 Apr 2025 10:00:47 -0400 Subject: [PATCH] refactor(callbacks): update code examples --- docs/callbacks/index.md | 86 +---- docs/callbacks/types-of-callbacks.md | 324 +----------------- .../callbacks/after_agent_callback.py | 59 ++++ .../callbacks/after_model_callback.py | 95 +++++ .../snippets/callbacks/after_tool_callback.py | 92 +++++ .../callbacks/before_agent_callback.py | 56 +++ .../callbacks/before_model_callback.py | 89 +++++ .../callbacks/before_tool_callback.py | 79 +++++ .../snippets/callbacks/callback_basic.py | 49 +++ 9 files changed, 528 insertions(+), 401 deletions(-) create mode 100644 examples/python/snippets/callbacks/after_agent_callback.py create mode 100644 examples/python/snippets/callbacks/after_model_callback.py create mode 100644 examples/python/snippets/callbacks/after_tool_callback.py create mode 100644 examples/python/snippets/callbacks/before_agent_callback.py create mode 100644 examples/python/snippets/callbacks/before_model_callback.py create mode 100644 examples/python/snippets/callbacks/before_tool_callback.py create mode 100644 examples/python/snippets/callbacks/callback_basic.py diff --git a/docs/callbacks/index.md b/docs/callbacks/index.md index 7426bee59e..4d4e3f3a35 100644 --- a/docs/callbacks/index.md +++ b/docs/callbacks/index.md @@ -23,28 +23,7 @@ Callbacks are a cornerstone feature of ADK, providing a powerful mechanism to ho **How are they added?** You register callbacks by passing your defined Python functions as arguments to the agent's constructor (`__init__`) when you create an instance of `Agent` or `LlmAgent`. ```py -# agents/llm_agent.py (Illustrative Snippet) -from google.adk.agents import LlmAgent, LlmRequest, LlmResponse, CallbackContext -from typing import Optional - -# --- Define your callback function --- -def my_before_model_logic( - callback_context: CallbackContext, llm_request: LlmRequest -) -> Optional[LlmResponse]: - print(f"Callback running before model call for agent: {callback_context.agent_name}") - # ... your custom logic here ... - return None # Allow the model call to proceed - -# --- Register it during Agent creation --- -my_agent = LlmAgent( - name="MyCallbackAgent", - model="gemini-2.0-flash-exp", # Or your desired model - instruction="Be helpful.", - # Other agent parameters... - before_model_callback=my_before_model_logic # Pass the function here -) - -print(f"Agent '{my_agent.name}' created with a before_model_callback.") +--8<-- "examples/python/snippets/callbacks/callback_basic.py:callback_basic" ``` ## The Callback Mechanism: Interception and Control @@ -76,68 +55,7 @@ When the ADK framework encounters a point where a callback can run (e.g., just b This example demonstrates the common pattern for a guardrail using `before_model_callback`. ```py -# agents/llm_agent.py (Illustrative Snippet) -from google.adk.agents import CallbackContext, LlmRequest, LlmResponse -from google.genai import types # For types.Content -from typing import Optional - -# Example for before_model_callback -def block_forbidden_input( - callback_context: CallbackContext, llm_request: LlmRequest -) -> Optional[LlmResponse]: - """ - Checks the last user message for a forbidden phrase. - If found, returns a predefined LlmResponse to block the LLM call. - Otherwise, returns None to allow the call. - """ - agent_name = callback_context.agent_name - last_user_message_text = "" - - # Safely get the last user message text - if llm_request.contents: - # Find the last content object with role 'user' - user_contents = [c for c in llm_request.contents if c.role == 'user'] - if user_contents: - last_content = user_contents[-1] - if last_content.parts and last_content.parts[0].text: - last_user_message_text = last_content.parts[0].text - - print(f"[Callback - {agent_name}] Checking input: '{last_user_message_text[:50]}...'") - - # --- Guardrail Logic --- - if "highly restricted topic" in last_user_message_text.lower(): - print(f"[Callback - {agent_name}] Forbidden topic detected! Blocking LLM call.") - - # Create the response object to return *instead* of calling the LLM - blocked_response = LlmResponse( - content=types.Content( - role="model", # Mimic a model response role - parts=[types.Part(text="I am unable to discuss that topic.")] - ) - # Optionally include error codes/messages if relevant - # error_code="POLICY_VIOLATION", - # error_message="Request blocked due to content policy." - ) - # Update state to track violations (optional) - callback_context.state['policy_violations'] = callback_context.state.get('policy_violations', 0) + 1 - - return blocked_response # <-- OVERRIDE: Skip LLM call, use this response. - - # --- No blocking condition met --- - else: - print(f"[Callback - {agent_name}] Input OK. Allowing LLM call.") - # Optionally modify llm_request here if needed before proceeding - # e.g., llm_request.config.temperature = 0.2 - - return None # <-- ALLOW: Proceed with the default behavior (call LLM). - -# Agent definition using the callback -guardrail_agent = LlmAgent( - name="GuardrailAgent", - model="gemini-2.0-flash-exp", - instruction="Answer user questions.", - before_model_callback=block_forbidden_input -) +--8<-- "examples/python/snippets/callbacks/before_model_callback.py" ``` By understanding this mechanism of returning `None` versus returning specific objects, you can precisely control the agent's execution path, making callbacks an essential tool for building sophisticated and reliable agents with ADK. diff --git a/docs/callbacks/types-of-callbacks.md b/docs/callbacks/types-of-callbacks.md index 44069a02e5..acb7ba975f 100644 --- a/docs/callbacks/types-of-callbacks.md +++ b/docs/callbacks/types-of-callbacks.md @@ -13,34 +13,7 @@ These callbacks are available on *any* agent that inherits from `BaseAgent` (inc **Purpose:** Ideal for setting up resources or state needed only for this specific agent's run, performing validation checks on the session state (callback\_context.state) before execution starts, logging the entry point of the agent's activity, or potentially modifying the invocation context before the core logic uses it. ```py - -from google.adk.agents.callback_context import CallbackContext - -# --- Define the Callback Function --- -def simple_before_agent_logger(callback_context: CallbackContext) -> Optional[types.Content]: - """Logs entry into an agent and checks a condition.""" - agent_name = callback_context.agent_name - invocation_id = callback_context.invocation_id - print(f"[Callback] Entering agent: {agent_name} (Invocation: {invocation_id})") - - # Example: Check a condition in state - if callback_context.state.get("skip_agent", False): - print(f"[Callback] Condition met: Skipping agent {agent_name}.") - # Return Content to skip the agent's run - return types.Content(parts=[types.Part(text=f"Agent {agent_name} was skipped by callback.")]) - else: - print(f"[Callback] Condition not met: Proceeding with agent {agent_name}.") - # Return None to allow the agent's run to execute - return None - -# Create LlmAgent and Assign Callback -my_llm_agent = LlmAgent( - name="SimpleLlmAgent", - model=GEMINI_2_FLASH, - instruction="You are a simple agent. Just say 'Hello!'", - description="An LLM agent demonstrating before_agent_callback", - before_agent_callback=simple_before_agent_logger - ) +--8<-- "examples/python/snippets/callbacks/before_agent_callback.py" ``` ### After Agent Callback @@ -50,38 +23,7 @@ my_llm_agent = LlmAgent( **Purpose:** Useful for cleanup tasks, post-execution validation, logging the completion of an agent's activity, modifying final state, or augmenting/replacing the agent's final output. ```py - -from google.adk.agents.callback_context import CallbackContext - -# --- Define the Callback Function --- -def simple_after_agent_logger(callback_context: CallbackContext) -> Optional[types.Content]: - """Logs exit from an agent and optionally appends a message.""" - agent_name = callback_context.agent_name - invocation_id = callback_context.invocation_id - print(f"[Callback] Exiting agent: {agent_name} (Invocation: {invocation_id})") - - # Example: Check state potentially modified during the agent's run - final_status = callback_context.state.get("agent_run_status", "Completed Normally") - print(f"[Callback] Agent run status from state: {final_status}") - - # Example: Optionally return Content to append a message - if callback_context.state.get("add_concluding_note", False): - print(f"[Callback] Adding concluding note for agent {agent_name}.") - # Return Content to append after the agent's own output - return types.Content(parts=[types.Part(text=f"Concluding note added by after_agent_callback.")]) - else: - print(f"[Callback] No concluding note added for agent {agent_name}.") - # Return None - no additional message appended - return None - -# Create LlmAgent and Assign Callback -my_llm_agent = LlmAgent( - name="SimpleLlmAgentWithAfter", - model=GEMINI_2_FLASH, - instruction="You are a simple agent. Just say 'Processing complete!'", - description="An LLM agent demonstrating after_agent_callback", - after_agent_callback=simple_after_agent_logger # Assign the function here - ) +--8<-- "examples/python/snippets/callbacks/after_agent_callback.py" ``` ## LLM Interaction Callbacks @@ -98,68 +40,7 @@ These callbacks are specific to `LlmAgent` and provide hooks around the interact If the callback returns `None`, the LLM continues its normal workflow. If the callback returns an `LlmResponse` object, then the call to the LLM is **skipped**. The returned `LlmResponse` is used directly as if it came from the model. This is powerful for implementing guardrails or caching. ```py -from google.adk.agents.callback_context import CallbackContext -from google.adk.models.llm_request import LlmRequest -from google.adk.models.llm_response import LlmResponse - - -# --- Define the Callback Function --- -def simple_before_model_modifier( - callback_context: CallbackContext, llm_request: LlmRequest -) -> Optional[LlmResponse]: - """Inspects/modifies the LLM request or skips the call.""" - agent_name = callback_context.agent_name - print(f"[Callback] Before model call for agent: {agent_name}") - - # Inspect the last user message in the request contents - last_user_message = "" - if llm_request.contents and llm_request.contents[-1].role == 'user': - if llm_request.contents[-1].parts: - last_user_message = llm_request.contents[-1].parts[0].text - print(f"[Callback] Inspecting last user message: '{last_user_message}'") - - # --- Modification Example --- - # Add a prefix to the system instruction - original_instruction = llm_request.config.system_instruction or types.Content(role="system", parts=[]) - prefix = "[Modified by Callback] " - # Ensure system_instruction is Content and parts list exists - if not isinstance(original_instruction, types.Content): - # Handle case where it might be a string (though config expects Content) - original_instruction = types.Content(role="system", parts=[types.Part(text=str(original_instruction))]) - if not original_instruction.parts: - original_instruction.parts.append(types.Part(text="")) # Add an empty part if none exist - - # Modify the text of the first part - modified_text = prefix + (original_instruction.parts[0].text or "") - original_instruction.parts[0].text = modified_text - llm_request.config.system_instruction = original_instruction - print(f"[Callback] Modified system instruction to: '{modified_text}'") - - # --- Skip Example --- - # Check if the last user message contains "BLOCK" - if "BLOCK" in last_user_message.upper(): - print("[Callback] 'BLOCK' keyword found. Skipping LLM call.") - # Return an LlmResponse to skip the actual LLM call - return LlmResponse( - content=types.Content( - role="model", - parts=[types.Part(text="LLM call was blocked by before_model_callback.")], - ) - ) - else: - print("[Callback] Proceeding with LLM call.") - # Return None to allow the (modified) request to go to the LLM - return None - - -# Create LlmAgent and Assign Callback -my_llm_agent = LlmAgent( - name="ModelCallbackAgent", - model=GEMINI_2_FLASH, - instruction="You are a helpful assistant.", # Base instruction - description="An LLM agent demonstrating before_model_callback", - before_model_callback=simple_before_model_modifier # Assign the function here - ) +--8<-- "examples/python/snippets/callbacks/before_model_callback.py" ``` ### After Model Callback @@ -175,75 +56,7 @@ my_llm_agent = LlmAgent( * or handling specific error codes. ```py - -from google.adk.agents.callback_context import CallbackContext -from google.adk.models.llm_request import LlmRequest -from google.adk.models.llm_response import LlmResponse - - -# --- Define the Callback Function --- -def simple_after_model_modifier( - callback_context: CallbackContext, llm_response: LlmResponse -) -> Optional[LlmResponse]: - """Inspects/modifies the LLM response after it's received.""" - agent_name = callback_context.agent_name - print(f"[Callback] After model call for agent: {agent_name}") - - # --- Inspection --- - original_text = "" - if llm_response.content and llm_response.content.parts: - # Assuming simple text response for this example - if llm_response.content.parts[0].text: - original_text = llm_response.content.parts[0].text - print(f"[Callback] Inspected original response text: '{original_text[:100]}...'") # Log snippet - elif llm_response.content.parts[0].function_call: - print(f"[Callback] Inspected response: Contains function call '{llm_response.content.parts[0].function_call.name}'. No text modification.") - return None # Don't modify tool calls in this example - else: - print("[Callback] Inspected response: No text content found.") - return None - elif llm_response.error_message: - print(f"[Callback] Inspected response: Contains error '{llm_response.error_message}'. No modification.") - return None - else: - print("[Callback] Inspected response: Empty LlmResponse.") - return None # Nothing to modify - - # --- Modification Example --- - # Replace "joke" with "funny story" (case-insensitive) - search_term = "joke" - replace_term = "funny story" - if search_term in original_text.lower(): - print(f"[Callback] Found '{search_term}'. Modifying response.") - modified_text = original_text.replace(search_term, replace_term) - modified_text = modified_text.replace(search_term.capitalize(), replace_term.capitalize()) # Handle capitalization - - # Create a NEW LlmResponse with the modified content - # Deep copy parts to avoid modifying original if other callbacks exist - modified_parts = [copy.deepcopy(part) for part in llm_response.content.parts] - modified_parts[0].text = modified_text # Update the text in the copied part - - new_response = LlmResponse( - content=types.Content(role="model", parts=modified_parts), - # Copy other relevant fields if necessary, e.g., grounding_metadata - grounding_metadata=llm_response.grounding_metadata - ) - print(f"[Callback] Returning modified response.") - return new_response # Return the modified response - else: - print(f"[Callback] '{search_term}' not found. Passing original response through.") - # Return None to use the original llm_response - return None - - -# Create LlmAgent and Assign Callback -my_llm_agent = LlmAgent( - name="AfterModelCallbackAgent", - model=GEMINI_2_FLASH, - instruction="You are a helpful assistant.", - description="An LLM agent demonstrating after_model_callback", - after_model_callback=simple_after_model_modifier # Assign the function here - ) +--8<-- "examples/python/snippets/callbacks/after_model_callback.py" ``` ## Tool Execution Callbacks @@ -262,67 +75,7 @@ These callbacks are also specific to `LlmAgent` and trigger around the execution 2. If a dictionary is returned, the tool's `run_async` method is **skipped**. The returned dictionary is used directly as the result of the tool call. This is useful for caching or overriding tool behavior. ```py - -from google.adk.agents.callback_context import CallbackContext -from google.adk.tools.base_tool import BaseTool -from google.adk.tools.tool_context import ToolContext -from google.adk.tools.function_tool import FunctionTool - -# --- Define a Simple Tool Function --- -def get_capital_city(country: str) -> str: - """Retrieves the capital city of a given country.""" - print(f"--- Tool 'get_capital_city' executing with country: {country} ---") - country_capitals = { - "united states": "Washington, D.C.", - "canada": "Ottawa", # Intentionally correct here - "france": "Paris", - "germany": "Berlin", - } - return country_capitals.get(country.lower(), f"Capital not found for {country}") - -# --- Wrap the function into a Tool --- -capital_tool = FunctionTool(func=get_capital_city) - -# --- Define the Callback Function --- -def simple_before_tool_modifier( - tool: BaseTool, args: Dict[str, Any], tool_context: ToolContext -) -> Optional[Dict]: - """Inspects/modifies tool args or skips the tool call.""" - agent_name = tool_context.agent_name - tool_name = tool.name - print(f"[Callback] Before tool call for tool '{tool_name}' in agent '{agent_name}'") - print(f"[Callback] Original args: {args}") - - # --- Modification Example --- - # If the tool is 'get_capital_city' and country is 'Canada', change it to 'France' - if tool_name == 'get_capital_city' and args.get('country', '').lower() == 'canada': - print("[Callback] Detected 'Canada'. Modifying args to 'France'.") - args['country'] = 'France' # Modify the args dictionary directly - print(f"[Callback] Modified args: {args}") - return None # Proceed with modified args - - # --- Skip Example --- - # If the tool is 'get_capital_city' and country is 'BLOCK' - if tool_name == 'get_capital_city' and args.get('country', '').upper() == 'BLOCK': - print("[Callback] Detected 'BLOCK'. Skipping tool execution.") - # Return a dictionary to be used as the tool result, skipping the actual tool call - return {"result": "Tool execution was blocked by before_tool_callback."} - - print("[Callback] Proceeding with original or previously modified args.") - # Return None to allow the tool to execute normally (with original or modified args) - return None - - - -# Create LlmAgent and Assign Callback -my_llm_agent = LlmAgent( - name="ToolCallbackAgent", - model=GEMINI_2_FLASH, - instruction="You are an agent that can find capital cities. Use the get_capital_city tool.", - description="An LLM agent demonstrating before_tool_callback", - tools=[capital_tool], # Add the tool here - before_tool_callback=simple_before_tool_modifier # Assign the callback here - ) +--8<-- "examples/python/snippets/callbacks/before_tool_callback.py" ``` ### After Tool Callback @@ -337,68 +90,5 @@ my_llm_agent = LlmAgent( 2. If a new dictionary is returned, it **replaces** the original `tool_response`. This allows modifying or filtering the result seen by the LLM. ```py - -from google.adk.agents.callback_context import CallbackContext -from google.adk.tools.base_tool import BaseTool -from google.adk.tools.tool_context import ToolContext -from google.adk.tools.function_tool import FunctionTool - - -# --- Define a Simple Tool Function (Same as before) --- -def get_capital_city(country: str) -> str: - """Retrieves the capital city of a given country.""" - print(f"--- Tool 'get_capital_city' executing with country: {country} ---") - country_capitals = { - "united states": "Washington, D.C.", - "canada": "Ottawa", - "france": "Paris", - "germany": "Berlin", - } - return {"result": country_capitals.get(country.lower(), f"Capital not found for {country}")} - -# --- Wrap the function into a Tool --- -capital_tool = FunctionTool(func=get_capital_city) - -# --- Define the Callback Function --- -def simple_after_tool_modifier( - tool: BaseTool, args: Dict[str, Any], tool_context: ToolContext, tool_response: Dict -) -> Optional[Dict]: - """Inspects/modifies the tool result after execution.""" - agent_name = tool_context.agent_name - tool_name = tool.name - print(f"[Callback] After tool call for tool '{tool_name}' in agent '{agent_name}'") - print(f"[Callback] Args used: {args}") - print(f"[Callback] Original tool_response: {tool_response}") - - # Default structure for function tool results is {"result": } - original_result_value = tool_response.get("result", "") - # original_result_value = tool_response - - # --- Modification Example --- - # If the tool was 'get_capital_city' and result is 'Washington, D.C.' - if tool_name == 'get_capital_city' and original_result_value == "Washington, D.C.": - print("[Callback] Detected 'Washington, D.C.'. Modifying tool response.") - - # IMPORTANT: Create a new dictionary or modify a copy - modified_response = copy.deepcopy(tool_response) - modified_response["result"] = f"{original_result_value} (Note: This is the capital of the USA)." - modified_response["note_added_by_callback"] = True # Add extra info if needed - - print(f"[Callback] Modified tool_response: {modified_response}") - return modified_response # Return the modified dictionary - - print("[Callback] Passing original tool response through.") - # Return None to use the original tool_response - return None - - -# Create LlmAgent and Assign Callback -my_llm_agent = LlmAgent( - name="AfterToolCallbackAgent", - model=GEMINI_2_FLASH, - instruction="You are an agent that finds capital cities using the get_capital_city tool. Report the result clearly.", - description="An LLM agent demonstrating after_tool_callback", - tools=[capital_tool], # Add the tool - after_tool_callback=simple_after_tool_modifier # Assign the callback - ) -``` +--8<-- "examples/python/snippets/callbacks/after_tool_callback.py" +``` \ No newline at end of file diff --git a/examples/python/snippets/callbacks/after_agent_callback.py b/examples/python/snippets/callbacks/after_agent_callback.py new file mode 100644 index 0000000000..943978652e --- /dev/null +++ b/examples/python/snippets/callbacks/after_agent_callback.py @@ -0,0 +1,59 @@ +from google.adk.agents import LlmAgent +from google.adk.agents.callback_context import CallbackContext +from google.adk.runners import Runner +from typing import Optional +from google.genai import types +from google.adk.sessions import InMemorySessionService + +GEMINI_2_FLASH="gemini-2.0-flash-exp" + +# --- Define the Callback Function --- +def simple_after_agent_logger(callback_context: CallbackContext) -> Optional[types.Content]: + """Logs exit from an agent and optionally appends a message.""" + agent_name = callback_context.agent_name + invocation_id = callback_context.invocation_id + print(f"[Callback] Exiting agent: {agent_name} (Invocation: {invocation_id})") + + # Example: Check state potentially modified during the agent's run + final_status = callback_context.state.get("agent_run_status", "Completed Normally") + print(f"[Callback] Agent run status from state: {final_status}") + + # Example: Optionally return Content to append a message + if callback_context.state.get("add_concluding_note", False): + print(f"[Callback] Adding concluding note for agent {agent_name}.") + # Return Content to append after the agent's own output + return types.Content(parts=[types.Part(text=f"Concluding note added by after_agent_callback.")]) + else: + print(f"[Callback] No concluding note added for agent {agent_name}.") + # Return None - no additional message appended + return None + +my_llm_agent = LlmAgent( + name="SimpleLlmAgentWithAfter", + model=GEMINI_2_FLASH, + instruction="You are a simple agent. Just say 'Processing complete!'", + description="An LLM agent demonstrating after_agent_callback", + after_agent_callback=simple_after_agent_logger # Assign the function here + ) + +APP_NAME = "guardrail_app" +USER_ID = "user_1" +SESSION_ID = "session_001" + +# Session and Runner +session_service = InMemorySessionService() +session = session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID) +runner = Runner(agent=my_llm_agent, app_name=APP_NAME, session_service=session_service) + + +# Agent Interaction +def call_agent(query): + content = types.Content(role='user', parts=[types.Part(text=query)]) + events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content) + + for event in events: + if event.is_final_response(): + final_response = event.content.parts[0].text + print("Agent Response: ", final_response) + +call_agent("callback example") \ No newline at end of file diff --git a/examples/python/snippets/callbacks/after_model_callback.py b/examples/python/snippets/callbacks/after_model_callback.py new file mode 100644 index 0000000000..693aa0e685 --- /dev/null +++ b/examples/python/snippets/callbacks/after_model_callback.py @@ -0,0 +1,95 @@ +from google.adk.agents import LlmAgent +from google.adk.agents.callback_context import CallbackContext +from google.adk.runners import Runner +from typing import Optional +from google.genai import types +from google.adk.sessions import InMemorySessionService +from google.adk.models import LlmResponse + +GEMINI_2_FLASH="gemini-2.0-flash-exp" + +# --- Define the Callback Function --- +def simple_after_model_modifier( + callback_context: CallbackContext, llm_response: LlmResponse +) -> Optional[LlmResponse]: + """Inspects/modifies the LLM response after it's received.""" + agent_name = callback_context.agent_name + print(f"[Callback] After model call for agent: {agent_name}") + + # --- Inspection --- + original_text = "" + if llm_response.content and llm_response.content.parts: + # Assuming simple text response for this example + if llm_response.content.parts[0].text: + original_text = llm_response.content.parts[0].text + print(f"[Callback] Inspected original response text: '{original_text[:100]}...'") # Log snippet + elif llm_response.content.parts[0].function_call: + print(f"[Callback] Inspected response: Contains function call '{llm_response.content.parts[0].function_call.name}'. No text modification.") + return None # Don't modify tool calls in this example + else: + print("[Callback] Inspected response: No text content found.") + return None + elif llm_response.error_message: + print(f"[Callback] Inspected response: Contains error '{llm_response.error_message}'. No modification.") + return None + else: + print("[Callback] Inspected response: Empty LlmResponse.") + return None # Nothing to modify + + # --- Modification Example --- + # Replace "joke" with "funny story" (case-insensitive) + search_term = "joke" + replace_term = "funny story" + if search_term in original_text.lower(): + print(f"[Callback] Found '{search_term}'. Modifying response.") + modified_text = original_text.replace(search_term, replace_term) + modified_text = modified_text.replace(search_term.capitalize(), replace_term.capitalize()) # Handle capitalization + + # Create a NEW LlmResponse with the modified content + # Deep copy parts to avoid modifying original if other callbacks exist + modified_parts = [copy.deepcopy(part) for part in llm_response.content.parts] + modified_parts[0].text = modified_text # Update the text in the copied part + + new_response = LlmResponse( + content=types.Content(role="model", parts=modified_parts), + # Copy other relevant fields if necessary, e.g., grounding_metadata + grounding_metadata=llm_response.grounding_metadata + ) + print(f"[Callback] Returning modified response.") + return new_response # Return the modified response + else: + print(f"[Callback] '{search_term}' not found. Passing original response through.") + # Return None to use the original llm_response + return None + + +# Create LlmAgent and Assign Callback +my_llm_agent = LlmAgent( + name="AfterModelCallbackAgent", + model=GEMINI_2_FLASH, + instruction="You are a helpful assistant.", + description="An LLM agent demonstrating after_model_callback", + after_model_callback=simple_after_model_modifier # Assign the function here +) + +APP_NAME = "guardrail_app" +USER_ID = "user_1" +SESSION_ID = "session_001" + +# Session and Runner +session_service = InMemorySessionService() +session = session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID) +runner = Runner(agent=my_llm_agent, app_name=APP_NAME, session_service=session_service) + + +# Agent Interaction +def call_agent(query): + content = types.Content(role='user', parts=[types.Part(text=query)]) + events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content) + + for event in events: + if event.is_final_response(): + final_response = event.content.parts[0].text + print("Agent Response: ", final_response) + +call_agent("callback example") \ No newline at end of file diff --git a/examples/python/snippets/callbacks/after_tool_callback.py b/examples/python/snippets/callbacks/after_tool_callback.py new file mode 100644 index 0000000000..4e185bb542 --- /dev/null +++ b/examples/python/snippets/callbacks/after_tool_callback.py @@ -0,0 +1,92 @@ +from google.adk.agents import LlmAgent +from google.adk.runners import Runner +from typing import Optional +from google.genai import types +from google.adk.sessions import InMemorySessionService +from google.adk.tools import FunctionTool +from google.adk.tools.tool_context import ToolContext +from google.adk.tools.base_tool import BaseTool +from typing import Dict, Any +from copy import copy + +GEMINI_2_FLASH="gemini-2.0-flash-exp" + +# --- Define a Simple Tool Function (Same as before) --- +def get_capital_city(country: str) -> str: + """Retrieves the capital city of a given country.""" + print(f"--- Tool 'get_capital_city' executing with country: {country} ---") + country_capitals = { + "united states": "Washington, D.C.", + "canada": "Ottawa", + "france": "Paris", + "germany": "Berlin", + } + return {"result": country_capitals.get(country.lower(), f"Capital not found for {country}")} + +# --- Wrap the function into a Tool --- +capital_tool = FunctionTool(func=get_capital_city) + +# --- Define the Callback Function --- +def simple_after_tool_modifier( + tool: BaseTool, args: Dict[str, Any], tool_context: ToolContext, tool_response: Dict +) -> Optional[Dict]: + """Inspects/modifies the tool result after execution.""" + agent_name = tool_context.agent_name + tool_name = tool.name + print(f"[Callback] After tool call for tool '{tool_name}' in agent '{agent_name}'") + print(f"[Callback] Args used: {args}") + print(f"[Callback] Original tool_response: {tool_response}") + + # Default structure for function tool results is {"result": } + original_result_value = tool_response.get("result", "") + # original_result_value = tool_response + + # --- Modification Example --- + # If the tool was 'get_capital_city' and result is 'Washington, D.C.' + if tool_name == 'get_capital_city' and original_result_value == "Washington, D.C.": + print("[Callback] Detected 'Washington, D.C.'. Modifying tool response.") + + # IMPORTANT: Create a new dictionary or modify a copy + modified_response = copy.deepcopy(tool_response) + modified_response["result"] = f"{original_result_value} (Note: This is the capital of the USA)." + modified_response["note_added_by_callback"] = True # Add extra info if needed + + print(f"[Callback] Modified tool_response: {modified_response}") + return modified_response # Return the modified dictionary + + print("[Callback] Passing original tool response through.") + # Return None to use the original tool_response + return None + + +# Create LlmAgent and Assign Callback +my_llm_agent = LlmAgent( + name="AfterToolCallbackAgent", + model=GEMINI_2_FLASH, + instruction="You are an agent that finds capital cities using the get_capital_city tool. Report the result clearly.", + description="An LLM agent demonstrating after_tool_callback", + tools=[capital_tool], # Add the tool + after_tool_callback=simple_after_tool_modifier # Assign the callback + ) + +APP_NAME = "guardrail_app" +USER_ID = "user_1" +SESSION_ID = "session_001" + +# Session and Runner +session_service = InMemorySessionService() +session = session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID) +runner = Runner(agent=my_llm_agent, app_name=APP_NAME, session_service=session_service) + + +# Agent Interaction +def call_agent(query): + content = types.Content(role='user', parts=[types.Part(text=query)]) + events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content) + + for event in events: + if event.is_final_response(): + final_response = event.content.parts[0].text + print("Agent Response: ", final_response) + +call_agent("callback example") \ No newline at end of file diff --git a/examples/python/snippets/callbacks/before_agent_callback.py b/examples/python/snippets/callbacks/before_agent_callback.py new file mode 100644 index 0000000000..a7e47570ce --- /dev/null +++ b/examples/python/snippets/callbacks/before_agent_callback.py @@ -0,0 +1,56 @@ +from google.adk.agents import LlmAgent +from google.adk.agents.callback_context import CallbackContext +from google.adk.runners import Runner +from typing import Optional +from google.genai import types +from google.adk.sessions import InMemorySessionService + +GEMINI_2_FLASH="gemini-2.0-flash-exp" + +# --- Define the Callback Function --- +def simple_before_agent_logger(callback_context: CallbackContext) -> Optional[types.Content]: + """Logs entry into an agent and checks a condition.""" + agent_name = callback_context.agent_name + invocation_id = callback_context.invocation_id + print(f"[Callback] Entering agent: {agent_name} (Invocation: {invocation_id})") + + # Example: Check a condition in state + if callback_context.state.get("skip_agent", False): + print(f"[Callback] Condition met: Skipping agent {agent_name}.") + # Return Content to skip the agent's run + return types.Content(parts=[types.Part(text=f"Agent {agent_name} was skipped by callback.")]) + else: + print(f"[Callback] Condition not met: Proceeding with agent {agent_name}.") + # Return None to allow the agent's run to execute + return None + +# Create LlmAgent and Assign Callback +my_llm_agent = LlmAgent( + name="SimpleLlmAgent", + model=GEMINI_2_FLASH, + instruction="You are a simple agent. Just say 'Hello!'", + description="An LLM agent demonstrating before_agent_callback", + before_agent_callback=simple_before_agent_logger + ) + +APP_NAME = "guardrail_app" +USER_ID = "user_1" +SESSION_ID = "session_001" + +# Session and Runner +session_service = InMemorySessionService() +session = session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID) +runner = Runner(agent=my_llm_agent, app_name=APP_NAME, session_service=session_service) + + +# Agent Interaction +def call_agent(query): + content = types.Content(role='user', parts=[types.Part(text=query)]) + events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content) + + for event in events: + if event.is_final_response(): + final_response = event.content.parts[0].text + print("Agent Response: ", final_response) + +call_agent("callback example") \ No newline at end of file diff --git a/examples/python/snippets/callbacks/before_model_callback.py b/examples/python/snippets/callbacks/before_model_callback.py new file mode 100644 index 0000000000..c033329792 --- /dev/null +++ b/examples/python/snippets/callbacks/before_model_callback.py @@ -0,0 +1,89 @@ +from google.adk.agents import LlmAgent +from google.adk.agents.callback_context import CallbackContext +from google.adk.models import LlmResponse, LlmRequest +from google.adk.runners import Runner +from typing import Optional +from google.genai import types +from google.adk.sessions import InMemorySessionService + +GEMINI_2_FLASH="gemini-2.0-flash-exp" + +# --- Define the Callback Function --- +def simple_before_model_modifier( + callback_context: CallbackContext, llm_request: LlmRequest +) -> Optional[LlmResponse]: + """Inspects/modifies the LLM request or skips the call.""" + agent_name = callback_context.agent_name + print(f"[Callback] Before model call for agent: {agent_name}") + + # Inspect the last user message in the request contents + last_user_message = "" + if llm_request.contents and llm_request.contents[-1].role == 'user': + if llm_request.contents[-1].parts: + last_user_message = llm_request.contents[-1].parts[0].text + print(f"[Callback] Inspecting last user message: '{last_user_message}'") + + # --- Modification Example --- + # Add a prefix to the system instruction + original_instruction = llm_request.config.system_instruction or types.Content(role="system", parts=[]) + prefix = "[Modified by Callback] " + # Ensure system_instruction is Content and parts list exists + if not isinstance(original_instruction, types.Content): + # Handle case where it might be a string (though config expects Content) + original_instruction = types.Content(role="system", parts=[types.Part(text=str(original_instruction))]) + if not original_instruction.parts: + original_instruction.parts.append(types.Part(text="")) # Add an empty part if none exist + + # Modify the text of the first part + modified_text = prefix + (original_instruction.parts[0].text or "") + original_instruction.parts[0].text = modified_text + llm_request.config.system_instruction = original_instruction + print(f"[Callback] Modified system instruction to: '{modified_text}'") + + # --- Skip Example --- + # Check if the last user message contains "BLOCK" + if "BLOCK" in last_user_message.upper(): + print("[Callback] 'BLOCK' keyword found. Skipping LLM call.") + # Return an LlmResponse to skip the actual LLM call + return LlmResponse( + content=types.Content( + role="model", + parts=[types.Part(text="LLM call was blocked by before_model_callback.")], + ) + ) + else: + print("[Callback] Proceeding with LLM call.") + # Return None to allow the (modified) request to go to the LLM + return None + + +# Create LlmAgent and Assign Callback +my_llm_agent = LlmAgent( + name="ModelCallbackAgent", + model=GEMINI_2_FLASH, + instruction="You are a helpful assistant.", # Base instruction + description="An LLM agent demonstrating before_model_callback", + before_model_callback=simple_before_model_modifier # Assign the function here +) + +APP_NAME = "guardrail_app" +USER_ID = "user_1" +SESSION_ID = "session_001" + +# Session and Runner +session_service = InMemorySessionService() +session = session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID) +runner = Runner(agent=my_llm_agent, app_name=APP_NAME, session_service=session_service) + + +# Agent Interaction +def call_agent(query): + content = types.Content(role='user', parts=[types.Part(text=query)]) + events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content) + + for event in events: + if event.is_final_response(): + final_response = event.content.parts[0].text + print("Agent Response: ", final_response) + +call_agent("callback example") \ No newline at end of file diff --git a/examples/python/snippets/callbacks/before_tool_callback.py b/examples/python/snippets/callbacks/before_tool_callback.py new file mode 100644 index 0000000000..0be21f8fd8 --- /dev/null +++ b/examples/python/snippets/callbacks/before_tool_callback.py @@ -0,0 +1,79 @@ +from google.adk.agents import LlmAgent +from google.adk.runners import Runner +from typing import Optional +from google.genai import types +from google.adk.sessions import InMemorySessionService +from google.adk.tools import FunctionTool +from google.adk.tools.tool_context import ToolContext +from google.adk.tools.base_tool import BaseTool +from typing import Dict, Any + + +GEMINI_2_FLASH="gemini-2.0-flash-exp" + +def get_capital_city(country: str) -> str: + """Retrieves the capital city of a given country.""" + print(f"--- Tool 'get_capital_city' executing with country: {country} ---") + country_capitals = { + "united states": "Washington, D.C.", + "canada": "Ottawa", + "france": "Paris", + "germany": "Berlin", + } + return country_capitals.get(country.lower(), f"Capital not found for {country}") + +capital_tool = FunctionTool(func=get_capital_city) + +def simple_before_tool_modifier( + tool: BaseTool, args: Dict[str, Any], tool_context: ToolContext +) -> Optional[Dict]: + """Inspects/modifies tool args or skips the tool call.""" + agent_name = tool_context.agent_name + tool_name = tool.name + print(f"[Callback] Before tool call for tool '{tool_name}' in agent '{agent_name}'") + print(f"[Callback] Original args: {args}") + + if tool_name == 'get_capital_city' and args.get('country', '').lower() == 'canada': + print("[Callback] Detected 'Canada'. Modifying args to 'France'.") + args['country'] = 'France' + print(f"[Callback] Modified args: {args}") + return None + + # If the tool is 'get_capital_city' and country is 'BLOCK' + if tool_name == 'get_capital_city' and args.get('country', '').upper() == 'BLOCK': + print("[Callback] Detected 'BLOCK'. Skipping tool execution.") + return {"result": "Tool execution was blocked by before_tool_callback."} + + print("[Callback] Proceeding with original or previously modified args.") + return None + +my_llm_agent = LlmAgent( + name="ToolCallbackAgent", + model=GEMINI_2_FLASH, + instruction="You are an agent that can find capital cities. Use the get_capital_city tool.", + description="An LLM agent demonstrating before_tool_callback", + tools=[capital_tool], + before_tool_callback=simple_before_tool_modifier +) + +APP_NAME = "guardrail_app" +USER_ID = "user_1" +SESSION_ID = "session_001" + +# Session and Runner +session_service = InMemorySessionService() +session = session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID) +runner = Runner(agent=my_llm_agent, app_name=APP_NAME, session_service=session_service) + + +# Agent Interaction +def call_agent(query): + content = types.Content(role='user', parts=[types.Part(text=query)]) + events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content) + + for event in events: + if event.is_final_response(): + final_response = event.content.parts[0].text + print("Agent Response: ", final_response) + +call_agent("callback example") \ No newline at end of file diff --git a/examples/python/snippets/callbacks/callback_basic.py b/examples/python/snippets/callbacks/callback_basic.py new file mode 100644 index 0000000000..99e4c6686e --- /dev/null +++ b/examples/python/snippets/callbacks/callback_basic.py @@ -0,0 +1,49 @@ +# --8<-- [start:callback_basic] +from google.adk.agents import LlmAgent +from google.adk.agents.callback_context import CallbackContext +from google.adk.models import LlmResponse, LlmRequest +from typing import Optional + +# --- Define your callback function --- +def my_before_model_logic( + callback_context: CallbackContext, llm_request: LlmRequest +) -> Optional[LlmResponse]: + print(f"Callback running before model call for agent: {callback_context.agent_name}") + # ... your custom logic here ... + return None # Allow the model call to proceed + +# --- Register it during Agent creation --- +my_agent = LlmAgent( + name="MyCallbackAgent", + model="gemini-2.0-flash-exp", # Or your desired model + instruction="Be helpful.", + # Other agent parameters... + before_model_callback=my_before_model_logic # Pass the function here +) +# --8<-- [end:callback_basic] + +APP_NAME = "guardrail_app" +USER_ID = "user_1" +SESSION_ID = "session_001" + +from google.adk.runners import Runner +from google.genai import types +from google.adk.sessions import InMemorySessionService + +# Session and Runner +session_service = InMemorySessionService() +session = session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID) +runner = Runner(agent=my_agent, app_name=APP_NAME, session_service=session_service) + + +# Agent Interaction +def call_agent(query): + content = types.Content(role='user', parts=[types.Part(text=query)]) + events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content) + + for event in events: + if event.is_final_response(): + final_response = event.content.parts[0].text + print("Agent Response: ", final_response) + +call_agent("callback example")