# 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. # # --- Setup Instructions --- # # 1. Install the ADK package: # !pip install google-adk # # Make sure to restart kernel if using colab/jupyter notebooks # # 2. Set up your Gemini API Key: # # - Get a key from Google AI Studio: https://aistudio.google.com/app/apikey # # - Set it as an environment variable: # import os # os.environ["GOOGLE_API_KEY"] = "YOUR_API_KEY_HERE" # <--- REPLACE with your actual key # # Or learn about other authentication methods (like Agent Platform): # # https://adk.dev/agents/models/ # ADK Imports from google.adk.agents import LlmAgent from google.adk.agents.callback_context import CallbackContext from google.adk.runners import InMemoryRunner # Use InMemoryRunner from google.genai import types # For types.Content from typing import Optional # Define the model - Use the specific model name requested GEMINI_2_FLASH = "gemini-2.0-flash" # --- 1. Define the Callback Function --- def modify_output_after_agent( callback_context: CallbackContext, ) -> Optional[types.Content]: """ Logs exit from an agent and checks 'add_concluding_note' in session state. If True, returns new Content to *replace* the agent's original output. If False or not present, returns None, allowing the agent's original output to be used. """ agent_name = callback_context.agent_name invocation_id = callback_context.invocation_id current_state = callback_context.state.to_dict() print(f"\n[Callback] Exiting agent: {agent_name} (Inv: {invocation_id})") print(f"[Callback] Current State: {current_state}") # Example: Check state to decide whether to modify the final output if current_state.get("add_concluding_note", False): print( f"[Callback] State condition 'add_concluding_note=True' met: Replacing agent {agent_name}'s output." ) # Return Content to *replace* the agent's own output return types.Content( parts=[ types.Part( text=f"Concluding note added by after_agent_callback, replacing original output." ) ], role="model", # Assign model role to the overriding response ) else: print( f"[Callback] State condition not met: Using agent {agent_name}'s original output." ) # Return None - the agent's output produced just before this callback will be used. return None # --- 2. Setup Agent with Callback --- llm_agent_with_after_cb = LlmAgent( name="MySimpleAgentWithAfter", model=GEMINI_2_FLASH, instruction="You are a simple agent. Just say 'Processing complete!'", description="An LLM agent demonstrating after_agent_callback for output modification", after_agent_callback=modify_output_after_agent, # Assign the callback here ) # --- 3. Setup Runner and Sessions using InMemoryRunner --- async def main(): app_name = "after_agent_demo" user_id = "test_user_after" session_id_normal = "session_run_normally" session_id_modify = "session_modify_output" # Use InMemoryRunner - it includes InMemorySessionService runner = InMemoryRunner(agent=llm_agent_with_after_cb, app_name=app_name) # Get the bundled session service to create sessions session_service = runner.session_service # Create session 1: Agent output will be used as is (default empty state) await session_service.create_session( app_name=app_name, user_id=user_id, session_id=session_id_normal, # No initial state means 'add_concluding_note' will be False in the callback check ) # print(f"Session '{session_id_normal}' created with default state.") # Create session 2: Agent output will be replaced by the callback await session_service.create_session( app_name=app_name, user_id=user_id, session_id=session_id_modify, state={"add_concluding_note": True}, # Set the state flag here ) # print(f"Session '{session_id_modify}' created with state={{'add_concluding_note': True}}.") # --- Scenario 1: Run where callback allows agent's original output --- print( "\n" + "=" * 20 + f" SCENARIO 1: Running Agent on Session '{session_id_normal}' (Should Use Original Output) " + "=" * 20 ) async for event in runner.run_async( user_id=user_id, session_id=session_id_normal, new_message=types.Content( role="user", parts=[types.Part(text="Process this please.")] ), ): # Print final output (either from LLM or callback override) if event.is_final_response() and event.content: print( f"Final Output: [{event.author}] {event.content.parts[0].text.strip()}" ) elif event.is_error(): print(f"Error Event: {event.error_details}") # --- Scenario 2: Run where callback replaces the agent's output --- print( "\n" + "=" * 20 + f" SCENARIO 2: Running Agent on Session '{session_id_modify}' (Should Replace Output) " + "=" * 20 ) async for event in runner.run_async( user_id=user_id, session_id=session_id_modify, new_message=types.Content( role="user", parts=[types.Part(text="Process this and add note.")] ), ): # Print final output (either from LLM or callback override) if event.is_final_response() and event.content: print( f"Final Output: [{event.author}] {event.content.parts[0].text.strip()}" ) elif event.is_error(): print(f"Error Event: {event.error_details}") # --- 4. Execute --- # In a Python script: # import asyncio # if __name__ == "__main__": # # Make sure GOOGLE_API_KEY environment variable is set if not using Agent Platform auth # # Or ensure Application Default Credentials (ADC) are configured for Agent Platform # asyncio.run(main()) # In a Jupyter Notebook or similar environment: await main()