diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index def9aba4cf..28399df5ec 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -2858,7 +2858,10 @@ def _finalize_tool_call_response( ) mapped_finish_reason = _map_finish_reason(finish_reason) llm_response.finish_reason = mapped_finish_reason - if mapped_finish_reason != types.FinishReason.STOP: + if ( + mapped_finish_reason is not None + and mapped_finish_reason != types.FinishReason.STOP + ): llm_response.error_code = mapped_finish_reason llm_response.error_message = _finish_reason_to_error_message( mapped_finish_reason @@ -2879,7 +2882,10 @@ def _finalize_text_response( ) mapped_finish_reason = _map_finish_reason(finish_reason) llm_response.finish_reason = mapped_finish_reason - if mapped_finish_reason != types.FinishReason.STOP: + if ( + mapped_finish_reason is not None + and mapped_finish_reason != types.FinishReason.STOP + ): llm_response.error_code = mapped_finish_reason llm_response.error_message = _finish_reason_to_error_message( mapped_finish_reason diff --git a/src/google/adk/tools/function_tool.py b/src/google/adk/tools/function_tool.py index 9514c9dd21..2f5ce6c0b0 100644 --- a/src/google/adk/tools/function_tool.py +++ b/src/google/adk/tools/function_tool.py @@ -14,6 +14,7 @@ from __future__ import annotations +import copy import inspect import logging from typing import Any @@ -29,6 +30,8 @@ import pydantic from typing_extensions import override +from ..features import FeatureName +from ..features import is_feature_enabled from ..utils._schema_utils import get_list_inner_type from ..utils._schema_utils import is_list_of_basemodel from ..utils.context_utils import Aclosing @@ -47,6 +50,19 @@ class FunctionTool(BaseTool): func: The function to wrap. """ + # Declaration built for the last seen function, ignored parameters, API + # variant, and schema representation. Declared on the class so instances + # created without __init__ (for example via copy.copy) still resolve it. + _declaration_cache: Optional[ + tuple[ + Callable[..., Any], + tuple[str, ...], + Any, + bool, + types.FunctionDeclaration, + ] + ] = None + def __init__( self, func: Callable[..., Any], @@ -92,17 +108,43 @@ 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, - ) + variant = self._api_variant + ignore_params = tuple(self._ignore_params) + json_schema_enabled = is_feature_enabled( + FeatureName.JSON_SCHEMA_FOR_FUNC_DECL ) - - return function_decl + cache = self._declaration_cache + if ( + cache is None + or cache[0] is not self.func + or cache[1] != ignore_params + or cache[2] != variant + or cache[3] != json_schema_enabled + ): + 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=variant, + ) + ) + cache = ( + self.func, + ignore_params, + variant, + json_schema_enabled, + function_decl, + ) + self._declaration_cache = cache + + # Callers are allowed to mutate the declaration they receive: + # LongRunningFunctionTool appends to its description, TransferToAgentTool + # writes an enum into its parameter schema, and BaseToolset rewrites its + # name when prefixing. Hand back a private copy so the cached value stays + # pristine. + return copy.deepcopy(cache[4]) 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_declaration_cache.py b/tests/unittests/tools/test_function_tool_declaration_cache.py new file mode 100644 index 0000000000..05d71fcfb8 --- /dev/null +++ b/tests/unittests/tools/test_function_tool_declaration_cache.py @@ -0,0 +1,254 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the FunctionTool declaration cache. + +`_get_declaration` is called once per tool on every LLM request, so its result +is cached. Callers are allowed to mutate the declaration they receive, so each +call must still get a private copy, and the cache must follow every input the +declaration depends on. +""" + +import copy + +from google.adk.features import FeatureName +from google.adk.features._feature_registry import temporary_feature_override +from google.adk.tools import _function_tool_declarations +from google.adk.tools.function_tool import FunctionTool +from google.adk.tools.long_running_tool import LongRunningFunctionTool +from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool +from google.adk.utils.variant_utils import GoogleLLMVariant +import pytest + + +def sample_tool(device_id: str, status: str = 'ON') -> str: + """Sets a device status. + + Args: + device_id: The device. + status: The desired status. + + Returns: + A confirmation message. + """ + return 'ok' + + +def poll_job(job_id: str) -> str: + """Polls a job. + + Args: + job_id: The job. + + Returns: + The job status. + """ + return 'PENDING' + + +def test_repeated_calls_return_equal_declarations(): + tool = FunctionTool(sample_tool) + + first = tool._get_declaration() + second = tool._get_declaration() + + assert first.model_dump() == second.model_dump() + + +def test_repeated_calls_return_distinct_objects(): + tool = FunctionTool(sample_tool) + + first = tool._get_declaration() + second = tool._get_declaration() + + assert first is not second + assert first.parameters_json_schema is not second.parameters_json_schema + + +def test_declaration_is_built_once_for_repeated_calls(): + tool = FunctionTool(sample_tool) + calls = [] + original = ( + _function_tool_declarations.build_function_declaration_with_json_schema + ) + + def counting(*args, **kwargs): + calls.append(1) + return original(*args, **kwargs) + + _function_tool_declarations.build_function_declaration_with_json_schema = ( + counting + ) + try: + for _ in range(5): + tool._get_declaration() + finally: + _function_tool_declarations.build_function_declaration_with_json_schema = ( + original + ) + + assert len(calls) == 1 + + +def test_mutating_a_returned_declaration_does_not_affect_later_calls(): + tool = FunctionTool(sample_tool) + + first = tool._get_declaration() + first.name = 'renamed' + first.description = 'mutated' + first.parameters_json_schema['properties']['device_id']['enum'] = ['x'] + + second = tool._get_declaration() + + assert second.name == 'sample_tool' + assert second.description != 'mutated' + assert 'enum' not in second.parameters_json_schema['properties']['device_id'] + + +def test_long_running_tool_description_is_stable_across_calls(): + """LongRunningFunctionTool appends a note to the declaration it is given.""" + tool = LongRunningFunctionTool(poll_job) + + descriptions = [tool._get_declaration().description for _ in range(3)] + + assert len(set(descriptions)) == 1 + assert descriptions[0].count('long-running operation') == 1 + + +def test_transfer_to_agent_tool_sets_enum_on_every_call(): + """TransferToAgentTool writes an enum into the parameter schema it is given.""" + tool = TransferToAgentTool(agent_names=['alpha', 'beta']) + + for _ in range(3): + declaration = tool._get_declaration() + assert declaration.parameters_json_schema['properties']['agent_name'][ + 'enum' + ] == ['alpha', 'beta'] + + +def test_toolset_name_prefixing_does_not_rename_the_original_tool(): + """BaseToolset prefixing writes .name onto the declaration it is given.""" + tool = FunctionTool(sample_tool) + prefixed = copy.copy(tool) + original_get_declaration = tool._get_declaration + + def prefixed_declaration(): + declaration = original_get_declaration() + declaration.name = 'prefix_sample_tool' + return declaration + + prefixed._get_declaration = prefixed_declaration + + assert prefixed._get_declaration().name == 'prefix_sample_tool' + assert tool._get_declaration().name == 'sample_tool' + + +def test_changing_ignore_params_rebuilds_the_declaration(): + tool = FunctionTool(sample_tool) + + before = tool._get_declaration().parameters_json_schema['properties'] + assert 'status' in before + + tool._ignore_params = list(tool._ignore_params) + ['status'] + after = tool._get_declaration().parameters_json_schema['properties'] + + assert 'status' not in after + + +def test_changing_func_rebuilds_the_declaration(): + tool = FunctionTool(sample_tool) + assert tool._get_declaration().name == 'sample_tool' + + tool.func = poll_job + + assert tool._get_declaration().name == 'poll_job' + + +@pytest.mark.parametrize( + 'variant', + [GoogleLLMVariant.GEMINI_API, GoogleLLMVariant.VERTEX_AI], +) +def test_changing_api_variant_rebuilds_the_declaration(monkeypatch, variant): + """The api variant is resolved from the environment on every call.""" + tool = FunctionTool(sample_tool) + + monkeypatch.setattr( + type(tool), + '_api_variant', + property(lambda self: GoogleLLMVariant.GEMINI_API), + ) + studio = tool._get_declaration() + + monkeypatch.setattr( + type(tool), '_api_variant', property(lambda self: variant) + ) + switched = tool._get_declaration() + + if variant is GoogleLLMVariant.GEMINI_API: + assert switched.model_dump() == studio.model_dump() + else: + assert switched.response_json_schema is not None + assert studio.response_json_schema is None + + +@pytest.mark.parametrize( + 'variant', + [GoogleLLMVariant.GEMINI_API, GoogleLLMVariant.VERTEX_AI], +) +@pytest.mark.parametrize('first_enabled', [False, True]) +def test_changing_schema_representation_rebuilds_the_declaration( + monkeypatch, variant, first_enabled +): + """The feature flag selects legacy or JSON-schema declaration fields.""" + tool = FunctionTool(sample_tool) + monkeypatch.setattr( + type(tool), '_api_variant', property(lambda self: variant) + ) + + with temporary_feature_override( + FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, first_enabled + ): + first = tool._get_declaration() + with temporary_feature_override( + FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, not first_enabled + ): + second = tool._get_declaration() + + declarations = {first_enabled: first, not first_enabled: second} + json_schema = declarations[True] + legacy = declarations[False] + assert json_schema.parameters_json_schema is not None + assert json_schema.parameters is None + assert legacy.parameters_json_schema is None + assert legacy.parameters is not None + + +def test_provider_environment_aliases_rebuild_the_declaration(monkeypatch): + """Both provider environment variables resolve into the cache's API key.""" + tool = FunctionTool(sample_tool) + monkeypatch.delenv('GOOGLE_GENAI_USE_ENTERPRISE', raising=False) + monkeypatch.delenv('GOOGLE_GENAI_USE_VERTEXAI', raising=False) + + studio = tool._get_declaration() + monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', '1') + deprecated_vertex = tool._get_declaration() + monkeypatch.setenv('GOOGLE_GENAI_USE_ENTERPRISE', '0') + enterprise_precedence = tool._get_declaration() + monkeypatch.setenv('GOOGLE_GENAI_USE_ENTERPRISE', 'true') + enterprise_vertex = tool._get_declaration() + + assert studio.response_json_schema is None + assert deprecated_vertex.response_json_schema is not None + assert enterprise_precedence.response_json_schema is None + assert enterprise_vertex.response_json_schema is not None