feat: support gemini-3.x models in loader and update default model to gemini-3.5-flash

packages/bigframes/bigframes/ml/llm.py

@@ -59,6 +59,8 @@
 _GEMINI_2P5_PRO_ENDPOINT = "gemini-2.5-pro"
 _GEMINI_2P5_FLASH_ENDPOINT = "gemini-2.5-flash"
 _GEMINI_2P5_FLASH_LITE_ENDPOINT = "gemini-2.5-flash-lite"
+_GEMINI_3P1_FLASH_LITE_ENDPOINT = "gemini-3.1-flash-lite"
+_GEMINI_3P5_FLASH_ENDPOINT = "gemini-3.5-flash"
 
 _GEMINI_ENDPOINTS = (
     _GEMINI_1P5_PRO_PREVIEW_ENDPOINT,
@@ -73,6 +75,8 @@
     _GEMINI_2P5_PRO_ENDPOINT,
     _GEMINI_2P5_FLASH_ENDPOINT,
     _GEMINI_2P5_FLASH_LITE_ENDPOINT,
+    _GEMINI_3P1_FLASH_LITE_ENDPOINT,
+    _GEMINI_3P5_FLASH_ENDPOINT,
 )
 _GEMINI_PREVIEW_ENDPOINTS = (
     _GEMINI_1P5_PRO_PREVIEW_ENDPOINT,
@@ -96,6 +100,8 @@
     _GEMINI_2P5_PRO_ENDPOINT,
     _GEMINI_2P5_FLASH_ENDPOINT,
     _GEMINI_2P5_FLASH_LITE_ENDPOINT,
+    _GEMINI_3P1_FLASH_LITE_ENDPOINT,
+    _GEMINI_3P5_FLASH_ENDPOINT,
 )
 
 _CLAUDE_3_SONNET_ENDPOINT = "claude-3-sonnet"
@@ -431,21 +437,22 @@ class GeminiTextGenerator(base.RetriableRemotePredictor):
     """Gemini text generator LLM model.
 
     .. note::
-        gemini-1.5-X are going to be deprecated. Use gemini-2.5-X (https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator) instead.
+        gemini-1.5-X are going to be deprecated. Use gemini-3.5-flash (https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator) instead.
 
     Args:
-        model_name (str, Default to "gemini-2.0-flash-001"):
+        model_name (str, Default to "gemini-3.5-flash"):
             The model for natural language tasks. Accepted values are
             "gemini-1.5-pro-preview-0514", "gemini-1.5-flash-preview-0514",
             "gemini-1.5-pro-001", "gemini-1.5-pro-002", "gemini-1.5-flash-001",
             "gemini-1.5-flash-002", "gemini-2.0-flash-exp",
             "gemini-2.0-flash-lite-001", "gemini-2.0-flash-001",
-            "gemini-2.5-pro", "gemini-2.5-flash" and "gemini-2.5-flash-lite".
-            If no setting is provided, "gemini-2.0-flash-001" will be used by
+            "gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-lite",
+            "gemini-3.1-flash-lite" and "gemini-3.5-flash".
+            If no setting is provided, "gemini-3.5-flash" will be used by
             default and a warning will be issued.
 
         .. note::
-            "gemini-1.5-X" is going to be deprecated. Please use gemini-2.5-X instead. For example, "gemini-2.5-flash".
+            "gemini-1.5-X" is going to be deprecated. Please use gemini-3.5-flash instead.
             "gemini-2.0-flash-exp", "gemini-1.5-pro-preview-0514" and "gemini-1.5-flash-preview-0514" is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the
             Service Specific Terms(https://cloud.google.com/terms/service-terms#1). Pre-GA products and features are available "as is"
             and might have limited support. For more information, see the launch stage descriptions
@@ -478,6 +485,8 @@ def __init__(
                 "gemini-2.5-pro",
                 "gemini-2.5-flash",
                 "gemini-2.5-flash-lite",
+                "gemini-3.1-flash-lite",
+                "gemini-3.5-flash",
             ]
         ] = None,
         session: Optional[bigframes.Session] = None,
@@ -496,7 +505,7 @@ def __init__(
             warnings.warn(msg, category=exceptions.PreviewWarning)
 
         if model_name is None:
-            model_name = "gemini-2.0-flash-001"
+            model_name = "gemini-3.5-flash"
             msg = exceptions.format_message(_REMOVE_DEFAULT_MODEL_WARNING)
             warnings.warn(msg, category=FutureWarning, stacklevel=2)
 
@@ -522,15 +531,21 @@ def _create_bqml_model(self):
                 )
             )
             warnings.warn(msg)
-        if self.model_name.startswith("gemini-1.5"):
+        if self.model_name.startswith(
+            ("gemini-2.0", "gemini-1.5")
+        ) or self.model_name in (
+            "gemini-2.5-flash",
+            "gemini-2.5-flash-lite",
+            "gemini-2.5-pro-preview-05-06",
+        ):
             msg = exceptions.format_message(
                 _MODEL_DEPRECATE_WARNING.format(
                     model_name=self.model_name,
-                    new_model_name="gemini-2.5-X",
+                    new_model_name="gemini-3.5-flash",
                     link="https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator",
                 )
             )
-            warnings.warn(msg)
+            warnings.warn(msg, category=exceptions.ApiDeprecationWarning)
 
         options = {"endpoint": self.model_name}
 

packages/bigframes/bigframes/ml/loader.py

@@ -70,6 +70,8 @@
         llm._GEMINI_2P5_FLASH_ENDPOINT: llm.GeminiTextGenerator,
         llm._GEMINI_2P5_FLASH_LITE_ENDPOINT: llm.GeminiTextGenerator,
         llm._GEMINI_2P5_PRO_ENDPOINT: llm.GeminiTextGenerator,
+        llm._GEMINI_3P1_FLASH_LITE_ENDPOINT: llm.GeminiTextGenerator,
+        llm._GEMINI_3P5_FLASH_ENDPOINT: llm.GeminiTextGenerator,
         llm._CLAUDE_3_HAIKU_ENDPOINT: llm.Claude3TextGenerator,
         llm._CLAUDE_3_SONNET_ENDPOINT: llm.Claude3TextGenerator,
         llm._CLAUDE_3_5_SONNET_ENDPOINT: llm.Claude3TextGenerator,