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"