mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
1ae944b39d
commit
bab57296d5
@@ -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)
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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).
|
||||
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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():
|
||||
|
||||
Reference in New Issue
Block a user