feat(utils): add GOOGLE_GENAI_USE_ENTERPRISE env var with deprecation fallback

Co-authored-by: Haran Rajkumar <haranrk@google.com>
PiperOrigin-RevId: 929280981

Author: haranrk

src/google/adk/utils/env_utils.py

@@ -21,6 +21,7 @@
 from __future__ import annotations
 
 import os
+import warnings
 
 
 def is_env_enabled(env_var_name: str, default: str = '0') -> bool:
@@ -57,3 +58,22 @@ def is_env_enabled(env_var_name: str, default: str = '0') -> bool:
     True
   """
   return os.environ.get(env_var_name, default).lower() in ['true', '1']
+
+
+def is_enterprise_mode_enabled() -> bool:
+  """Check if Google GenAI Enterprise mode is enabled via environment variables.
+
+  Returns:
+    True if enabled, False otherwise.
+  """
+  if 'GOOGLE_GENAI_USE_ENTERPRISE' in os.environ:
+    return is_env_enabled('GOOGLE_GENAI_USE_ENTERPRISE')
+  if 'GOOGLE_GENAI_USE_VERTEXAI' in os.environ:
+    warnings.warn(
+        'GOOGLE_GENAI_USE_VERTEXAI is deprecated, please use'
+        ' GOOGLE_GENAI_USE_ENTERPRISE instead',
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return is_env_enabled('GOOGLE_GENAI_USE_VERTEXAI')
+  return False

src/google/adk/utils/variant_utils.py

@@ -22,7 +22,7 @@
 
 from enum import Enum
 
-from .env_utils import is_env_enabled
+from .env_utils import is_enterprise_mode_enabled
 
 _GOOGLE_LLM_VARIANT_VERTEX_AI = 'VERTEX_AI'
 _GOOGLE_LLM_VARIANT_GEMINI_API = 'GEMINI_API'
@@ -41,8 +41,6 @@ class GoogleLLMVariant(Enum):
 
 
 def get_google_llm_variant() -> GoogleLLMVariant:
-  return (
-      GoogleLLMVariant.VERTEX_AI
-      if is_env_enabled('GOOGLE_GENAI_USE_VERTEXAI')
-      else GoogleLLMVariant.GEMINI_API
-  )
+  if is_enterprise_mode_enabled():
+    return GoogleLLMVariant.VERTEX_AI
+  return GoogleLLMVariant.GEMINI_API

src/google/adk/utils/vertex_ai_utils.py

@@ -23,7 +23,7 @@
 import os
 from typing import Optional
 
-from ..utils.env_utils import is_env_enabled
+from .env_utils import is_enterprise_mode_enabled
 
 
 def get_express_mode_api_key(
@@ -37,7 +37,7 @@ def get_express_mode_api_key(
         'Cannot specify project or location and express_mode_api_key. '
         'Either use project and location, or just the express_mode_api_key.'
     )
-  if is_env_enabled('GOOGLE_GENAI_USE_VERTEXAI'):
+  if is_enterprise_mode_enabled():
     return express_mode_api_key or os.environ.get('GOOGLE_API_KEY', None)
   else:
     return None

tests/unittests/utils/test_env_utils.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
+
+from google.adk.utils.env_utils import is_enterprise_mode_enabled
 from google.adk.utils.env_utils import is_env_enabled
 import pytest
 
@@ -47,3 +50,35 @@ def test_is_env_enabled_with_defaults(monkeypatch, default, expected):
   """Test is_env_enabled when env var is not set with different defaults."""
   monkeypatch.delenv('TEST_FLAG', raising=False)
   assert is_env_enabled('TEST_FLAG', default=default) is expected
+
+
+def test_is_enterprise_mode_enabled_via_enterprise_env(monkeypatch):
+  """Enterprise mode is on when GOOGLE_GENAI_USE_ENTERPRISE is truthy."""
+  monkeypatch.setenv('GOOGLE_GENAI_USE_ENTERPRISE', 'true')
+
+  assert is_enterprise_mode_enabled() is True
+
+
+def test_is_enterprise_mode_enabled_falls_back_to_vertexai_with_warning(
+    monkeypatch,
+):
+  """The deprecated GOOGLE_GENAI_USE_VERTEXAI still enables enterprise mode and warns."""
+  monkeypatch.delenv('GOOGLE_GENAI_USE_ENTERPRISE', raising=False)
+  monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
+
+  with warnings.catch_warnings(record=True) as caught:
+    warnings.simplefilter('always')
+    result = is_enterprise_mode_enabled()
+
+  assert result is True
+  assert len(caught) == 1
+  assert issubclass(caught[-1].category, DeprecationWarning)
+  assert 'GOOGLE_GENAI_USE_VERTEXAI is deprecated' in str(caught[-1].message)
+
+
+def test_is_enterprise_mode_enabled_defaults_to_false(monkeypatch):
+  """Enterprise mode is off when no relevant env var is set."""
+  monkeypatch.delenv('GOOGLE_GENAI_USE_ENTERPRISE', raising=False)
+  monkeypatch.delenv('GOOGLE_GENAI_USE_VERTEXAI', raising=False)
+
+  assert is_enterprise_mode_enabled() is False

tests/unittests/utils/test_output_schema_utils.py

@@ -69,9 +69,9 @@ def test_can_use_output_schema_with_tools(
 ) -> None:
   """Test can_use_output_schema_with_tools."""
   if env_value is not None:
-    monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", env_value)
+    monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", env_value)
   else:
-    monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False)
+    monkeypatch.delenv("GOOGLE_GENAI_USE_ENTERPRISE", raising=False)
   assert can_use_output_schema_with_tools(model) == expected
 
 
@@ -90,9 +90,9 @@ def test_can_use_output_schema_with_tools_claude(
   """Test can_use_output_schema_with_tools with Claude models."""
   claude_model = _make_claude(model)
   if env_value is not None:
-    monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", env_value)
+    monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", env_value)
   else:
-    monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False)
+    monkeypatch.delenv("GOOGLE_GENAI_USE_ENTERPRISE", raising=False)
   assert can_use_output_schema_with_tools(claude_model) == expected
 
 
@@ -113,7 +113,7 @@ def test_can_use_output_schema_with_tools_litellm(
   """Test can_use_output_schema_with_tools with LiteLLM models."""
   litellm_model = _make_litellm(model)
   if env_value is not None:
-    monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", env_value)
+    monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", env_value)
   else:
-    monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False)
+    monkeypatch.delenv("GOOGLE_GENAI_USE_ENTERPRISE", raising=False)
   assert can_use_output_schema_with_tools(litellm_model) == expected

tests/unittests/utils/test_variant_utils.py

@@ -0,0 +1,43 @@
+# 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.
+
+"""Tests for variant_utils."""
+
+import warnings
+
+from google.adk.utils import variant_utils
+from google.adk.utils.variant_utils import GoogleLLMVariant
+
+
+def test_get_google_llm_variant_enterprise(monkeypatch):
+  monkeypatch.setenv('GOOGLE_GENAI_USE_ENTERPRISE', 'true')
+  assert variant_utils.get_google_llm_variant() == GoogleLLMVariant.VERTEX_AI
+
+
+def test_get_google_llm_variant_vertexai_fallback(monkeypatch):
+  monkeypatch.delenv('GOOGLE_GENAI_USE_ENTERPRISE', raising=False)
+  monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter('always')
+    result = variant_utils.get_google_llm_variant()
+    assert result == GoogleLLMVariant.VERTEX_AI
+    assert len(w) == 1
+    assert issubclass(w[-1].category, DeprecationWarning)
+    assert 'GOOGLE_GENAI_USE_VERTEXAI is deprecated' in str(w[-1].message)
+
+
+def test_get_google_llm_variant_default(monkeypatch):
+  monkeypatch.delenv('GOOGLE_GENAI_USE_ENTERPRISE', raising=False)
+  monkeypatch.delenv('GOOGLE_GENAI_USE_VERTEXAI', raising=False)
+  assert variant_utils.get_google_llm_variant() == GoogleLLMVariant.GEMINI_API

tests/unittests/utils/test_vertex_ai_utils.py

@@ -15,6 +15,7 @@
 """Tests for vertex_utils."""
 
 from unittest import mock
+import warnings
 
 from google.adk.utils import vertex_ai_utils
 import pytest
@@ -77,7 +78,7 @@ def test_get_express_mode_api_key(
 ):
   env_vars = {}
   if use_vertexai_env:
-    env_vars['GOOGLE_GENAI_USE_VERTEXAI'] = use_vertexai_env
+    env_vars['GOOGLE_GENAI_USE_ENTERPRISE'] = use_vertexai_env
   if google_api_key_env:
     env_vars['GOOGLE_API_KEY'] = google_api_key_env
   with mock.patch.dict('os.environ', env_vars, clear=True):
@@ -89,3 +90,26 @@ def test_get_express_mode_api_key(
         )
         == expected
     )
+
+
+def test_get_express_mode_api_key_enterprise(monkeypatch):
+  monkeypatch.setenv('GOOGLE_GENAI_USE_ENTERPRISE', 'true')
+  monkeypatch.setenv('GOOGLE_API_KEY', 'google_key')
+  assert (
+      vertex_ai_utils.get_express_mode_api_key(None, None, None) == 'google_key'
+  )
+
+
+def test_get_express_mode_api_key_vertexai_fallback_warning(monkeypatch):
+  monkeypatch.delenv('GOOGLE_GENAI_USE_ENTERPRISE', raising=False)
+  monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
+  monkeypatch.setenv('GOOGLE_API_KEY', 'google_key')
+  # Should trigger a deprecation warning and return the key
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter('always')
+    result = vertex_ai_utils.get_express_mode_api_key(None, None, None)
+
+    assert result == 'google_key'
+    assert len(w) == 1
+    assert issubclass(w[-1].category, DeprecationWarning)
+    assert 'GOOGLE_GENAI_USE_VERTEXAI is deprecated' in str(w[-1].message)