Codex/function tool cache v1

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.

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"