Skip to content

Commit 8dc03a0

Browse files
committed
perf(tools): cache the FunctionTool declaration between LLM requests
FunctionTool._get_declaration() rebuilds the declaration from the function signature on every call, and LlmRequest.append_tools() calls it once per tool on every LLM request. The build runs pydantic create_model() plus model_json_schema() for the parameters and a TypeAdapter().json_schema() for the return type, which costs 0.53-0.71 ms of CPU per tool for the five tools of contributing/samples/evaluation/home_automation_agent. Cache the built declaration per tool, keyed on the function, the ignored parameters and the API variant, since all three can change during the life of a tool. Callers mutate the declaration they are given - LongRunningFunctionTool appends to its description, TransferToAgentTool writes an enum into its parameter schema and BaseToolset rewrites its name when prefixing - so each call still returns a private copy and the cached value stays pristine.
1 parent 8addc44 commit 8dc03a0

2 files changed

Lines changed: 235 additions & 11 deletions

File tree

src/google/adk/tools/function_tool.py

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
from __future__ import annotations
1616

17+
import copy
1718
import inspect
1819
import logging
1920
from typing import Any
@@ -47,6 +48,13 @@ class FunctionTool(BaseTool):
4748
func: The function to wrap.
4849
"""
4950

51+
# Declaration built for the last seen (func, ignore_params, api variant).
52+
# Declared on the class so instances created without __init__ (for example
53+
# via copy.copy) still resolve it.
54+
_declaration_cache: Optional[
55+
tuple[Callable[..., Any], tuple[str, ...], Any, types.FunctionDeclaration]
56+
] = None
57+
5058
def __init__(
5159
self,
5260
func: Callable[..., Any],
@@ -92,17 +100,33 @@ def __init__(
92100

93101
@override
94102
def _get_declaration(self) -> Optional[types.FunctionDeclaration]:
95-
function_decl = types.FunctionDeclaration.model_validate(
96-
build_function_declaration(
97-
func=self.func,
98-
# The model doesn't understand the function context.
99-
# input_stream is for streaming tool
100-
ignore_params=self._ignore_params,
101-
variant=self._api_variant,
102-
)
103-
)
104-
105-
return function_decl
103+
variant = self._api_variant
104+
ignore_params = tuple(self._ignore_params)
105+
cache = self._declaration_cache
106+
if (
107+
cache is None
108+
or cache[0] is not self.func
109+
or cache[1] != ignore_params
110+
or cache[2] != variant
111+
):
112+
function_decl = types.FunctionDeclaration.model_validate(
113+
build_function_declaration(
114+
func=self.func,
115+
# The model doesn't understand the function context.
116+
# input_stream is for streaming tool
117+
ignore_params=self._ignore_params,
118+
variant=variant,
119+
)
120+
)
121+
cache = (self.func, ignore_params, variant, function_decl)
122+
self._declaration_cache = cache
123+
124+
# Callers are allowed to mutate the declaration they receive:
125+
# LongRunningFunctionTool appends to its description, TransferToAgentTool
126+
# writes an enum into its parameter schema, and BaseToolset rewrites its
127+
# name when prefixing. Hand back a private copy so the cached value stays
128+
# pristine.
129+
return copy.deepcopy(cache[3])
106130

107131
def _preprocess_args(self, args: dict[str, Any]) -> dict[str, Any]:
108132
"""Preprocess and convert function arguments before invocation.
Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,200 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Tests for the FunctionTool declaration cache.
16+
17+
`_get_declaration` is called once per tool on every LLM request, so its result
18+
is cached. Callers are allowed to mutate the declaration they receive, so each
19+
call must still get a private copy, and the cache must follow every input the
20+
declaration depends on.
21+
"""
22+
23+
import copy
24+
25+
from google.adk.tools import _function_tool_declarations
26+
from google.adk.tools.function_tool import FunctionTool
27+
from google.adk.tools.long_running_tool import LongRunningFunctionTool
28+
from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool
29+
from google.adk.utils.variant_utils import GoogleLLMVariant
30+
import pytest
31+
32+
33+
def sample_tool(device_id: str, status: str = 'ON') -> str:
34+
"""Sets a device status.
35+
36+
Args:
37+
device_id: The device.
38+
status: The desired status.
39+
40+
Returns:
41+
A confirmation message.
42+
"""
43+
return 'ok'
44+
45+
46+
def poll_job(job_id: str) -> str:
47+
"""Polls a job.
48+
49+
Args:
50+
job_id: The job.
51+
52+
Returns:
53+
The job status.
54+
"""
55+
return 'PENDING'
56+
57+
58+
def test_repeated_calls_return_equal_declarations():
59+
tool = FunctionTool(sample_tool)
60+
61+
first = tool._get_declaration()
62+
second = tool._get_declaration()
63+
64+
assert first.model_dump() == second.model_dump()
65+
66+
67+
def test_repeated_calls_return_distinct_objects():
68+
tool = FunctionTool(sample_tool)
69+
70+
first = tool._get_declaration()
71+
second = tool._get_declaration()
72+
73+
assert first is not second
74+
assert first.parameters_json_schema is not second.parameters_json_schema
75+
76+
77+
def test_declaration_is_built_once_for_repeated_calls():
78+
tool = FunctionTool(sample_tool)
79+
calls = []
80+
original = (
81+
_function_tool_declarations.build_function_declaration_with_json_schema
82+
)
83+
84+
def counting(*args, **kwargs):
85+
calls.append(1)
86+
return original(*args, **kwargs)
87+
88+
_function_tool_declarations.build_function_declaration_with_json_schema = (
89+
counting
90+
)
91+
try:
92+
for _ in range(5):
93+
tool._get_declaration()
94+
finally:
95+
_function_tool_declarations.build_function_declaration_with_json_schema = (
96+
original
97+
)
98+
99+
assert len(calls) == 1
100+
101+
102+
def test_mutating_a_returned_declaration_does_not_affect_later_calls():
103+
tool = FunctionTool(sample_tool)
104+
105+
first = tool._get_declaration()
106+
first.name = 'renamed'
107+
first.description = 'mutated'
108+
first.parameters_json_schema['properties']['device_id']['enum'] = ['x']
109+
110+
second = tool._get_declaration()
111+
112+
assert second.name == 'sample_tool'
113+
assert second.description != 'mutated'
114+
assert 'enum' not in second.parameters_json_schema['properties']['device_id']
115+
116+
117+
def test_long_running_tool_description_is_stable_across_calls():
118+
"""LongRunningFunctionTool appends a note to the declaration it is given."""
119+
tool = LongRunningFunctionTool(poll_job)
120+
121+
descriptions = [tool._get_declaration().description for _ in range(3)]
122+
123+
assert len(set(descriptions)) == 1
124+
assert descriptions[0].count('long-running operation') == 1
125+
126+
127+
def test_transfer_to_agent_tool_sets_enum_on_every_call():
128+
"""TransferToAgentTool writes an enum into the parameter schema it is given."""
129+
tool = TransferToAgentTool(agent_names=['alpha', 'beta'])
130+
131+
for _ in range(3):
132+
declaration = tool._get_declaration()
133+
assert declaration.parameters_json_schema['properties']['agent_name'][
134+
'enum'
135+
] == ['alpha', 'beta']
136+
137+
138+
def test_toolset_name_prefixing_does_not_rename_the_original_tool():
139+
"""BaseToolset prefixing writes .name onto the declaration it is given."""
140+
tool = FunctionTool(sample_tool)
141+
prefixed = copy.copy(tool)
142+
original_get_declaration = tool._get_declaration
143+
144+
def prefixed_declaration():
145+
declaration = original_get_declaration()
146+
declaration.name = 'prefix_sample_tool'
147+
return declaration
148+
149+
prefixed._get_declaration = prefixed_declaration
150+
151+
assert prefixed._get_declaration().name == 'prefix_sample_tool'
152+
assert tool._get_declaration().name == 'sample_tool'
153+
154+
155+
def test_changing_ignore_params_rebuilds_the_declaration():
156+
tool = FunctionTool(sample_tool)
157+
158+
before = tool._get_declaration().parameters_json_schema['properties']
159+
assert 'status' in before
160+
161+
tool._ignore_params = list(tool._ignore_params) + ['status']
162+
after = tool._get_declaration().parameters_json_schema['properties']
163+
164+
assert 'status' not in after
165+
166+
167+
def test_changing_func_rebuilds_the_declaration():
168+
tool = FunctionTool(sample_tool)
169+
assert tool._get_declaration().name == 'sample_tool'
170+
171+
tool.func = poll_job
172+
173+
assert tool._get_declaration().name == 'poll_job'
174+
175+
176+
@pytest.mark.parametrize(
177+
'variant',
178+
[GoogleLLMVariant.GEMINI_API, GoogleLLMVariant.VERTEX_AI],
179+
)
180+
def test_changing_api_variant_rebuilds_the_declaration(monkeypatch, variant):
181+
"""The api variant is resolved from the environment on every call."""
182+
tool = FunctionTool(sample_tool)
183+
184+
monkeypatch.setattr(
185+
type(tool),
186+
'_api_variant',
187+
property(lambda self: GoogleLLMVariant.GEMINI_API),
188+
)
189+
studio = tool._get_declaration()
190+
191+
monkeypatch.setattr(
192+
type(tool), '_api_variant', property(lambda self: variant)
193+
)
194+
switched = tool._get_declaration()
195+
196+
if variant is GoogleLLMVariant.GEMINI_API:
197+
assert switched.model_dump() == studio.model_dump()
198+
else:
199+
assert switched.response_json_schema is not None
200+
assert studio.response_json_schema is None

0 commit comments

Comments
 (0)