chore: Migrate Google tools to use the new feature decorator

xuanyang15 · copybara-github · commit bab57296d553 · 2025-12-09T11:21:06.000-08:00
Co-authored-by: Xuan Yang &lt;xygoogle@google.com&gt;
PiperOrigin-RevId: 842321126
diff --git a/src/google/adk/features/_feature_registry.py b/src/google/adk/features/_feature_registry.py
@@ -24,9 +24,16 @@
 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)
diff --git a/src/google/adk/tools/_google_credentials.py b/src/google/adk/tools/_google_credentials.py
@@ -33,11 +33,12 @@
 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).
 
diff --git a/src/google/adk/tools/bigquery/bigquery_credentials.py b/src/google/adk/tools/bigquery/bigquery_credentials.py
@@ -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).
 
diff --git a/src/google/adk/tools/bigquery/bigquery_toolset.py b/src/google/adk/tools/bigquery/bigquery_toolset.py
@@ -24,16 +24,17 @@
 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."""
 
diff --git a/src/google/adk/tools/bigquery/config.py b/src/google/adk/tools/bigquery/config.py
@@ -21,7 +21,8 @@
 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."""
 
diff --git a/src/google/adk/tools/bigtable/bigtable_credentials.py b/src/google/adk/tools/bigtable/bigtable_credentials.py
@@ -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 @@
 ]
 
 
-@experimental
+@experimental(FeatureName.GOOGLE_CREDENTIALS_CONFIG)
 class BigtableCredentialsConfig(BaseGoogleCredentialsConfig):
   """Bigtable Credentials Configuration for Google API tools (Experimental).
 
diff --git a/src/google/adk/tools/bigtable/settings.py b/src/google/adk/tools/bigtable/settings.py
@@ -16,10 +16,11 @@
 
 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."""
 
diff --git a/src/google/adk/tools/google_tool.py b/src/google/adk/tools/google_tool.py
@@ -23,14 +23,15 @@
 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.
 
diff --git a/src/google/adk/tools/spanner/settings.py b/src/google/adk/tools/spanner/settings.py
@@ -22,7 +22,8 @@
 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 @@ def __post_init__(self):
     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."""
 
diff --git a/src/google/adk/tools/spanner/spanner_credentials.py b/src/google/adk/tools/spanner/spanner_credentials.py
@@ -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 @@
 ]
 
 
-@experimental
+@experimental(FeatureName.GOOGLE_CREDENTIALS_CONFIG)
 class SpannerCredentialsConfig(BaseGoogleCredentialsConfig):
   """Spanner Credentials Configuration for Google API tools (Experimental).
 
diff --git a/src/google/adk/tools/spanner/spanner_toolset.py b/src/google/adk/tools/spanner/spanner_toolset.py
@@ -24,19 +24,20 @@
 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
 
 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.
 
diff --git a/tests/unittests/tools/bigquery/test_bigquery_tool_config.py b/tests/unittests/tools/bigquery/test_bigquery_tool_config.py
@@ -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
 
 
diff --git a/tests/unittests/tools/spanner/test_spanner_tool_settings.py b/tests/unittests/tools/spanner/test_spanner_tool_settings.py
@@ -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 @@
 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():
