8000 Refactor function callback handling and update type signatures · PaveLuchkov/adk-python@08ac9a1 · GitHub
[go: up one dir, main page]
Skip to content

Commit 08ac9a1

Browse files
AlankritVerma01AlankritVerma01
committed
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
1 parent 2173606 commit 08ac9a1

File tree

2 files changed

+18
-24
lines changed

2 files changed

+18
-24
lines changed

‎src/google/adk/agents/llm_agent.py

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,7 @@
1515
from __future__ import annotations
1616

1717
import logging
18-
from typing import Any
19-
from typing import AsyncGenerator
20-
from typing import Callable
21-
from typing import Literal
22-
from typing import Optional
23-
from typing import Union
18+
from typing import Any, AsyncGenerator, Awaitable, Callable, Literal, Optional, Union
2419

2520
from google.genai import types
2621
from pydantic import BaseModel
@@ -62,11 +57,11 @@
6257
]
6358
BeforeToolCallback: TypeAlias = Callable[
6459
[BaseTool, dict[str, Any], ToolContext],
65-
Optional[dict],
60+
Awaitable[Optional[dict]],
6661
]
6762
AfterToolCallback: TypeAlias = Callable[
6863
[BaseTool, dict[str, Any], ToolContext, dict],
69-
Optional[dict],
64+
Awaitable[Optional[dict]],
7065
]
7166

7267
InstructionProvider: TypeAlias = Callable[[ReadonlyContext], str]

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

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -237,12 +237,11 @@ async def handle_function_calls_live(
237237
# tool, function_args, tool_context
238238
# )
239239
if agent.before_tool_callback:
240-
_maybe = agent.before_tool_callback(
241-
tool=tool, args=function_args, tool_context=tool_context
242-
)
243-
if inspect.isawaitable(_maybe):
244-
_maybe = await _maybe
245-
function_response = _maybe
240+
function_response = agent.before_tool_callback(
241+
tool=tool, args=function_args, tool_context=tool_context
242+
)
243+
if inspect.isawaitable(function_response):
244+
function_response = await function_response
246245

247246
if not function_response:
248247
function_response = await _process_function_live_helper(
@@ -260,16 +259,16 @@ async def handle_function_calls_live(
260259
# if new_response:
261260
# function_response = new_response
262261
if agent.after_tool_callback:
263-
_maybe2 = agent.after_tool_callback(
264-
tool=tool,
265-
args=function_args,
266-
tool_context=tool_context,
267-
tool_response=function_response,
268-
)
269-
if inspect.isawaitable(_maybe2):
270-
_maybe2 = await _maybe2
271-
if _maybe2 is not None:
272-
function_response = _maybe2
262+
altered_function_response = agent.after_tool_callback(
263+
tool=tool,
264+
args=function_args,
265+
tool_context=tool_context,
266+
tool_response=function_response,
267+
)
268+
if inspect.isawaitable(altered_function_response):
269+
altered_function_response = await altered_function_response
270+
if altered_function_response is not None:
271+
function_response = altered_function_response
273272

274273
if tool.is_long_running:
275274
# Allow async function to return None to not provide function response.

0 commit comments

Comments
 (0)
0