Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions src/google/adk/tools/load_artifacts_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand Down
14 changes: 14 additions & 0 deletions src/google/adk/tools/load_memory_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
18 changes: 18 additions & 0 deletions src/google/adk/tools/retrieval/base_retrieval_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,16 +12,34 @@
# 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


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,
Expand Down
67 changes: 67 additions & 0 deletions tests/unittests/tools/retrieval/test_base_retrieval_tool.py
Original file line number Diff line number Diff line change
@@ -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.',
},
},
}
20 changes: 20 additions & 0 deletions tests/unittests/tools/test_load_artifacts_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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'},
},
},
}
46 changes: 46 additions & 0 deletions tests/unittests/tools/test_load_memory_tool.py
Original file line number Diff line number Diff line change
@@ -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'],
}
Loading