Copybara import of the project:

--
21736067f9 by Alankrit Verma <alankrit386@gmail.com>:

feat(llm_flows): support async before/after tool callbacks

Previously, callbacks were treated as purely synchronous,
so passing an async coroutine caused “was never awaited”
errors and Pydantic serialization failures.

Now we detect awaitable return values from
before_tool_callback and after_tool_callback,
and `await` them if necessary.

Fixes: #380

--
08ac9a117e by Alankrit Verma <alankrit386@gmail.com>:

Refactor function callback handling and update type signatures

- Simplify variable names in `functions.py`: always use `function_response` and `altered_function_response`
- Update LlmAgent callback type aliases to support async:
  - Import `Awaitable`
  - Change `BeforeToolCallback` and `AfterToolCallback` signatures to return `Awaitable[Optional[dict]]`
- Ensure `after_tool_callback` uses `await` when necessary

--
fcbf57466e by Alankrit Verma <alankrit386@gmail.com>:

refactor: update callback type signatures to support sync and async responses
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/434 from AlankritVerma01:support-async-tool-callbacks 926b0ef1a6
PiperOrigin-RevId: 753005846
This commit is contained in:
Alankrit Verma
2025-04-29 21:40:05 -07:00
committed by Copybara-Service
parent dbbeb190b0
commit 27ce65ff50
4 changed files with 150 additions and 23 deletions

View File

@@ -15,12 +15,7 @@
from __future__ import annotations
import logging
from typing import Any
from typing import AsyncGenerator
from typing import Callable
from typing import Literal
from typing import Optional
from typing import Union
from typing import Any, AsyncGenerator, Awaitable, Callable, Literal, Optional, Union
from google.genai import types
from pydantic import BaseModel
@@ -62,11 +57,11 @@ AfterModelCallback: TypeAlias = Callable[
]
BeforeToolCallback: TypeAlias = Callable[
[BaseTool, dict[str, Any], ToolContext],
Optional[dict],
Union[Awaitable[Optional[dict]], Optional[dict]],
]
AfterToolCallback: TypeAlias = Callable[
[BaseTool, dict[str, Any], ToolContext, dict],
Optional[dict],
Union[Awaitable[Optional[dict]], Optional[dict]],
]
InstructionProvider: TypeAlias = Callable[[ReadonlyContext], str]

View File

@@ -151,28 +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.
function_response: Optional[dict] = None
# before_tool_callback (sync or async)
if agent.before_tool_callback:
function_response = agent.before_tool_callback(
tool=tool, args=function_args, tool_context=tool_context
)
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:
new_response = 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 new_response:
function_response = new_response
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.
@@ -223,11 +228,17 @@ async def handle_function_calls_live(
# 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.
# # 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, function_args, tool_context
# )
if agent.before_tool_callback:
function_response = agent.before_tool_callback(
tool, function_args, tool_context
tool=tool, args=function_args, tool_context=tool_context
)
if inspect.isawaitable(function_response):
function_response = await function_response
if not function_response:
function_response = await _process_function_live_helper(
@@ -235,15 +246,26 @@ async def handle_function_calls_live(
)
# Calls after_tool_callback if it exists.
# if agent.after_tool_callback:
# new_response = agent.after_tool_callback(
# tool,
# function_args,
# tool_context,
# function_response,
# )
# if new_response:
# function_response = new_response
if agent.after_tool_callback:
new_response = agent.after_tool_callback(
tool,
function_args,
tool_context,
function_response,
altered_function_response = agent.after_tool_callback(
tool=tool,
args=function_args,
tool_context=tool_context,
tool_response=function_response,
)
if new_response:
function_response = new_response
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 async function to return None to not provide function response.