adk callbacks

from google.adk.agents.llm_agent import Agent, BaseAgent
from google.genai import types
from google.adk.agents.context import Context
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.agents.callback_context import CallbackContext
from typing import Optional
from google.adk.tools.tool_context import ToolContext
from google.adk.tools.base_tool import BaseTool
from typing import Any



async def before_agent_callback(callback_context: CallbackContext):
    print("Before agent callback triggered")
    print(callback_context)

async def after_agent_callback(callback_context: CallbackContext):
    print("After agent callback triggered")
    print(callback_context)

def before_model_callback(callback_context: Context, llm_request: LlmRequest):
    print("Before model callback triggered")
    print(callback_context)
    print(llm_request)
    
def after_model_callback(callback_context: Context, llm_response: LlmResponse):
    print("After model callback triggered")
    print(callback_context)
    print(llm_response)

def before_tool_callback(tool: BaseTool, args: dict[str, Any], tool_context: ToolContext) -> Optional[dict]:
    print("Before tool callback triggered")
    print(tool.name, args)
    if False:
        return {"result": "Tool execution was blocked by before_tool_callback."}


def after_tool_callback(tool: BaseTool, args: dict[str, Any], tool_context: ToolContext, tool_response: dict) -> Optional[dict]:
    print("After tool callback triggered")
    print(tool.name, args, tool_response)


def get_temperature(city: str):
    """
    A dummy tool to get the temperature of a city.
    """
    temprature = len(city) * 3  # Dummy temperature based on city name length
    return f"The current temperature in {city} is {temprature} degrees Celsius."


root_agent = Agent(
    model='gemini-2.5-flash',
    name='root_agent',
    description='A helpful assistant for user questions.',
    before_agent_callback=before_agent_callback,
    after_agent_callback=after_agent_callback,
    before_model_callback=before_model_callback,
    after_model_callback=after_model_callback,
    before_tool_callback=before_tool_callback,
    after_tool_callback=after_tool_callback,
    instruction='You are a weather assistant. Use get_temperature tool to answer user questions about the weather.',
    tools=[get_temperature]
)

async def main():
    from google.adk.sessions import InMemorySessionService
    from google.adk.runners import Runner

    session_service = InMemorySessionService()

    runner = Runner(
        agent=root_agent,
        app_name="app",
        session_service=session_service, 
    )
    await runner.run_debug("how warm is it in munich?", verbose=True)

if __name__ == "__main__":
    import asyncio
    import dotenv
    dotenv.load_dotenv()  
    asyncio.run(main())