Merge branch 'main' into main

Author: DeanChensj

src/google/adk/cli/cli_deploy.py

@@ -67,6 +67,8 @@
 """
 
 _AGENT_ENGINE_APP_TEMPLATE: Final[str] = """
+import os
+import vertexai
 from vertexai.agent_engines import AdkApp
 
 if {is_config_agent}:
@@ -81,9 +83,12 @@
   from .agent import {adk_app_object}
 
 if {express_mode}: # Whether or not to use Express Mode
-  import os
-  import vertexai
   vertexai.init(api_key=os.environ.get("GOOGLE_API_KEY"))
+else:
+  vertexai.init(
+    project=os.environ.get("GOOGLE_CLOUD_PROJECT"),
+    location=os.environ.get("GOOGLE_CLOUD_LOCATION"),
+  )
 
 adk_app = AdkApp(
     {adk_app_type}={adk_app_object},

src/google/adk/features/__init__.py

@@ -0,0 +1,13 @@
+# 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.

src/google/adk/features/feature_registry.py

@@ -0,0 +1,158 @@
+# 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 __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+import warnings
+
+from ..utils.env_utils import is_env_enabled
+
+
+class FeatureName(str, Enum):
+  """Feature names."""
+
+  JSON_SCHEMA_FOR_FUNC_DECL = "JSON_SCHEMA_FOR_FUNC_DECL"
+  COMPUTER_USE = "COMPUTER_USE"
+
+
+class FeatureStage(Enum):
+  """Feature lifecycle stages.
+
+  Attributes:
+    WIP: Work in progress, not functioning completely. ADK internal development
+      only.
+    EXPERIMENTAL: Feature works but API may change.
+    STABLE: Production-ready, no breaking changes without MAJOR version bump.
+  """
+
+  WIP = "wip"
+  EXPERIMENTAL = "experimental"
+  STABLE = "stable"
+
+
+@dataclass
+class FeatureConfig:
+  """Feature configuration.
+
+  Attributes:
+    stage: The feature stage.
+    default_on: Whether the feature is enabled by default.
+  """
+
+  stage: FeatureStage
+  default_on: bool = False
+
+
+# Central registry: FeatureName -> FeatureConfig
+_FEATURE_REGISTRY: dict[FeatureName, FeatureConfig] = {
+    FeatureName.JSON_SCHEMA_FOR_FUNC_DECL: FeatureConfig(
+        FeatureStage.WIP, default_on=False
+    ),
+    FeatureName.COMPUTER_USE: FeatureConfig(
+        FeatureStage.EXPERIMENTAL, default_on=True
+    ),
+}
+
+# Track which experimental features have already warned (warn only once)
+_WARNED_FEATURES: set[FeatureName] = set()
+
+
+def _get_feature_config(
+    feature_name: FeatureName,
+) -> FeatureConfig | None:
+  """Get the stage of a feature from the registry.
+
+  Args:
+    feature_name: The feature name.
+
+  Returns:
+    The feature config from the registry, or None if not found.
+  """
+  return _FEATURE_REGISTRY.get(feature_name, None)
+
+
+def _register_feature(
+    feature_name: FeatureName,
+    config: FeatureConfig,
+) -> None:
+  """Register a feature with a specific config.
+
+  Args:
+    feature_name: The feature name.
+    config: The feature config to register.
+  """
+  _FEATURE_REGISTRY[feature_name] = config
+
+
+def is_feature_enabled(feature_name: FeatureName) -> bool:
+  """Check if a feature is enabled at runtime.
+
+  This function is used for runtime behavior gating within stable features.
+  It allows you to conditionally enable new behavior based on feature flags.
+
+  Args:
+    feature_name: The feature name (e.g., FeatureName.RESUMABILITY).
+
+  Returns:
+    True if the feature is enabled, False otherwise.
+
+  Example:
+    ```python
+    def _execute_agent_loop():
+      if is_feature_enabled(FeatureName.RESUMABILITY):
+        # New behavior: save checkpoints for resuming
+        return _execute_with_checkpoints()
+      else:
+        # Old behavior: run without checkpointing
+        return _execute_standard()
+    ```
+  """
+  config = _get_feature_config(feature_name)
+  if config is None:
+    raise ValueError(f"Feature {feature_name} is not registered.")
+
+  # Check environment variables first (highest priority)
+  enable_var = f"ADK_ENABLE_{feature_name}"
+  disable_var = f"ADK_DISABLE_{feature_name}"
+  if is_env_enabled(enable_var):
+    if config.stage != FeatureStage.STABLE:
+      _emit_non_stable_warning_once(feature_name, config.stage)
+    return True
+  if is_env_enabled(disable_var):
+    return False
+
+  # Fall back to registry config
+  if config.stage != FeatureStage.STABLE and config.default_on:
+    _emit_non_stable_warning_once(feature_name, config.stage)
+  return config.default_on
+
+
+def _emit_non_stable_warning_once(
+    feature_name: FeatureName,
+    feature_stage: FeatureStage,
+) -> None:
+  """Emit a warning for a non-stable feature, but only once per feature.
+
+  Args:
+    feature_name: The feature name.
+    feature_stage: The feature stage.
+  """
+  if feature_name not in _WARNED_FEATURES:
+    _WARNED_FEATURES.add(feature_name)
+    full_message = (
+        f"[{feature_stage.name.upper()}] feature {feature_name} is enabled."
+    )
+    warnings.warn(full_message, category=UserWarning, stacklevel=4)

src/google/adk/tools/bigquery/config.py

@@ -61,6 +61,14 @@ class BigQueryToolConfig(BaseModel):
   change in future versions.
   """
 
+  maximum_bytes_billed: Optional[int] = None
+  """Maximum number of bytes to bill for a query.
+
+  In BigQuery on-demand pricing, charges are rounded up to the nearest MB, with
+  a minimum 10 MB data processed per table referenced by the query, and with a
+  minimum 10 MB data processed per query. So this value must be set >=10485760.
+  """
+
   max_query_result_rows: int = 50
   """Maximum number of rows to return from a query.
 
@@ -91,6 +99,19 @@ class BigQueryToolConfig(BaseModel):
   locations, see https://cloud.google.com/bigquery/docs/locations.
   """
 
+  @field_validator('maximum_bytes_billed')
+  @classmethod
+  def validate_maximum_bytes_billed(cls, v):
+    """Validate the maximum bytes billed."""
+    if v and v < 10_485_760:
+      raise ValueError(
+          'In BigQuery on-demand pricing, charges are rounded up to the nearest'
+          ' MB, with a minimum 10 MB data processed per table referenced by the'
+          ' query, and with a minimum 10 MB data processed per query. So'
+          ' max_bytes_billed must be set >=10485760.'
+      )
+    return v
+
   @field_validator('application_name')
   @classmethod
   def validate_application_name(cls, v):

src/google/adk/tools/bigquery/query_tool.py

@@ -152,12 +152,15 @@ def _execute_sql(
       return {"status": "SUCCESS", "dry_run_info": dry_run_job.to_api_repr()}
 
     # Finally execute the query, fetch the result, and return it
+    job_config = bigquery.QueryJobConfig(
+        connection_properties=bq_connection_properties,
+        labels=bq_job_labels,
+    )
+    if settings.maximum_bytes_billed:
+      job_config.maximum_bytes_billed = settings.maximum_bytes_billed
     row_iterator = bq_client.query_and_wait(
         query,
-        job_config=bigquery.QueryJobConfig(
-            connection_properties=bq_connection_properties,
-            labels=bq_job_labels,
-        ),
+        job_config=job_config,
         project=project_id,
         max_results=settings.max_query_result_rows,
     )
@@ -1090,21 +1093,23 @@ def analyze_contribution(
   """
 
   # Create a session and run the create model query.
-  original_write_mode = settings.write_mode
   try:
-    if original_write_mode == WriteMode.BLOCKED:
+    execute_sql_settings = settings
+    if execute_sql_settings.write_mode == WriteMode.BLOCKED:
       raise ValueError("analyze_contribution is not allowed in this session.")
-    elif original_write_mode != WriteMode.PROTECTED:
+    elif execute_sql_settings.write_mode != WriteMode.PROTECTED:
       # Running create temp model requires a session. So we set the write mode
       # to PROTECTED to run the create model query and job query in the same
       # session.
-      settings.write_mode = WriteMode.PROTECTED
+      execute_sql_settings = settings.model_copy(
+          update={"write_mode": WriteMode.PROTECTED}
+      )
 
     result = _execute_sql(
         project_id=project_id,
         query=create_model_query,
         credentials=credentials,
-        settings=settings,
+        settings=execute_sql_settings,
         tool_context=tool_context,
         caller_id="analyze_contribution",
     )
@@ -1115,18 +1120,15 @@ def analyze_contribution(
         project_id=project_id,
         query=get_insights_query,
         credentials=credentials,
-        settings=settings,
+        settings=execute_sql_settings,
         tool_context=tool_context,
         caller_id="analyze_contribution",
     )
   except Exception as ex:  # pylint: disable=broad-except
     return {
         "status": "ERROR",
-        "error_details": f"Error during analyze_contribution: {str(ex)}",
+        "error_details": f"Error during analyze_contribution: {repr(ex)}",
     }
-  finally:
-    # Restore the original write mode.
-    settings.write_mode == original_write_mode
 
   return result
 
@@ -1324,21 +1326,23 @@ def detect_anomalies(
     """
 
   # Create a session and run the create model query.
-  original_write_mode = settings.write_mode
   try:
-    if settings.write_mode == WriteMode.BLOCKED:
+    execute_sql_settings = settings
+    if execute_sql_settings.write_mode == WriteMode.BLOCKED:
       raise ValueError("anomaly detection is not allowed in this session.")
-    elif original_write_mode != WriteMode.PROTECTED:
+    elif execute_sql_settings.write_mode != WriteMode.PROTECTED:
       # Running create temp model requires a session. So we set the write mode
       # to PROTECTED to run the create model query and job query in the same
       # session.
-      settings.write_mode = WriteMode.PROTECTED
+      execute_sql_settings = settings.model_copy(
+          update={"write_mode": WriteMode.PROTECTED}
+      )
 
     result = _execute_sql(
         project_id=project_id,
         query=create_model_query,
         credentials=credentials,
-        settings=settings,
+        settings=execute_sql_settings,
         tool_context=tool_context,
         caller_id="detect_anomalies",
     )
@@ -1349,17 +1353,14 @@ def detect_anomalies(
         project_id=project_id,
         query=anomaly_detection_query,
         credentials=credentials,
-        settings=settings,
+        settings=execute_sql_settings,
         tool_context=tool_context,
         caller_id="detect_anomalies",
     )
   except Exception as ex:  # pylint: disable=broad-except
     return {
         "status": "ERROR",
-        "error_details": f"Error during anomaly detection: {str(ex)}",
+        "error_details": f"Error during anomaly detection: {repr(ex)}",
     }
-  finally:
-    # Restore the original write mode.
-    settings.write_mode == original_write_mode
 
   return result