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):
|
class FeatureName(str, Enum):
|
||||||
"""Feature names."""
|
"""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"
|
COMPUTER_USE = "COMPUTER_USE"
|
||||||
|
GOOGLE_CREDENTIALS_CONFIG = "GOOGLE_CREDENTIALS_CONFIG"
|
||||||
|
GOOGLE_TOOL = "GOOGLE_TOOL"
|
||||||
JSON_SCHEMA_FOR_FUNC_DECL = "JSON_SCHEMA_FOR_FUNC_DECL"
|
JSON_SCHEMA_FOR_FUNC_DECL = "JSON_SCHEMA_FOR_FUNC_DECL"
|
||||||
PROGRESSIVE_SSE_STREAMING = "PROGRESSIVE_SSE_STREAMING"
|
PROGRESSIVE_SSE_STREAMING = "PROGRESSIVE_SSE_STREAMING"
|
||||||
|
SPANNER_TOOLSET = "SPANNER_TOOLSET"
|
||||||
|
SPANNER_TOOL_SETTINGS = "SPANNER_TOOL_SETTINGS"
|
||||||
|
|
||||||
|
|
||||||
class FeatureStage(Enum):
|
class FeatureStage(Enum):
|
||||||
@@ -59,15 +66,36 @@ class FeatureConfig:
|
|||||||
|
|
||||||
# Central registry: FeatureName -> FeatureConfig
|
# Central registry: FeatureName -> FeatureConfig
|
||||||
_FEATURE_REGISTRY: dict[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(
|
FeatureName.COMPUTER_USE: FeatureConfig(
|
||||||
FeatureStage.EXPERIMENTAL, default_on=True
|
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(
|
FeatureName.JSON_SCHEMA_FOR_FUNC_DECL: FeatureConfig(
|
||||||
FeatureStage.WIP, default_on=False
|
FeatureStage.WIP, default_on=False
|
||||||
),
|
),
|
||||||
FeatureName.PROGRESSIVE_SSE_STREAMING: FeatureConfig(
|
FeatureName.PROGRESSIVE_SSE_STREAMING: FeatureConfig(
|
||||||
FeatureStage.WIP, default_on=False
|
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)
|
# 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 AuthCredentialTypes
|
||||||
from ..auth.auth_credential import OAuth2Auth
|
from ..auth.auth_credential import OAuth2Auth
|
||||||
from ..auth.auth_tool import AuthConfig
|
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
|
from .tool_context import ToolContext
|
||||||
|
|
||||||
|
|
||||||
@experimental
|
@experimental(FeatureName.GOOGLE_CREDENTIALS_CONFIG)
|
||||||
class BaseGoogleCredentialsConfig(BaseModel):
|
class BaseGoogleCredentialsConfig(BaseModel):
|
||||||
"""Base Google Credentials Configuration for Google API tools (Experimental).
|
"""Base Google Credentials Configuration for Google API tools (Experimental).
|
||||||
|
|
||||||
|
|||||||
@@ -14,14 +14,15 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from ...utils.feature_decorator import experimental
|
from ...features import experimental
|
||||||
|
from ...features import FeatureName
|
||||||
from .._google_credentials import BaseGoogleCredentialsConfig
|
from .._google_credentials import BaseGoogleCredentialsConfig
|
||||||
|
|
||||||
BIGQUERY_TOKEN_CACHE_KEY = "bigquery_token_cache"
|
BIGQUERY_TOKEN_CACHE_KEY = "bigquery_token_cache"
|
||||||
BIGQUERY_DEFAULT_SCOPE = ["https://www.googleapis.com/auth/bigquery"]
|
BIGQUERY_DEFAULT_SCOPE = ["https://www.googleapis.com/auth/bigquery"]
|
||||||
|
|
||||||
|
|
||||||
@experimental
|
@experimental(FeatureName.GOOGLE_CREDENTIALS_CONFIG)
|
||||||
class BigQueryCredentialsConfig(BaseGoogleCredentialsConfig):
|
class BigQueryCredentialsConfig(BaseGoogleCredentialsConfig):
|
||||||
"""BigQuery Credentials Configuration for Google API tools (Experimental).
|
"""BigQuery Credentials Configuration for Google API tools (Experimental).
|
||||||
|
|
||||||
|
|||||||
@@ -24,16 +24,17 @@ from typing_extensions import override
|
|||||||
from . import data_insights_tool
|
from . import data_insights_tool
|
||||||
from . import metadata_tool
|
from . import metadata_tool
|
||||||
from . import query_tool
|
from . import query_tool
|
||||||
|
from ...features import experimental
|
||||||
|
from ...features import FeatureName
|
||||||
from ...tools.base_tool import BaseTool
|
from ...tools.base_tool import BaseTool
|
||||||
from ...tools.base_toolset import BaseToolset
|
from ...tools.base_toolset import BaseToolset
|
||||||
from ...tools.base_toolset import ToolPredicate
|
from ...tools.base_toolset import ToolPredicate
|
||||||
from ...tools.google_tool import GoogleTool
|
from ...tools.google_tool import GoogleTool
|
||||||
from ...utils.feature_decorator import experimental
|
|
||||||
from .bigquery_credentials import BigQueryCredentialsConfig
|
from .bigquery_credentials import BigQueryCredentialsConfig
|
||||||
from .config import BigQueryToolConfig
|
from .config import BigQueryToolConfig
|
||||||
|
|
||||||
|
|
||||||
@experimental
|
@experimental(FeatureName.BIG_QUERY_TOOLSET)
|
||||||
class BigQueryToolset(BaseToolset):
|
class BigQueryToolset(BaseToolset):
|
||||||
"""BigQuery Toolset contains tools for interacting with BigQuery data and metadata."""
|
"""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 ConfigDict
|
||||||
from pydantic import field_validator
|
from pydantic import field_validator
|
||||||
|
|
||||||
from ...utils.feature_decorator import experimental
|
from ...features import experimental
|
||||||
|
from ...features import FeatureName
|
||||||
|
|
||||||
|
|
||||||
class WriteMode(Enum):
|
class WriteMode(Enum):
|
||||||
@@ -47,7 +48,7 @@ class WriteMode(Enum):
|
|||||||
"""All write operations are allowed."""
|
"""All write operations are allowed."""
|
||||||
|
|
||||||
|
|
||||||
@experimental('Config defaults may have breaking change in the future.')
|
@experimental(FeatureName.BIG_QUERY_TOOL_CONFIG)
|
||||||
class BigQueryToolConfig(BaseModel):
|
class BigQueryToolConfig(BaseModel):
|
||||||
"""Configuration for BigQuery tools."""
|
"""Configuration for BigQuery tools."""
|
||||||
|
|
||||||
|
|||||||
@@ -14,7 +14,8 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from ...utils.feature_decorator import experimental
|
from ...features import experimental
|
||||||
|
from ...features import FeatureName
|
||||||
from .._google_credentials import BaseGoogleCredentialsConfig
|
from .._google_credentials import BaseGoogleCredentialsConfig
|
||||||
|
|
||||||
BIGTABLE_TOKEN_CACHE_KEY = "bigtable_token_cache"
|
BIGTABLE_TOKEN_CACHE_KEY = "bigtable_token_cache"
|
||||||
@@ -24,7 +25,7 @@ BIGTABLE_DEFAULT_SCOPE = [
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
@experimental
|
@experimental(FeatureName.GOOGLE_CREDENTIALS_CONFIG)
|
||||||
class BigtableCredentialsConfig(BaseGoogleCredentialsConfig):
|
class BigtableCredentialsConfig(BaseGoogleCredentialsConfig):
|
||||||
"""Bigtable Credentials Configuration for Google API tools (Experimental).
|
"""Bigtable Credentials Configuration for Google API tools (Experimental).
|
||||||
|
|
||||||
|
|||||||
@@ -16,10 +16,11 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from pydantic import BaseModel
|
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):
|
class BigtableToolSettings(BaseModel):
|
||||||
"""Settings for Bigtable tools."""
|
"""Settings for Bigtable tools."""
|
||||||
|
|
||||||
|
|||||||
@@ -23,14 +23,15 @@ from google.auth.credentials import Credentials
|
|||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from typing_extensions import override
|
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 BaseGoogleCredentialsConfig
|
||||||
from ._google_credentials import GoogleCredentialsManager
|
from ._google_credentials import GoogleCredentialsManager
|
||||||
from .function_tool import FunctionTool
|
from .function_tool import FunctionTool
|
||||||
from .tool_context import ToolContext
|
from .tool_context import ToolContext
|
||||||
|
|
||||||
|
|
||||||
@experimental
|
@experimental(FeatureName.GOOGLE_TOOL)
|
||||||
class GoogleTool(FunctionTool):
|
class GoogleTool(FunctionTool):
|
||||||
"""GoogleTool class for tools that call Google APIs.
|
"""GoogleTool class for tools that call Google APIs.
|
||||||
|
|
||||||
|
|||||||
@@ -22,7 +22,8 @@ from typing import Optional
|
|||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from pydantic import model_validator
|
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.
|
# Vector similarity search nearest neighbors search algorithms.
|
||||||
EXACT_NEAREST_NEIGHBORS = "EXACT_NEAREST_NEIGHBORS"
|
EXACT_NEAREST_NEIGHBORS = "EXACT_NEAREST_NEIGHBORS"
|
||||||
@@ -138,7 +139,7 @@ class SpannerVectorStoreSettings(BaseModel):
|
|||||||
return self
|
return self
|
||||||
|
|
||||||
|
|
||||||
@experimental("Tool settings defaults may have breaking change in the future.")
|
@experimental(FeatureName.SPANNER_TOOL_SETTINGS)
|
||||||
class SpannerToolSettings(BaseModel):
|
class SpannerToolSettings(BaseModel):
|
||||||
"""Settings for Spanner tools."""
|
"""Settings for Spanner tools."""
|
||||||
|
|
||||||
|
|||||||
@@ -14,7 +14,8 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from ...utils.feature_decorator import experimental
|
from ...features import experimental
|
||||||
|
from ...features import FeatureName
|
||||||
from .._google_credentials import BaseGoogleCredentialsConfig
|
from .._google_credentials import BaseGoogleCredentialsConfig
|
||||||
|
|
||||||
SPANNER_TOKEN_CACHE_KEY = "spanner_token_cache"
|
SPANNER_TOKEN_CACHE_KEY = "spanner_token_cache"
|
||||||
@@ -24,7 +25,7 @@ SPANNER_DEFAULT_SCOPE = [
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
@experimental
|
@experimental(FeatureName.GOOGLE_CREDENTIALS_CONFIG)
|
||||||
class SpannerCredentialsConfig(BaseGoogleCredentialsConfig):
|
class SpannerCredentialsConfig(BaseGoogleCredentialsConfig):
|
||||||
"""Spanner Credentials Configuration for Google API tools (Experimental).
|
"""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 google.adk.tools.spanner import search_tool
|
||||||
from typing_extensions import override
|
from typing_extensions import override
|
||||||
|
|
||||||
|
from ...features import experimental
|
||||||
|
from ...features import FeatureName
|
||||||
from ...tools.base_tool import BaseTool
|
from ...tools.base_tool import BaseTool
|
||||||
from ...tools.base_toolset import BaseToolset
|
from ...tools.base_toolset import BaseToolset
|
||||||
from ...tools.base_toolset import ToolPredicate
|
from ...tools.base_toolset import ToolPredicate
|
||||||
from ...tools.google_tool import GoogleTool
|
from ...tools.google_tool import GoogleTool
|
||||||
from ...utils.feature_decorator import experimental
|
|
||||||
from .settings import Capabilities
|
from .settings import Capabilities
|
||||||
from .settings import SpannerToolSettings
|
from .settings import SpannerToolSettings
|
||||||
from .spanner_credentials import SpannerCredentialsConfig
|
from .spanner_credentials import SpannerCredentialsConfig
|
||||||
@@ -36,7 +37,7 @@ from .spanner_credentials import SpannerCredentialsConfig
|
|||||||
DEFAULT_SPANNER_TOOL_NAME_PREFIX = "spanner"
|
DEFAULT_SPANNER_TOOL_NAME_PREFIX = "spanner"
|
||||||
|
|
||||||
|
|
||||||
@experimental
|
@experimental(FeatureName.SPANNER_TOOLSET)
|
||||||
class SpannerToolset(BaseToolset):
|
class SpannerToolset(BaseToolset):
|
||||||
"""Spanner Toolset contains tools for interacting with Spanner data, database and table information.
|
"""Spanner Toolset contains tools for interacting with Spanner data, database and table information.
|
||||||
|
|
||||||
|
|||||||
@@ -14,17 +14,25 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import warnings
|
||||||
|
|
||||||
|
from google.adk.features._feature_registry import _WARNED_FEATURES
|
||||||
from google.adk.tools.bigquery.config import BigQueryToolConfig
|
from google.adk.tools.bigquery.config import BigQueryToolConfig
|
||||||
import pytest
|
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():
|
def test_bigquery_tool_config_experimental_warning():
|
||||||
"""Test BigQueryToolConfig experimental warning."""
|
"""Test BigQueryToolConfig experimental warning."""
|
||||||
with pytest.warns(
|
with warnings.catch_warnings(record=True) as w:
|
||||||
UserWarning,
|
|
||||||
match="Config defaults may have breaking change in the future.",
|
|
||||||
):
|
|
||||||
BigQueryToolConfig()
|
BigQueryToolConfig()
|
||||||
|
assert len(w) == 1
|
||||||
|
assert "BIG_QUERY_TOOL_CONFIG is enabled." in str(w[0].message)
|
||||||
|
|
||||||
|
|
||||||
def test_bigquery_tool_config_invalid_property():
|
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():
|
def test_bigquery_tool_config_max_query_result_rows_default():
|
||||||
"""Test BigQueryToolConfig max_query_result_rows default value."""
|
"""Test BigQueryToolConfig max_query_result_rows default value."""
|
||||||
with pytest.warns(UserWarning):
|
config = BigQueryToolConfig()
|
||||||
config = BigQueryToolConfig()
|
|
||||||
assert config.max_query_result_rows == 50
|
assert config.max_query_result_rows == 50
|
||||||
|
|
||||||
|
|
||||||
def test_bigquery_tool_config_max_query_result_rows_custom():
|
def test_bigquery_tool_config_max_query_result_rows_custom():
|
||||||
"""Test BigQueryToolConfig max_query_result_rows custom value."""
|
"""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
|
assert config.max_query_result_rows == 100
|
||||||
|
|
||||||
|
|
||||||
def test_bigquery_tool_config_valid_maximum_bytes_billed():
|
def test_bigquery_tool_config_valid_maximum_bytes_billed():
|
||||||
"""Test BigQueryToolConfig raises exception with valid max 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
|
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):
|
def test_bigquery_tool_config_valid_labels(labels):
|
||||||
"""Test BigQueryToolConfig accepts valid 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
|
assert config.job_labels == labels
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -14,6 +14,9 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
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 Capabilities
|
||||||
from google.adk.tools.spanner.settings import QueryResultMode
|
from google.adk.tools.spanner.settings import QueryResultMode
|
||||||
from google.adk.tools.spanner.settings import SpannerToolSettings
|
from google.adk.tools.spanner.settings import SpannerToolSettings
|
||||||
@@ -22,6 +25,12 @@ from pydantic import ValidationError
|
|||||||
import pytest
|
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):
|
def common_spanner_vector_store_settings(vector_length=None):
|
||||||
return {
|
return {
|
||||||
"project_id": "test-project",
|
"project_id": "test-project",
|
||||||
@@ -36,11 +45,10 @@ def common_spanner_vector_store_settings(vector_length=None):
|
|||||||
|
|
||||||
def test_spanner_tool_settings_experimental_warning():
|
def test_spanner_tool_settings_experimental_warning():
|
||||||
"""Test SpannerToolSettings experimental warning."""
|
"""Test SpannerToolSettings experimental warning."""
|
||||||
with pytest.warns(
|
with warnings.catch_warnings(record=True) as w:
|
||||||
UserWarning,
|
|
||||||
match="Tool settings defaults may have breaking change in the future.",
|
|
||||||
):
|
|
||||||
SpannerToolSettings()
|
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():
|
def test_spanner_vector_store_settings_all_fields_present():
|
||||||
|
|||||||
Reference in New Issue
Block a user