chore: Migrate Google tools to use the new feature decorator

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 842321126
This commit is contained in:
Xuan Yang
2025-12-09 11:21:06 -08:00
committed by Copybara-Service
parent 1ae944b39d
commit bab57296d5
13 changed files with 86 additions and 36 deletions
@@ -24,9 +24,16 @@ from ..utils.env_utils import is_env_enabled
class FeatureName(str, Enum):
"""Feature names."""
BIG_QUERY_TOOLSET = "BIG_QUERY_TOOLSET"
BIG_QUERY_TOOL_CONFIG = "BIG_QUERY_TOOL_CONFIG"
BIGTABLE_TOOL_SETTINGS = "BIGTABLE_TOOL_SETTINGS"
COMPUTER_USE = "COMPUTER_USE"
GOOGLE_CREDENTIALS_CONFIG = "GOOGLE_CREDENTIALS_CONFIG"
GOOGLE_TOOL = "GOOGLE_TOOL"
JSON_SCHEMA_FOR_FUNC_DECL = "JSON_SCHEMA_FOR_FUNC_DECL"
PROGRESSIVE_SSE_STREAMING = "PROGRESSIVE_SSE_STREAMING"
SPANNER_TOOLSET = "SPANNER_TOOLSET"
SPANNER_TOOL_SETTINGS = "SPANNER_TOOL_SETTINGS"
class FeatureStage(Enum):
@@ -59,15 +66,36 @@ class FeatureConfig:
# Central registry: FeatureName -> FeatureConfig
_FEATURE_REGISTRY: dict[FeatureName, FeatureConfig] = {
FeatureName.BIG_QUERY_TOOLSET: FeatureConfig(
FeatureStage.EXPERIMENTAL, default_on=True
),
FeatureName.BIG_QUERY_TOOL_CONFIG: FeatureConfig(
FeatureStage.EXPERIMENTAL, default_on=True
),
FeatureName.BIGTABLE_TOOL_SETTINGS: FeatureConfig(
FeatureStage.EXPERIMENTAL, default_on=True
),
FeatureName.COMPUTER_USE: FeatureConfig(
FeatureStage.EXPERIMENTAL, default_on=True
),
FeatureName.GOOGLE_CREDENTIALS_CONFIG: FeatureConfig(
FeatureStage.EXPERIMENTAL, default_on=True
),
FeatureName.GOOGLE_TOOL: FeatureConfig(
FeatureStage.EXPERIMENTAL, default_on=True
),
FeatureName.JSON_SCHEMA_FOR_FUNC_DECL: FeatureConfig(
FeatureStage.WIP, default_on=False
),
FeatureName.PROGRESSIVE_SSE_STREAMING: FeatureConfig(
FeatureStage.WIP, default_on=False
),
FeatureName.SPANNER_TOOLSET: FeatureConfig(
FeatureStage.EXPERIMENTAL, default_on=True
),
FeatureName.SPANNER_TOOL_SETTINGS: FeatureConfig(
FeatureStage.EXPERIMENTAL, default_on=True
),
}
# Track which experimental features have already warned (warn only once)
+3 -2
View File
@@ -33,11 +33,12 @@ from ..auth.auth_credential import AuthCredential
from ..auth.auth_credential import AuthCredentialTypes
from ..auth.auth_credential import OAuth2Auth
from ..auth.auth_tool import AuthConfig
from ..utils.feature_decorator import experimental
from ..features import experimental
from ..features import FeatureName
from .tool_context import ToolContext
@experimental
@experimental(FeatureName.GOOGLE_CREDENTIALS_CONFIG)
class BaseGoogleCredentialsConfig(BaseModel):
"""Base Google Credentials Configuration for Google API tools (Experimental).
@@ -14,14 +14,15 @@
from __future__ import annotations
from ...utils.feature_decorator import experimental
from ...features import experimental
from ...features import FeatureName
from .._google_credentials import BaseGoogleCredentialsConfig
BIGQUERY_TOKEN_CACHE_KEY = "bigquery_token_cache"
BIGQUERY_DEFAULT_SCOPE = ["https://www.googleapis.com/auth/bigquery"]
@experimental
@experimental(FeatureName.GOOGLE_CREDENTIALS_CONFIG)
class BigQueryCredentialsConfig(BaseGoogleCredentialsConfig):
"""BigQuery Credentials Configuration for Google API tools (Experimental).
@@ -24,16 +24,17 @@ from typing_extensions import override
from . import data_insights_tool
from . import metadata_tool
from . import query_tool
from ...features import experimental
from ...features import FeatureName
from ...tools.base_tool import BaseTool
from ...tools.base_toolset import BaseToolset
from ...tools.base_toolset import ToolPredicate
from ...tools.google_tool import GoogleTool
from ...utils.feature_decorator import experimental
from .bigquery_credentials import BigQueryCredentialsConfig
from .config import BigQueryToolConfig
@experimental
@experimental(FeatureName.BIG_QUERY_TOOLSET)
class BigQueryToolset(BaseToolset):
"""BigQuery Toolset contains tools for interacting with BigQuery data and metadata."""
+3 -2
View File
@@ -21,7 +21,8 @@ from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import field_validator
from ...utils.feature_decorator import experimental
from ...features import experimental
from ...features import FeatureName
class WriteMode(Enum):
@@ -47,7 +48,7 @@ class WriteMode(Enum):
"""All write operations are allowed."""
@experimental('Config defaults may have breaking change in the future.')
@experimental(FeatureName.BIG_QUERY_TOOL_CONFIG)
class BigQueryToolConfig(BaseModel):
"""Configuration for BigQuery tools."""
@@ -14,7 +14,8 @@
from __future__ import annotations
from ...utils.feature_decorator import experimental
from ...features import experimental
from ...features import FeatureName
from .._google_credentials import BaseGoogleCredentialsConfig
BIGTABLE_TOKEN_CACHE_KEY = "bigtable_token_cache"
@@ -24,7 +25,7 @@ BIGTABLE_DEFAULT_SCOPE = [
]
@experimental
@experimental(FeatureName.GOOGLE_CREDENTIALS_CONFIG)
class BigtableCredentialsConfig(BaseGoogleCredentialsConfig):
"""Bigtable Credentials Configuration for Google API tools (Experimental).
+3 -2
View File
@@ -16,10 +16,11 @@ from __future__ import annotations
from pydantic import BaseModel
from ...utils.feature_decorator import experimental
from ...features import experimental
from ...features import FeatureName
@experimental('Tool settings defaults may have breaking change in the future.')
@experimental(FeatureName.BIGTABLE_TOOL_SETTINGS)
class BigtableToolSettings(BaseModel):
"""Settings for Bigtable tools."""
+3 -2
View File
@@ -23,14 +23,15 @@ from google.auth.credentials import Credentials
from pydantic import BaseModel
from typing_extensions import override
from ..utils.feature_decorator import experimental
from ..features import experimental
from ..features import FeatureName
from ._google_credentials import BaseGoogleCredentialsConfig
from ._google_credentials import GoogleCredentialsManager
from .function_tool import FunctionTool
from .tool_context import ToolContext
@experimental
@experimental(FeatureName.GOOGLE_TOOL)
class GoogleTool(FunctionTool):
"""GoogleTool class for tools that call Google APIs.
+3 -2
View File
@@ -22,7 +22,8 @@ from typing import Optional
from pydantic import BaseModel
from pydantic import model_validator
from ...utils.feature_decorator import experimental
from ...features import experimental
from ...features import FeatureName
# Vector similarity search nearest neighbors search algorithms.
EXACT_NEAREST_NEIGHBORS = "EXACT_NEAREST_NEIGHBORS"
@@ -138,7 +139,7 @@ class SpannerVectorStoreSettings(BaseModel):
return self
@experimental("Tool settings defaults may have breaking change in the future.")
@experimental(FeatureName.SPANNER_TOOL_SETTINGS)
class SpannerToolSettings(BaseModel):
"""Settings for Spanner tools."""
@@ -14,7 +14,8 @@
from __future__ import annotations
from ...utils.feature_decorator import experimental
from ...features import experimental
from ...features import FeatureName
from .._google_credentials import BaseGoogleCredentialsConfig
SPANNER_TOKEN_CACHE_KEY = "spanner_token_cache"
@@ -24,7 +25,7 @@ SPANNER_DEFAULT_SCOPE = [
]
@experimental
@experimental(FeatureName.GOOGLE_CREDENTIALS_CONFIG)
class SpannerCredentialsConfig(BaseGoogleCredentialsConfig):
"""Spanner Credentials Configuration for Google API tools (Experimental).
@@ -24,11 +24,12 @@ from google.adk.tools.spanner import query_tool
from google.adk.tools.spanner import search_tool
from typing_extensions import override
from ...features import experimental
from ...features import FeatureName
from ...tools.base_tool import BaseTool
from ...tools.base_toolset import BaseToolset
from ...tools.base_toolset import ToolPredicate
from ...tools.google_tool import GoogleTool
from ...utils.feature_decorator import experimental
from .settings import Capabilities
from .settings import SpannerToolSettings
from .spanner_credentials import SpannerCredentialsConfig
@@ -36,7 +37,7 @@ from .spanner_credentials import SpannerCredentialsConfig
DEFAULT_SPANNER_TOOL_NAME_PREFIX = "spanner"
@experimental
@experimental(FeatureName.SPANNER_TOOLSET)
class SpannerToolset(BaseToolset):
"""Spanner Toolset contains tools for interacting with Spanner data, database and table information.
@@ -14,17 +14,25 @@
from __future__ import annotations
import warnings
from google.adk.features._feature_registry import _WARNED_FEATURES
from google.adk.tools.bigquery.config import BigQueryToolConfig
import pytest
@pytest.fixture(autouse=True)
def reset_warned_features():
"""Reset warned features before each test."""
_WARNED_FEATURES.clear()
def test_bigquery_tool_config_experimental_warning():
"""Test BigQueryToolConfig experimental warning."""
with pytest.warns(
UserWarning,
match="Config defaults may have breaking change in the future.",
):
with warnings.catch_warnings(record=True) as w:
BigQueryToolConfig()
assert len(w) == 1
assert "BIG_QUERY_TOOL_CONFIG is enabled." in str(w[0].message)
def test_bigquery_tool_config_invalid_property():
@@ -46,22 +54,19 @@ def test_bigquery_tool_config_invalid_application_name():
def test_bigquery_tool_config_max_query_result_rows_default():
"""Test BigQueryToolConfig max_query_result_rows default value."""
with pytest.warns(UserWarning):
config = BigQueryToolConfig()
config = BigQueryToolConfig()
assert config.max_query_result_rows == 50
def test_bigquery_tool_config_max_query_result_rows_custom():
"""Test BigQueryToolConfig max_query_result_rows custom value."""
with pytest.warns(UserWarning):
config = BigQueryToolConfig(max_query_result_rows=100)
config = BigQueryToolConfig(max_query_result_rows=100)
assert config.max_query_result_rows == 100
def test_bigquery_tool_config_valid_maximum_bytes_billed():
"""Test BigQueryToolConfig raises exception with valid max bytes billed."""
with pytest.warns(UserWarning):
config = BigQueryToolConfig(maximum_bytes_billed=10_485_760)
config = BigQueryToolConfig(maximum_bytes_billed=10_485_760)
assert config.maximum_bytes_billed == 10_485_760
@@ -98,8 +103,7 @@ def test_bigquery_tool_config_invalid_maximum_bytes_billed():
)
def test_bigquery_tool_config_valid_labels(labels):
"""Test BigQueryToolConfig accepts valid labels."""
with pytest.warns(UserWarning):
config = BigQueryToolConfig(job_labels=labels)
config = BigQueryToolConfig(job_labels=labels)
assert config.job_labels == labels
@@ -14,6 +14,9 @@
from __future__ import annotations
import warnings
from google.adk.features._feature_registry import _WARNED_FEATURES
from google.adk.tools.spanner.settings import Capabilities
from google.adk.tools.spanner.settings import QueryResultMode
from google.adk.tools.spanner.settings import SpannerToolSettings
@@ -22,6 +25,12 @@ from pydantic import ValidationError
import pytest
@pytest.fixture(autouse=True)
def reset_warned_features():
"""Reset warned features before each test."""
_WARNED_FEATURES.clear()
def common_spanner_vector_store_settings(vector_length=None):
return {
"project_id": "test-project",
@@ -36,11 +45,10 @@ def common_spanner_vector_store_settings(vector_length=None):
def test_spanner_tool_settings_experimental_warning():
"""Test SpannerToolSettings experimental warning."""
with pytest.warns(
UserWarning,
match="Tool settings defaults may have breaking change in the future.",
):
with warnings.catch_warnings(record=True) as w:
SpannerToolSettings()
assert len(w) == 1
assert "SPANNER_TOOL_SETTINGS is enabled." in str(w[0].message)
def test_spanner_vector_store_settings_all_fields_present():