diff --git a/src/google/adk/tools/load_artifacts_tool.py b/src/google/adk/tools/load_artifacts_tool.py index dbdc1f26f2..8f82d4df7a 100644 --- a/src/google/adk/tools/load_artifacts_tool.py +++ b/src/google/adk/tools/load_artifacts_tool.py @@ -24,6 +24,8 @@ from google.genai import types from typing_extensions import override +from ..features import FeatureName +from ..features import is_feature_enabled from .base_tool import BaseTool # MIME types Gemini accepts for inline data in requests. @@ -132,6 +134,20 @@ def __init__(self): ) def _get_declaration(self) -> types.FunctionDeclaration | None: + if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL): + return types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'artifact_names': { + 'type': 'array', + 'items': {'type': 'string'}, + }, + }, + }, + ) return types.FunctionDeclaration( name=self.name, description=self.description, diff --git a/src/google/adk/tools/load_memory_tool.py b/src/google/adk/tools/load_memory_tool.py index 8410e41141..66e2ffabc3 100644 --- a/src/google/adk/tools/load_memory_tool.py +++ b/src/google/adk/tools/load_memory_tool.py @@ -21,6 +21,8 @@ from pydantic import Field from typing_extensions import override +from ..features import FeatureName +from ..features import is_feature_enabled from ..memory.memory_entry import MemoryEntry from .function_tool import FunctionTool from .tool_context import ToolContext @@ -59,6 +61,18 @@ def __init__(self): @override def _get_declaration(self) -> types.FunctionDeclaration | None: + if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL): + return types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'query': {'type': 'string'}, + }, + 'required': ['query'], + }, + ) return types.FunctionDeclaration( name=self.name, description=self.description, diff --git a/src/google/adk/tools/retrieval/base_retrieval_tool.py b/src/google/adk/tools/retrieval/base_retrieval_tool.py index 64f3ec91db..f101be9de3 100644 --- a/src/google/adk/tools/retrieval/base_retrieval_tool.py +++ b/src/google/adk/tools/retrieval/base_retrieval_tool.py @@ -12,9 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + from google.genai import types from typing_extensions import override +from ...features import FeatureName +from ...features import is_feature_enabled from ..base_tool import BaseTool @@ -22,6 +26,20 @@ class BaseRetrievalTool(BaseTool): @override def _get_declaration(self) -> types.FunctionDeclaration: + if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL): + return types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'query': { + 'type': 'string', + 'description': 'The query to retrieve.', + }, + }, + }, + ) return types.FunctionDeclaration( name=self.name, description=self.description, diff --git a/tests/unittests/tools/retrieval/test_base_retrieval_tool.py b/tests/unittests/tools/retrieval/test_base_retrieval_tool.py new file mode 100644 index 0000000000..8c40f49265 --- /dev/null +++ b/tests/unittests/tools/retrieval/test_base_retrieval_tool.py @@ -0,0 +1,67 @@ +# Copyright 2025 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. + +from google.adk.features import FeatureName +from google.adk.features._feature_registry import temporary_feature_override +from google.adk.tools.retrieval.base_retrieval_tool import BaseRetrievalTool +from google.genai import types + + +class _TestRetrievalTool(BaseRetrievalTool): + """Concrete implementation of BaseRetrievalTool for testing.""" + + def __init__(self): + super().__init__( + name='test_retrieval', + description='A test retrieval tool.', + ) + + async def run_async(self, *, args, tool_context): + return {'result': 'test'} + + +def test_get_declaration_with_json_schema_feature_disabled(): + """Test that _get_declaration uses parameters when feature is disabled.""" + tool = _TestRetrievalTool() + + with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, False): + declaration = tool._get_declaration() + + assert declaration.name == 'test_retrieval' + assert declaration.description == 'A test retrieval tool.' + assert declaration.parameters_json_schema is None + assert isinstance(declaration.parameters, types.Schema) + assert declaration.parameters.type == types.Type.OBJECT + assert 'query' in declaration.parameters.properties + + +def test_get_declaration_with_json_schema_feature_enabled(): + """Test that _get_declaration uses parameters_json_schema when feature is enabled.""" + tool = _TestRetrievalTool() + + with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True): + declaration = tool._get_declaration() + + assert declaration.name == 'test_retrieval' + assert declaration.description == 'A test retrieval tool.' + assert declaration.parameters is None + assert declaration.parameters_json_schema == { + 'type': 'object', + 'properties': { + 'query': { + 'type': 'string', + 'description': 'The query to retrieve.', + }, + }, + } diff --git a/tests/unittests/tools/test_load_artifacts_tool.py b/tests/unittests/tools/test_load_artifacts_tool.py index 1ea50bb33c..6a420574f0 100644 --- a/tests/unittests/tools/test_load_artifacts_tool.py +++ b/tests/unittests/tools/test_load_artifacts_tool.py @@ -14,6 +14,8 @@ import base64 +from google.adk.features import FeatureName +from google.adk.features._feature_registry import temporary_feature_override from google.adk.models.llm_request import LlmRequest from google.adk.tools.load_artifacts_tool import _maybe_base64_to_bytes from google.adk.tools.load_artifacts_tool import load_artifacts_tool @@ -160,3 +162,21 @@ def test_maybe_base64_to_bytes_returns_none_for_invalid(): """Invalid base64 strings return None.""" # Single character is invalid (base64 requires length % 4 == 0 after padding) assert _maybe_base64_to_bytes('x') is None + + +def test_get_declaration_with_json_schema_feature_enabled(): + """Test that _get_declaration uses parameters_json_schema when feature is enabled.""" + with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True): + declaration = load_artifacts_tool._get_declaration() + + assert declaration.name == 'load_artifacts' + assert declaration.parameters is None + assert declaration.parameters_json_schema == { + 'type': 'object', + 'properties': { + 'artifact_names': { + 'type': 'array', + 'items': {'type': 'string'}, + }, + }, + } diff --git a/tests/unittests/tools/test_load_memory_tool.py b/tests/unittests/tools/test_load_memory_tool.py new file mode 100644 index 0000000000..a3affd3ed3 --- /dev/null +++ b/tests/unittests/tools/test_load_memory_tool.py @@ -0,0 +1,46 @@ +# Copyright 2025 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. + +from google.adk.features import FeatureName +from google.adk.features._feature_registry import temporary_feature_override +from google.adk.tools.load_memory_tool import load_memory_tool +from google.genai import types + + +def test_get_declaration_with_json_schema_feature_disabled(): + """Test that _get_declaration uses parameters when feature is disabled.""" + with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, False): + declaration = load_memory_tool._get_declaration() + + assert declaration.name == 'load_memory' + assert declaration.parameters_json_schema is None + assert isinstance(declaration.parameters, types.Schema) + assert declaration.parameters.type == types.Type.OBJECT + assert 'query' in declaration.parameters.properties + + +def test_get_declaration_with_json_schema_feature_enabled(): + """Test that _get_declaration uses parameters_json_schema when feature is enabled.""" + with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True): + declaration = load_memory_tool._get_declaration() + + assert declaration.name == 'load_memory' + assert declaration.parameters is None + assert declaration.parameters_json_schema == { + 'type': 'object', + 'properties': { + 'query': {'type': 'string'}, + }, + 'required': ['query'], + }