chore: Checks gemini version for 2 and above for gemini-builtin tools

PiperOrigin-RevId: 820848897

Author: Jacksunwei

src/google/adk/code_executors/built_in_code_executor.py

@@ -19,7 +19,7 @@
 
 from ..agents.invocation_context import InvocationContext
 from ..models import LlmRequest
-from ..utils.model_name_utils import is_gemini_2_model
+from ..utils.model_name_utils import is_gemini_2_or_above
 from .base_code_executor import BaseCodeExecutor
 from .code_execution_utils import CodeExecutionInput
 from .code_execution_utils import CodeExecutionResult
@@ -42,7 +42,7 @@ def execute_code(
 
   def process_llm_request(self, llm_request: LlmRequest) -> None:
     """Pre-process the LLM request for Gemini 2.0+ models to use the code execution tool."""
-    if is_gemini_2_model(llm_request.model):
+    if is_gemini_2_or_above(llm_request.model):
       llm_request.config = llm_request.config or types.GenerateContentConfig()
       llm_request.config.tools = llm_request.config.tools or []
       llm_request.config.tools.append(

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

@@ -24,7 +24,7 @@
 from typing_extensions import override
 from vertexai.preview import rag
 
-from ...utils.model_name_utils import is_gemini_2_model
+from ...utils.model_name_utils import is_gemini_2_or_above
 from ..tool_context import ToolContext
 from .base_retrieval_tool import BaseRetrievalTool
 
@@ -63,7 +63,7 @@ async def process_llm_request(
       llm_request: LlmRequest,
   ) -> None:
     # Use Gemini built-in Vertex AI RAG tool for Gemini 2 models.
-    if is_gemini_2_model(llm_request.model):
+    if is_gemini_2_or_above(llm_request.model):
       llm_request.config = (
           types.GenerateContentConfig()
           if not llm_request.config

src/google/adk/tools/url_context_tool.py

@@ -20,7 +20,7 @@
 from typing_extensions import override
 
 from ..utils.model_name_utils import is_gemini_1_model
-from ..utils.model_name_utils import is_gemini_2_model
+from ..utils.model_name_utils import is_gemini_2_or_above
 from .base_tool import BaseTool
 from .tool_context import ToolContext
 
@@ -50,7 +50,7 @@ async def process_llm_request(
     llm_request.config.tools = llm_request.config.tools or []
     if is_gemini_1_model(llm_request.model):
       raise ValueError('Url context tool can not be used in Gemini 1.x.')
-    elif is_gemini_2_model(llm_request.model):
+    elif is_gemini_2_or_above(llm_request.model):
       llm_request.config.tools.append(
           types.Tool(url_context=types.UrlContext())
       )

src/google/adk/utils/model_name_utils.py

@@ -19,6 +19,9 @@
 import re
 from typing import Optional
 
+from packaging.version import InvalidVersion
+from packaging.version import Version
+
 
 def extract_model_name(model_string: str) -> str:
   """Extract the actual model name from either simple or path-based format.
@@ -74,17 +77,29 @@ def is_gemini_1_model(model_string: Optional[str]) -> bool:
   return re.match(r'^gemini-1\.\d+', model_name) is not None
 
 
-def is_gemini_2_model(model_string: Optional[str]) -> bool:
-  """Check if the model is a Gemini 2.x model using regex patterns.
+def is_gemini_2_or_above(model_string: Optional[str]) -> bool:
+  """Check if the model is a Gemini 2.0 or newer model using semantic versions.
 
   Args:
     model_string: Either a simple model name or path-based model name
 
   Returns:
-    True if it's a Gemini 2.x model, False otherwise
+    True if it's a Gemini 2.0+ model, False otherwise
   """
   if not model_string:
     return False
 
   model_name = extract_model_name(model_string)
-  return re.match(r'^gemini-2\.\d+', model_name) is not None
+  if not model_name.startswith('gemini-'):
+    return False
+
+  version_string = model_name[len('gemini-') :].split('-', 1)[0]
+  if not version_string:
+    return False
+
+  try:
+    parsed_version = Version(version_string)
+  except InvalidVersion:
+    return False
+
+  return parsed_version.major >= 2

tests/unittests/utils/test_model_name_utils.py

@@ -16,7 +16,7 @@
 
 from google.adk.utils.model_name_utils import extract_model_name
 from google.adk.utils.model_name_utils import is_gemini_1_model
-from google.adk.utils.model_name_utils import is_gemini_2_model
+from google.adk.utils.model_name_utils import is_gemini_2_or_above
 from google.adk.utils.model_name_utils import is_gemini_model
 
 
@@ -165,46 +165,51 @@ def test_is_gemini_1_model_edge_cases(self):
 
 
 class TestIsGemini2Model:
-  """Test the is_gemini_2_model function."""
-
-  def test_is_gemini_2_model_simple_names(self):
-    """Test Gemini 2.x model detection with simple model names."""
-    assert is_gemini_2_model('gemini-2.0-flash') is True
-    assert is_gemini_2_model('gemini-2.5-pro') is True
-    assert is_gemini_2_model('gemini-2.0-flash-001') is True
-    assert is_gemini_2_model('gemini-2.9-experimental') is True
-    assert is_gemini_2_model('gemini-1.5-flash') is False
-    assert is_gemini_2_model('gemini-1.0-pro') is False
-    assert is_gemini_2_model('gemini-3.0-pro') is False  # Only 2.x versions
-    assert is_gemini_2_model('claude-3-sonnet') is False
-
-  def test_is_gemini_2_model_path_based_names(self):
-    """Test Gemini 2.x model detection with path-based model names."""
+  """Test the is_gemini_2_or_above function."""
+
+  def test_is_gemini_2_or_above_simple_names(self):
+    """Test Gemini 2.0+ model detection with simple model names."""
+    assert is_gemini_2_or_above('gemini-2.0-flash') is True
+    assert is_gemini_2_or_above('gemini-2.5-pro') is True
+    assert is_gemini_2_or_above('gemini-2.0-flash-001') is True
+    assert is_gemini_2_or_above('gemini-2.9-experimental') is True
+    assert is_gemini_2_or_above('gemini-2-pro') is True
+    assert is_gemini_2_or_above('gemini-2') is True
+    assert is_gemini_2_or_above('gemini-3.0-pro') is True
+    assert is_gemini_2_or_above('gemini-1.5-flash') is False
+    assert is_gemini_2_or_above('gemini-1.0-pro') is False
+    assert is_gemini_2_or_above('claude-3-sonnet') is False
+
+  def test_is_gemini_2_or_above_path_based_names(self):
+    """Test Gemini 2.0+ model detection with path-based model names."""
     gemini_2_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-2.0-flash-001'
-    assert is_gemini_2_model(gemini_2_path) is True
+    assert is_gemini_2_or_above(gemini_2_path) is True
 
     gemini_2_path_2 = 'projects/12345/locations/us-east1/publishers/google/models/gemini-2.5-pro-preview'
-    assert is_gemini_2_model(gemini_2_path_2) is True
+    assert is_gemini_2_or_above(gemini_2_path_2) is True
 
     gemini_1_path = 'projects/265104255505/locations/us-central1/publishers/google/models/gemini-1.5-flash-001'
-    assert is_gemini_2_model(gemini_1_path) is False
+    assert is_gemini_2_or_above(gemini_1_path) is False
 
-  def test_is_gemini_2_model_edge_cases(self):
-    """Test edge cases for Gemini 2.x model detection."""
+    gemini_3_path = 'projects/12345/locations/us-east1/publishers/google/models/gemini-3.0-pro'
+    assert is_gemini_2_or_above(gemini_3_path) is True
+
+  def test_is_gemini_2_or_above_edge_cases(self):
+    """Test edge cases for Gemini 2.0+ model detection."""
     # Test with None
-    assert is_gemini_2_model(None) is False
+    assert is_gemini_2_or_above(None) is False
 
     # Test with empty string
-    assert is_gemini_2_model('') is False
+    assert is_gemini_2_or_above('') is False
 
     # Test with model names containing gemini-2 but not starting with it
-    assert is_gemini_2_model('my-gemini-2.5-model') is False
-    assert is_gemini_2_model('custom-gemini-2.0-flash') is False
+    assert is_gemini_2_or_above('my-gemini-2.5-model') is False
+    assert is_gemini_2_or_above('custom-gemini-2.0-flash') is False
 
     # Test with invalid versions
-    assert is_gemini_2_model('gemini-2') is False  # Missing dot
-    assert is_gemini_2_model('gemini-2-pro') is False  # Missing dot
-    assert is_gemini_2_model('gemini-2.') is False  # Missing version number
+    assert is_gemini_2_or_above('gemini-2.') is False  # Missing version number
+    assert is_gemini_2_or_above('gemini-0.9-test') is False
+    assert is_gemini_2_or_above('gemini-one') is False
 
 
 class TestModelNameUtilsIntegration:
@@ -216,32 +221,34 @@ def test_model_classification_consistency(self):
         'gemini-1.5-flash',
         'gemini-2.0-flash',
         'gemini-2.5-pro',
+        'gemini-3.0-pro',
         'projects/123/locations/us-central1/publishers/google/models/gemini-1.5-pro',
         'projects/123/locations/us-central1/publishers/google/models/gemini-2.0-flash',
+        'projects/123/locations/us-central1/publishers/google/models/gemini-3.0-pro',
         'claude-3-sonnet',
         'gpt-4',
     ]
 
     for model in test_models:
-      # A model can only be either Gemini 1.x or Gemini 2.x, not both
+      # A model can only be either Gemini 1.x or Gemini 2.0+, not both
       if is_gemini_1_model(model):
-        assert not is_gemini_2_model(
+        assert not is_gemini_2_or_above(
             model
-        ), f'Model {model} classified as both Gemini 1.x and 2.x'
+        ), f'Model {model} classified as both Gemini 1.x and 2.0+'
         assert is_gemini_model(
             model
         ), f'Model {model} is Gemini 1.x but not classified as Gemini'
 
-      if is_gemini_2_model(model):
+      if is_gemini_2_or_above(model):
         assert not is_gemini_1_model(
             model
-        ), f'Model {model} classified as both Gemini 1.x and 2.x'
+        ), f'Model {model} classified as both Gemini 1.x and 2.0+'
         assert is_gemini_model(
             model
-        ), f'Model {model} is Gemini 2.x but not classified as Gemini'
+        ), f'Model {model} is Gemini 2.0+ but not classified as Gemini'
 
-      # If it's neither Gemini 1.x nor 2.x, it should not be classified as Gemini
-      if not is_gemini_1_model(model) and not is_gemini_2_model(model):
+      # If it's neither Gemini 1.x nor 2.0+, it should not be classified as Gemini
+      if not is_gemini_1_model(model) and not is_gemini_2_or_above(model):
         if model and 'gemini-' not in extract_model_name(model):
           assert not is_gemini_model(
               model
@@ -262,6 +269,10 @@ def test_path_vs_simple_model_consistency(self):
             'gemini-2.5-pro',
             'projects/123/locations/us-central1/publishers/google/models/gemini-2.5-pro',
         ),
+        (
+            'gemini-3.0-pro',
+            'projects/123/locations/us-central1/publishers/google/models/gemini-3.0-pro',
+        ),
         (
             'claude-3-sonnet',
             'projects/123/locations/us-central1/publishers/google/models/claude-3-sonnet',
@@ -278,7 +289,9 @@ def test_path_vs_simple_model_consistency(self):
           f'Inconsistent Gemini 1.x classification for {simple_model} vs'
           f' {path_model}'
       )
-      assert is_gemini_2_model(simple_model) == is_gemini_2_model(path_model), (
-          f'Inconsistent Gemini 2.x classification for {simple_model} vs'
+      assert is_gemini_2_or_above(simple_model) == is_gemini_2_or_above(
+          path_model
+      ), (
+          f'Inconsistent Gemini 2.0+ classification for {simple_model} vs'
           f' {path_model}'
       )