From e8501cd5f8dc20f999acef2c0d1f212b5f647006 Mon Sep 17 00:00:00 2001 From: George Weale Date: Tue, 28 Jul 2026 14:43:58 -0700 Subject: [PATCH] perf: cache the FunctionTool declaration across LLM calls FunctionTool._get_declaration ran pydantic create_model + JSON-schema generation for every tool on every LLM step, even though the result is fixed once the tool is constructed. Memoize the build keyed by (func, ignored params, API variant, feature flag) and return a copy so callers (e.g. toolset prefixing) can still mutate the result. The cache is module-level rather than per-instance because a bare callable in LlmAgent.tools is re-wrapped into a fresh FunctionTool on every step, which defeats an instance cache for exactly the callers this costs the most. Co-authored-by: George Weale PiperOrigin-RevId: 955481497 (cherry picked from commit 57f3af24a00de46096089bd3279791a6a98b5c48) --- src/google/adk/tools/function_tool.py | 47 ++++++++++++++++----- tests/unittests/tools/test_function_tool.py | 28 ++++++++++++ 2 files changed, 65 insertions(+), 10 deletions(-) diff --git a/src/google/adk/tools/function_tool.py b/src/google/adk/tools/function_tool.py index ee63fc3873..bd9e0eed10 100644 --- a/src/google/adk/tools/function_tool.py +++ b/src/google/adk/tools/function_tool.py @@ -16,6 +16,7 @@ from contextlib import contextmanager import contextvars +import functools import inspect import logging from typing import Any @@ -32,7 +33,10 @@ import pydantic from typing_extensions import override +from ..features import FeatureName +from ..features import is_feature_enabled from ..utils.context_utils import find_context_parameter +from ..utils.variant_utils import GoogleLLMVariant from ._automatic_function_calling_util import build_function_declaration from .base_tool import BaseTool from .tool_context import ToolContext @@ -63,6 +67,30 @@ def _use_sync_callable_runner( _SYNC_CALLABLE_RUNNER.reset(token) +@functools.lru_cache(maxsize=1024) +def _build_declaration_cached( + func: Callable[..., Any], + ignore_params: tuple[str, ...], + variant: GoogleLLMVariant, + json_schema_enabled: bool, +) -> types.FunctionDeclaration: + """Builds (and caches) a tool's FunctionDeclaration. + + The build runs pydantic ``create_model`` + JSON-schema generation, which is + expensive and otherwise re-run for every tool on every LLM call even though + the result depends only on these (static) inputs. ``json_schema_enabled`` is + part of the key so toggling the feature flag rebuilds. + """ + del json_schema_enabled # Only participates in the cache key. + return types.FunctionDeclaration.model_validate( + build_function_declaration( + func=func, + ignore_params=list(ignore_params), + variant=variant, + ) + ) + + class FunctionTool(BaseTool): """A tool that wraps a user-defined Python function. @@ -115,17 +143,16 @@ def __init__( @override def _get_declaration(self) -> Optional[types.FunctionDeclaration]: - function_decl = types.FunctionDeclaration.model_validate( - build_function_declaration( - func=self.func, - # The model doesn't understand the function context. - # input_stream is for streaming tool - ignore_params=self._ignore_params, - variant=self._api_variant, - ) + # `ignore_params` drops the function context and input_stream (for streaming + # tools), which the model doesn't understand. Return a copy: the cached + # declaration is shared and callers (e.g. toolset prefixing) mutate it. + declaration = _build_declaration_cached( + self.func, + tuple(self._ignore_params), + self._api_variant, + is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL), ) - - return function_decl + return declaration.model_copy(deep=True) def _preprocess_args(self, args: dict[str, Any]) -> dict[str, Any]: """Preprocess and convert function arguments before invocation. diff --git a/tests/unittests/tools/test_function_tool.py b/tests/unittests/tools/test_function_tool.py index fbb16821fc..5323bde2d7 100644 --- a/tests/unittests/tools/test_function_tool.py +++ b/tests/unittests/tools/test_function_tool.py @@ -20,6 +20,7 @@ from google.adk.agents.context import Context from google.adk.agents.invocation_context import InvocationContext from google.adk.sessions.session import Session +from google.adk.tools.function_tool import _build_declaration_cached from google.adk.tools.function_tool import FunctionTool from google.adk.tools.tool_confirmation import ToolConfirmation from google.adk.tools.tool_context import ToolContext @@ -644,3 +645,30 @@ async def streaming_tool_req(req_param: str): assert isinstance(result, dict) assert "error" in result assert "mandatory input parameters are not present" in result["error"] + + +def test_get_declaration_is_cached_and_returns_independent_copies(): + """_get_declaration caches the build and hands out independent copies.""" + + def sample_tool(a: int, b: str) -> str: + """A sample tool.""" + return b * a + + _build_declaration_cached.cache_clear() + tool = FunctionTool(func=sample_tool) + + d1 = tool._get_declaration() # pylint: disable=protected-access + d2 = tool._get_declaration() # pylint: disable=protected-access + + # The expensive build runs once; the second call is served from cache. + info = _build_declaration_cached.cache_info() + assert info.misses == 1 + assert info.hits >= 1 + + assert d1.name == d2.name == "sample_tool" + + # Callers (e.g. toolset prefixing) mutate the returned declaration, so each + # call must return an independent copy rather than the shared cached object. + d1.name = "prefixed_sample_tool" + d3 = tool._get_declaration() # pylint: disable=protected-access + assert d3.name == "sample_tool"