Support chaining for tool callbacks

(before/after) tool callbacks are invoked throughout the provided chain until one callback does not return None. Callbacks can be async and sync.

PiperOrigin-RevId: 756526507
This commit is contained in:
Selcuk Gun
2025-05-08 17:38:04 -07:00
committed by Copybara-Service
parent 0299020cc4
commit 2cbbf88135
5 changed files with 282 additions and 17 deletions
+7 -6
View File
@@ -153,22 +153,22 @@ async def handle_function_calls_async(
function_args = function_call.args or {}
function_response: Optional[dict] = None
# before_tool_callback (sync or async)
if agent.before_tool_callback:
function_response = agent.before_tool_callback(
for callback in agent.canonical_before_tool_callbacks:
function_response = callback(
tool=tool, args=function_args, tool_context=tool_context
)
if inspect.isawaitable(function_response):
function_response = await function_response
if function_response:
break
if not function_response:
function_response = await __call_tool_async(
tool, args=function_args, tool_context=tool_context
)
# after_tool_callback (sync or async)
if agent.after_tool_callback:
altered_function_response = agent.after_tool_callback(
for callback in agent.canonical_after_tool_callbacks:
altered_function_response = callback(
tool=tool,
args=function_args,
tool_context=tool_context,
@@ -178,6 +178,7 @@ async def handle_function_calls_async(
altered_function_response = await altered_function_response
if altered_function_response is not None:
function_response = altered_function_response
break
if tool.is_long_running:
# Allow long running function to return None to not provide function response.