Skip to content

Commit b4ce3b1

Browse files
seanzhougooglecopybara-github
authored andcommitted
fix: incorrect logic in LlmRequest.append_tools and make BaseTool to call it
PiperOrigin-RevId: 792642882
1 parent 98a5730 commit b4ce3b1

File tree

3 files changed

+346
-46
lines changed

3 files changed

+346
-46
lines changed

src/google/adk/models/llm_request.py

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,24 @@
2424
from ..tools.base_tool import BaseTool
2525

2626

27+
def _find_tool_with_function_declarations(
28+
llm_request: LlmRequest,
29+
) -> Optional[types.Tool]:
30+
"""Find an existing Tool with function_declarations in the LlmRequest."""
31+
# TODO: add individual tool with declaration and merge in google_llm.py
32+
if not llm_request.config or not llm_request.config.tools:
33+
return None
34+
35+
return next(
36+
(
37+
tool
38+
for tool in llm_request.config.tools
39+
if isinstance(tool, types.Tool) and tool.function_declarations
40+
),
41+
None,
42+
)
43+
44+
2745
class LlmRequest(BaseModel):
2846
"""LLM request class that allows passing in tools, output schema and system
2947
@@ -81,15 +99,26 @@ def append_tools(self, tools: list[BaseTool]) -> None:
8199
return
82100
declarations = []
83101
for tool in tools:
84-
if isinstance(tool, BaseTool):
85-
declaration = tool._get_declaration()
86-
else:
87-
declaration = tool.get_declaration()
102+
declaration = tool._get_declaration()
88103
if declaration:
89104
declarations.append(declaration)
90105
self.tools_dict[tool.name] = tool
91106
if declarations:
92-
self.config.tools.append(types.Tool(function_declarations=declarations))
107+
if self.config.tools is None:
108+
self.config.tools = []
109+
110+
# Find existing tool with function_declarations and append to it
111+
if tool_with_function_declarations := _find_tool_with_function_declarations(
112+
self
113+
):
114+
if tool_with_function_declarations.function_declarations is None:
115+
tool_with_function_declarations.function_declarations = []
116+
tool_with_function_declarations.function_declarations.extend(
117+
declarations
118+
)
119+
else:
120+
# No existing tool with function_declarations, create new one
121+
self.config.tools.append(types.Tool(function_declarations=declarations))
93122

94123
def set_output_schema(self, base_model: type[BaseModel]) -> None:
95124
"""Sets the output schema for the request.

src/google/adk/tools/base_tool.py

Lines changed: 2 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -125,30 +125,8 @@ async def process_llm_request(
125125
tool_context: The context of the tool.
126126
llm_request: The outgoing LLM request, mutable this method.
127127
"""
128-
if (function_declaration := self._get_declaration()) is None:
129-
return
130-
131-
llm_request.tools_dict[self.name] = self
132-
if tool_with_function_declarations := _find_tool_with_function_declarations(
133-
llm_request
134-
):
135-
if tool_with_function_declarations.function_declarations is None:
136-
tool_with_function_declarations.function_declarations = []
137-
tool_with_function_declarations.function_declarations.append(
138-
function_declaration
139-
)
140-
else:
141-
llm_request.config = (
142-
types.GenerateContentConfig()
143-
if not llm_request.config
144-
else llm_request.config
145-
)
146-
llm_request.config.tools = (
147-
[] if not llm_request.config.tools else llm_request.config.tools
148-
)
149-
llm_request.config.tools.append(
150-
types.Tool(function_declarations=[function_declaration])
151-
)
128+
# Use the consolidated logic in LlmRequest.append_tools
129+
llm_request.append_tools([self])
152130

153131
@property
154132
def _api_variant(self) -> GoogleLLMVariant:
@@ -232,20 +210,3 @@ def from_config(
232210
else:
233211
logger.warning("Unsupported parsing for argument: %s.", param_name)
234212
return cls(**kwargs)
235-
236-
237-
def _find_tool_with_function_declarations(
238-
llm_request: LlmRequest,
239-
) -> Optional[types.Tool]:
240-
# TODO: add individual tool with declaration and merge in google_llm.py
241-
if not llm_request.config or not llm_request.config.tools:
242-
return None
243-
244-
return next(
245-
(
246-
tool
247-
for tool in llm_request.config.tools
248-
if isinstance(tool, types.Tool) and tool.function_declarations
249-
),
250-
None,
251-
)

0 commit comments

Comments
 (0)