refactor: update callback type signatures to support sync and async responses

src/google/adk/agents/llm_agent.py

@@ -57,11 +57,11 @@ AfterModelCallback: TypeAlias = Callable[
 ]
 BeforeToolCallback: TypeAlias = Callable[
     [BaseTool, dict[str, Any], ToolContext],
-    Awaitable[Optional[dict]],
+    Union[Awaitable[Optional[dict]], Optional[dict]],
 ]
 AfterToolCallback: TypeAlias = Callable[
     [BaseTool, dict[str, Any], ToolContext, dict],
-    Awaitable[Optional[dict]],
+    Union[Awaitable[Optional[dict]], Optional[dict]],
 ]
 
 InstructionProvider: TypeAlias = Callable[[ReadonlyContext], str]

src/google/adk/flows/llm_flows/functions.py

@@ -151,36 +151,33 @@ async def handle_function_calls_async(
     # do not use "args" as the variable name, because it is a reserved keyword
     # in python debugger.
     function_args = function_call.args or {}
-    function_response = None
-    # # Calls the tool if before_tool_callback does not exist or returns None.
-    # if agent.before_tool_callback:
-    #   function_response = agent.before_tool_callback(
-    #       tool=tool, args=function_args, tool_context=tool_context
-    #   )
-    # Short-circuit via before_tool_callback (sync *or* async)
+    function_response: Optional[dict] = None
+
+    # before_tool_callback (sync or async)
     if agent.before_tool_callback:
-      _maybe = agent.before_tool_callback(
+      function_response = agent.before_tool_callback(
           tool=tool, args=function_args, tool_context=tool_context
       )
-      if inspect.isawaitable(_maybe):
-          _maybe = await _maybe
-      function_response = _maybe
+      if inspect.isawaitable(function_response):
+        function_response = await function_response
+
     if not function_response:
       function_response = await __call_tool_async(
           tool, args=function_args, tool_context=tool_context
       )
-    # Calls after_tool_callback if it exists.
+
+    # after_tool_callback (sync or async)
     if agent.after_tool_callback:
-      _maybe2 = agent.after_tool_callback(
+      altered_function_response = agent.after_tool_callback(
           tool=tool,
           args=function_args,
           tool_context=tool_context,
           tool_response=function_response,
       )
-      if inspect.isawaitable(_maybe2):
-        _maybe2 = await _maybe2
-      if _maybe2 is not None:
-        function_response = _maybe2
+      if inspect.isawaitable(altered_function_response):
+        altered_function_response = await altered_function_response
+      if altered_function_response is not None:
+        function_response = altered_function_response
 
     if tool.is_long_running:
       # Allow long running function to return None to not provide function response.

tests/unittests/flows/llm_flows/test_async_tool_callbacks.py

@@ -27,6 +27,7 @@ from ... import utils
 
 
 class AsyncBeforeToolCallback:
+
   def __init__(self, mock_response: Dict[str, Any]):
     self.mock_response = mock_response
 
@@ -40,6 +41,7 @@ class AsyncBeforeToolCallback:
 
 
 class AsyncAfterToolCallback:
+
   def __init__(self, mock_response: Dict[str, Any]):
     self.mock_response = mock_response