callback-docs-fix (#132)

Author: lavinigam-gcp

docs/assets/callback_flow.png

@@ -4,11 +4,15 @@
 
 Callbacks are a cornerstone feature of ADK, providing a powerful mechanism to hook into an agent's execution process. They allow you to observe, customize, and even control the agent's behavior at specific, predefined points without modifying the core ADK framework code.
 
-**What are they?** In essence, callbacks are standard Python functions that you define. You then associate these functions with an agent when you create it. The ADK framework automatically calls your functions at key stages in the agent's lifecycle, such as:
+**What are they?** In essence, callbacks are standard Python functions that you define. You then associate these functions with an agent when you create it. The ADK framework automatically calls your functions at key stages, letting you observe or intervene. Think of it like checkpoints during the agent's process:
+
+* **Before the agent starts its main work on a request, and after it finishes:** When you ask an agent to do something (e.g., answer a question), it runs its internal logic to figure out the response.
+  * The `before_agent` callback executes *right before* this main work begins for that specific request.
+  * The `after_agent` callback executes *right after* the agent has finished all its steps for that request and has prepared the final result, but just before the result is returned.
+  * This "main work" encompasses the agent's *entire* process for handling that single request. This might involve deciding to call an LLM, actually calling the LLM, deciding to use a tool, using the tool, processing the results, and finally putting together the answer. These callbacks essentially wrap the whole sequence from receiving the input to producing the final output for that one interaction.
+* **Before sending a request to, or after receiving a response from, the Large Language Model (LLM):** These callbacks (`before_model`, `after_model`) allow you to inspect or modify the data going to and coming from the LLM specifically.
+* **Before executing a tool (like a Python function or another agent) or after it finishes:** Similarly, `before_tool` and `after_tool` callbacks give you control points specifically around the execution of tools invoked by the agent.
 
-* Before or after the agent's main processing logic runs.  
-* Before sending a request to, or after receiving a response from, the Large Language Model (LLM).  
-* Before executing a tool (like a Python function or another agent) or after it finishes.
 
 ![intro_components.png](../assets/callback_flow.png)
 

docs/callbacks/index.md

@@ -18,6 +18,18 @@ These callbacks are available on *any* agent that inherits from `BaseAgent` (inc
     --8<-- "examples/python/snippets/callbacks/before_agent_callback.py"
     ```
 
+**Note on the `before_agent_callback` Example:**
+
+* **What it Shows:** This example demonstrates the `before_agent_callback`. This callback runs *right before* the agent's main processing logic starts for a given request.
+* **How it Works:** The callback function (`check_if_agent_should_run`) looks at a flag (`skip_llm_agent`) in the session's state.
+    * If the flag is `True`, the callback returns a `types.Content` object. This tells the ADK framework to **skip** the agent's main execution entirely and use the callback's returned content as the final response.
+    * If the flag is `False` (or not set), the callback returns `None`. This tells the ADK framework to **proceed** with the agent's normal execution (calling the LLM in this case).
+* **Expected Outcome:** You'll see two scenarios:
+    1. In the session *with* the `skip_llm_agent: True` state, the agent's LLM call is bypassed, and the output comes directly from the callback ("Agent... skipped...").
+    2. In the session *without* that state flag, the callback allows the agent to run, and you see the actual response from the LLM (e.g., "Hello!").
+* **Understanding Callbacks:** This highlights how `before_` callbacks act as **gatekeepers**, allowing you to intercept execution *before* a major step and potentially prevent it based on checks (like state, input validation, permissions).
+
+
 ### After Agent Callback
 
 **When:** Called *immediately after* the agent's `_run_async_impl` (or `_run_live_impl`) method successfully completes. It does *not* run if the agent was skipped due to `before_agent_callback` returning content or if `end_invocation` was set during the agent's run.
@@ -30,6 +42,17 @@ These callbacks are available on *any* agent that inherits from `BaseAgent` (inc
     --8<-- "examples/python/snippets/callbacks/after_agent_callback.py"
     ```
 
+**Note on the `after_agent_callback` Example:**
+
+* **What it Shows:** This example demonstrates the `after_agent_callback`. This callback runs *right after* the agent's main processing logic has finished and produced its result, but *before* that result is finalized and returned.
+* **How it Works:** The callback function (`modify_output_after_agent`) checks a flag (`add_concluding_note`) in the session's state.
+    * If the flag is `True`, the callback returns a *new* `types.Content` object. This tells the ADK framework to **replace** the agent's original output with the content returned by the callback.
+    * If the flag is `False` (or not set), the callback returns `None`. This tells the ADK framework to **use** the original output generated by the agent.
+*   **Expected Outcome:** You'll see two scenarios:
+    1. In the session *without* the `add_concluding_note: True` state, the callback allows the agent's original output ("Processing complete!") to be used.
+    2. In the session *with* that state flag, the callback intercepts the agent's original output and replaces it with its own message ("Concluding note added...").
+* **Understanding Callbacks:** This highlights how `after_` callbacks allow **post-processing** or **modification**. You can inspect the result of a step (the agent's run) and decide whether to let it pass through, change it, or completely replace it based on your logic.
+
 ## LLM Interaction Callbacks
 
 These callbacks are specific to `LlmAgent` and provide hooks around the interaction with the Large Language Model.

docs/callbacks/types-of-callbacks.md

@@ -1,59 +1,127 @@
+# # --- 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 Vertex AI):
+# # https://google.github.io/adk-docs/agents/models/
+
+
+# ADK Imports
 from google.adk.agents import LlmAgent
 from google.adk.agents.callback_context import CallbackContext
-from google.adk.runners import Runner
+from google.adk.runners import InMemoryRunner # Use InMemoryRunner
+from google.genai import types # For types.Content
 from typing import Optional
-from google.genai import types 
-from google.adk.sessions import InMemorySessionService
 
+# Define the model - Use the specific model name requested
 GEMINI_2_FLASH="gemini-2.0-flash"
 
-# --- 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."""
+# --- 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
-    print(f"[Callback] Exiting agent: {agent_name} (Invocation: {invocation_id})")
+    current_state = callback_context.state.to_dict()
 
-    # 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}")
+    print(f"\n[Callback] Exiting agent: {agent_name} (Inv: {invocation_id})")
+    print(f"[Callback] Current State: {current_state}")
 
-    # 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.")])
+    # 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] No concluding note added for agent {agent_name}.")
-        # Return None - no additional message appended
+        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
 
-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
+# --- 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)
+    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.")
 
-APP_NAME = "guardrail_app"
-USER_ID = "user_1"
-SESSION_ID = "session_001"
+    # Create session 2: Agent output will be replaced by the callback
+    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}}.")
 
-# 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)
 
+    # --- 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}")
 
-# 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)
+    # --- 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}")
 
-  for event in events:
-      if event.is_final_response():
-          final_response = event.content.parts[0].text
-          print("Agent Response: ", final_response)
+# --- 4. Execute ---
+# In a Python script:
+# import asyncio
+# if __name__ == "__main__":
+#     # Make sure GOOGLE_API_KEY environment variable is set if not using Vertex AI auth
+#     # Or ensure Application Default Credentials (ADC) are configured for Vertex AI
+#     asyncio.run(main())
 
-call_agent("callback example")
+# In a Jupyter Notebook or similar environment:
+await main()