Merge branch 'main' into main

Author: theyashl

src/google/adk/cli/cli_deploy.py

@@ -689,6 +689,7 @@ def to_agent_engine(
     adk_app: str,
     staging_bucket: Optional[str] = None,
     trace_to_cloud: Optional[bool] = None,
+    otel_to_cloud: Optional[bool] = None,
     api_key: Optional[str] = None,
     adk_app_object: Optional[str] = None,
     agent_engine_id: Optional[str] = None,
@@ -733,6 +734,8 @@ def to_agent_engine(
     staging_bucket (str): Deprecated. This argument is no longer required or
       used.
     trace_to_cloud (bool): Whether to enable Cloud Trace.
+    otel_to_cloud (bool): Whether to enable exporting OpenTelemetry signals
+      to Google Cloud.
     api_key (str): Optional. The API key to use for Express Mode.
       If not provided, the API key from the GOOGLE_API_KEY environment variable
       will be used. It will only be used if GOOGLE_GENAI_USE_VERTEXAI is true.
@@ -910,6 +913,14 @@ def to_agent_engine(
       if 'GOOGLE_API_KEY' in env_vars:
         api_key = env_vars['GOOGLE_API_KEY']
         click.echo(f'api_key set by GOOGLE_API_KEY in {env_file}')
+    if otel_to_cloud:
+      if 'GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY' in env_vars:
+        click.secho(
+            'Ignoring GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in .env'
+            ' as `--otel_to_cloud` was explicitly passed and takes precedence',
+            fg='yellow',
+        )
+      env_vars['GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY'] = 'true'
     if env_vars:
       if 'env_vars' in agent_config:
         click.echo(

src/google/adk/cli/cli_tools_click.py

@@ -50,28 +50,44 @@
 )
 
 
-def _apply_feature_overrides(enable_features: tuple[str, ...]) -> None:
+def _apply_feature_overrides(
+    *,
+    enable_features: tuple[str, ...] = (),
+    disable_features: tuple[str, ...] = (),
+) -> None:
   """Apply feature overrides from CLI flags.
 
   Args:
     enable_features: Tuple of feature names to enable.
+    disable_features: Tuple of feature names to disable.
   """
+  feature_overrides: dict[str, bool] = {}
+
   for features_str in enable_features:
     for feature_name_str in features_str.split(","):
       feature_name_str = feature_name_str.strip()
-      if not feature_name_str:
-        continue
-      try:
-        feature_name = FeatureName(feature_name_str)
-        override_feature_enabled(feature_name, True)
-      except ValueError:
-        valid_names = ", ".join(f.value for f in FeatureName)
-        click.secho(
-            f"WARNING: Unknown feature name '{feature_name_str}'. "
-            f"Valid names are: {valid_names}",
-            fg="yellow",
-            err=True,
-        )
+      if feature_name_str:
+        feature_overrides[feature_name_str] = True
+
+  for features_str in disable_features:
+    for feature_name_str in features_str.split(","):
+      feature_name_str = feature_name_str.strip()
+      if feature_name_str:
+        feature_overrides[feature_name_str] = False
+
+  # Apply all overrides
+  for feature_name_str, enabled in feature_overrides.items():
+    try:
+      feature_name = FeatureName(feature_name_str)
+      override_feature_enabled(feature_name, enabled)
+    except ValueError:
+      valid_names = ", ".join(f.value for f in FeatureName)
+      click.secho(
+          f"WARNING: Unknown feature name '{feature_name_str}'. "
+          f"Valid names are: {valid_names}",
+          fg="yellow",
+          err=True,
+      )
 
 
 def feature_options():
@@ -88,11 +104,25 @@ def decorator(func):
         ),
         multiple=True,
     )
+    @click.option(
+        "--disable_features",
+        help=(
+            "Optional. Comma-separated list of feature names to disable. "
+            "This provides an alternative to environment variables for "
+            "disabling features. Example: "
+            "--disable_features=JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING"
+        ),
+        multiple=True,
+    )
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
       enable_features = kwargs.pop("enable_features", ())
-      if enable_features:
-        _apply_feature_overrides(enable_features)
+      disable_features = kwargs.pop("disable_features", ())
+      if enable_features or disable_features:
+        _apply_feature_overrides(
+            enable_features=enable_features,
+            disable_features=disable_features,
+        )
       return func(*args, **kwargs)
 
     return wrapper
@@ -1751,6 +1781,14 @@ def cli_migrate_session(
     default=None,
     help="Optional. Whether to enable Cloud Trace for Agent Engine.",
 )
+@click.option(
+    "--otel_to_cloud",
+    type=bool,
+    is_flag=True,
+    show_default=True,
+    default=None,
+    help="Optional. Whether to enable OpenTelemetry for Agent Engine.",
+)
 @click.option(
     "--display_name",
     type=str,
@@ -1842,6 +1880,7 @@ def cli_deploy_agent_engine(
     staging_bucket: Optional[str],
     agent_engine_id: Optional[str],
     trace_to_cloud: Optional[bool],
+    otel_to_cloud: Optional[bool],
     api_key: Optional[str],
     display_name: str,
     description: str,
@@ -1872,6 +1911,7 @@ def cli_deploy_agent_engine(
         region=region,
         agent_engine_id=agent_engine_id,
         trace_to_cloud=trace_to_cloud,
+        otel_to_cloud=otel_to_cloud,
         api_key=api_key,
         adk_app_object=adk_app_object,
         display_name=display_name,

src/google/adk/sessions/migration/_schema_check_utils.py

@@ -82,8 +82,28 @@ def get_db_schema_version_from_connection(connection) -> str:
   return _get_schema_version_impl(inspector, connection)
 
 
-def _to_sync_url(db_url: str) -> str:
-  """Removes '+driver' from SQLAlchemy URL."""
+def to_sync_url(db_url: str) -> str:
+  """Removes '+driver' from SQLAlchemy URL.
+
+  This is useful when you need to use a synchronous SQLAlchemy engine with
+  a database URL that specifies an async driver (e.g., postgresql+asyncpg://
+  or sqlite+aiosqlite://).
+
+  Args:
+    db_url: The database URL, potentially with a driver specification.
+
+  Returns:
+    The database URL with the driver specification removed (e.g.,
+    'postgresql+asyncpg://host/db' becomes 'postgresql://host/db').
+
+  Examples:
+    >>> to_sync_url('postgresql+asyncpg://localhost/mydb')
+    'postgresql://localhost/mydb'
+    >>> to_sync_url('sqlite+aiosqlite:///path/to/db.sqlite')
+    'sqlite:///path/to/db.sqlite'
+    >>> to_sync_url('mysql://localhost/mydb')  # No driver, returns unchanged
+    'mysql://localhost/mydb'
+  """
   if "://" in db_url:
     scheme, _, rest = db_url.partition("://")
     if "+" in scheme:
@@ -106,7 +126,7 @@ def get_db_schema_version(db_url: str) -> str:
   """
   engine = None
   try:
-    engine = create_sync_engine(_to_sync_url(db_url))
+    engine = create_sync_engine(to_sync_url(db_url))
     with engine.connect() as connection:
       inspector = inspect(connection)
       return _get_schema_version_impl(inspector, connection)

src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py

@@ -165,17 +165,23 @@ def _get_state_dict(state_val: Any) -> dict:
 # --- Migration Logic ---
 def migrate(source_db_url: str, dest_db_url: str):
   """Migrates data from old pickle schema to new JSON schema."""
+  # Convert async driver URLs to sync URLs for SQLAlchemy's synchronous engine.
+  # This allows users to provide URLs like 'postgresql+asyncpg://...' and have
+  # them automatically converted to 'postgresql://...' for migration.
+  source_sync_url = _schema_check_utils.to_sync_url(source_db_url)
+  dest_sync_url = _schema_check_utils.to_sync_url(dest_db_url)
+
   logger.info(f"Connecting to source database: {source_db_url}")
   try:
-    source_engine = create_engine(source_db_url)
+    source_engine = create_engine(source_sync_url)
     SourceSession = sessionmaker(bind=source_engine)
   except Exception as e:
     logger.error(f"Failed to connect to source database: {e}")
     raise RuntimeError(f"Failed to connect to source database: {e}") from e
 
   logger.info(f"Connecting to destination database: {dest_db_url}")
   try:
-    dest_engine = create_engine(dest_db_url)
+    dest_engine = create_engine(dest_sync_url)
     v1.Base.metadata.create_all(dest_engine)
     DestSession = sessionmaker(bind=dest_engine)
   except Exception as e:

src/google/adk/sessions/migration/migrate_from_sqlalchemy_sqlite.py

@@ -23,6 +23,7 @@
 import sys
 
 from google.adk.sessions import sqlite_session_service as sss
+from google.adk.sessions.migration import _schema_check_utils
 from google.adk.sessions.schemas import v0 as v0_schema
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
@@ -32,9 +33,14 @@
 
 def migrate(source_db_url: str, dest_db_path: str):
   """Migrates data from a SQLAlchemy-based SQLite DB to the new schema."""
+  # Convert async driver URLs to sync URLs for SQLAlchemy's synchronous engine.
+  # This allows users to provide URLs like 'sqlite+aiosqlite://...' and have
+  # them automatically converted to 'sqlite://...' for migration.
+  source_sync_url = _schema_check_utils.to_sync_url(source_db_url)
+
   logger.info(f"Connecting to source database: {source_db_url}")
   try:
-    engine = create_engine(source_db_url)
+    engine = create_engine(source_sync_url)
     v0_schema.Base.metadata.create_all(
         engine
     )  # Ensure tables exist for inspection

src/google/adk/tools/load_artifacts_tool.py

@@ -24,6 +24,8 @@
 from google.genai import types
 from typing_extensions import override
 
+from ..features import FeatureName
+from ..features import is_feature_enabled
 from .base_tool import BaseTool
 
 # MIME types Gemini accepts for inline data in requests.
@@ -132,6 +134,20 @@ def __init__(self):
     )
 
   def _get_declaration(self) -> types.FunctionDeclaration | None:
+    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
+      return types.FunctionDeclaration(
+          name=self.name,
+          description=self.description,
+          parameters_json_schema={
+              'type': 'object',
+              'properties': {
+                  'artifact_names': {
+                      'type': 'array',
+                      'items': {'type': 'string'},
+                  },
+              },
+          },
+      )
     return types.FunctionDeclaration(
         name=self.name,
         description=self.description,

src/google/adk/tools/load_memory_tool.py

@@ -21,6 +21,8 @@
 from pydantic import Field
 from typing_extensions import override
 
+from ..features import FeatureName
+from ..features import is_feature_enabled
 from ..memory.memory_entry import MemoryEntry
 from .function_tool import FunctionTool
 from .tool_context import ToolContext
@@ -59,6 +61,18 @@ def __init__(self):
 
   @override
   def _get_declaration(self) -> types.FunctionDeclaration | None:
+    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
+      return types.FunctionDeclaration(
+          name=self.name,
+          description=self.description,
+          parameters_json_schema={
+              'type': 'object',
+              'properties': {
+                  'query': {'type': 'string'},
+              },
+              'required': ['query'],
+          },
+      )
     return types.FunctionDeclaration(
         name=self.name,
         description=self.description,

src/google/adk/tools/retrieval/base_retrieval_tool.py

@@ -12,16 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from google.genai import types
 from typing_extensions import override
 
+from ...features import FeatureName
+from ...features import is_feature_enabled
 from ..base_tool import BaseTool
 
 
 class BaseRetrievalTool(BaseTool):
 
   @override
   def _get_declaration(self) -> types.FunctionDeclaration:
+    if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
+      return types.FunctionDeclaration(
+          name=self.name,
+          description=self.description,
+          parameters_json_schema={
+              'type': 'object',
+              'properties': {
+                  'query': {
+                      'type': 'string',
+                      'description': 'The query to retrieve.',
+                  },
+              },
+          },
+      )
     return types.FunctionDeclaration(
         name=self.name,
         description=self.description,