Skip to content

Commit 3619b29

Browse files
shuoweilshuoweil
authored
feat: support gemini-3.x models in loader and update default model to gemini-3.5-flash (#17557)
** STILL IN CLEAN UP STAGE" - Updated the default model in BQML `GeminiTextGenerator` to `gemini-3.5-flash` (instead of `gemini-2.5-pro`). - Registered `gemini-3.1-flash-lite` and `gemini-3.5-flash` endpoints in BQML model loader to prevent `KeyError` during model loading. - Updated all related docstrings, deprecation notes, warnings, and tests to align with the new defaults. - Updated Gemini model references across all sample notebooks to `gemini-3.5-flash` and fixed notebook formatting syntax to pass pre-commit checks. Fixes #<519302924> 🦕
1 parent 17bef41 commit 3619b29

32 files changed

Lines changed: 3770 additions & 3728 deletions

File tree

‎packages/bigframes/bigframes/bigquery/_operations/ai.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ def generate(
8787
Specifies the connection to use to communicate with the model. For example, ``myproject.us.myconnection``.
8888
If not provided, the query uses your end-user credential.
8989
endpoint (str, optional):
90-
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-2.5-flash"``. You can specify any
90+
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-3.5-flash"``. You can specify any
9191
generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
9292
uses the full endpoint of the model. If you don't specify an ENDPOINT value, BigQuery ML selects a recent stable
9393
version of Gemini to use.
@@ -178,7 +178,7 @@ def generate_bool(
178178
Specifies the connection to use to communicate with the model. For example, ``myproject.us.myconnection``.
179179
If not provided, the query uses your end-user credential.
180180
endpoint (str, optional):
181-
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-2.5-flash"``. You can specify any
181+
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-3.5-flash"``. You can specify any
182182
generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
183183
uses the full endpoint of the model. If you don't specify an ENDPOINT value, BigQuery ML selects a recent stable
184184
version of Gemini to use.
@@ -252,7 +252,7 @@ def generate_int(
252252
Specifies the connection to use to communicate with the model. For example, ``myproject.us.myconnection``.
253253
If not provided, the query uses your end-user credential.
254254
endpoint (str, optional):
255-
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-2.5-flash"``. You can specify any
255+
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-3.5-flash"``. You can specify any
256256
generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
257257
uses the full endpoint of the model. If you don't specify an ENDPOINT value, BigQuery ML selects a recent stable
258258
version of Gemini to use.
@@ -326,7 +326,7 @@ def generate_double(
326326
Specifies the connection to use to communicate with the model. For example, ``myproject.us.myconnection``.
327327
If not provided, the query uses your end-user credential.
328328
endpoint (str, optional):
329-
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-2.5-flash"``. You can specify any
329+
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-3.5-flash"``. You can specify any
330330
generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
331331
uses the full endpoint of the model. If you don't specify an ENDPOINT value, BigQuery ML selects a recent stable
332332
version of Gemini to use.
@@ -809,7 +809,7 @@ def if_(
809809
Specifies the connection to use to communicate with the model. For example, ``myproject.us.myconnection``.
810810
If not provided, the query uses your end-user credential.
811811
endpoint (str, optional):
812-
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-2.5-flash"``. You can specify any
812+
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-3.5-flash"``. You can specify any
813813
generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
814814
uses the full endpoint of the model. If you don't specify an ENDPOINT value, BigQuery ML dynamically chooses a model based on your query to have the
815815
best cost to quality tradeoff for the task.
@@ -960,7 +960,7 @@ def score(
960960
Specifies the connection to use to communicate with the model. For example, ``myproject.us.myconnection``.
961961
If not provided, the query uses your end-user credential.
962962
endpoint (str, optional):
963-
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-2.5-flash"``. You can specify any
963+
Specifies the Vertex AI endpoint to use for the model. For example ``"gemini-3.5-flash"``. You can specify any
964964
generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
965965
uses the full endpoint of the model. If you don't specify an endpoint value, BigQuery ML dynamically chooses a model
966966
based on your query to have the best cost to quality tradeoff for the task.

‎packages/bigframes/bigframes/ml/llm.py‎

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,8 @@
5959
_GEMINI_2P5_PRO_ENDPOINT = "gemini-2.5-pro"
6060
_GEMINI_2P5_FLASH_ENDPOINT = "gemini-2.5-flash"
6161
_GEMINI_2P5_FLASH_LITE_ENDPOINT = "gemini-2.5-flash-lite"
62+
_GEMINI_3P1_FLASH_LITE_ENDPOINT = "gemini-3.1-flash-lite"
63+
_GEMINI_3P5_FLASH_ENDPOINT = "gemini-3.5-flash"
6264

6365
_GEMINI_ENDPOINTS = (
6466
_GEMINI_1P5_PRO_PREVIEW_ENDPOINT,
@@ -73,6 +75,8 @@
7375
_GEMINI_2P5_PRO_ENDPOINT,
7476
_GEMINI_2P5_FLASH_ENDPOINT,
7577
_GEMINI_2P5_FLASH_LITE_ENDPOINT,
78+
_GEMINI_3P1_FLASH_LITE_ENDPOINT,
79+
_GEMINI_3P5_FLASH_ENDPOINT,
7680
)
7781
_GEMINI_PREVIEW_ENDPOINTS = (
7882
_GEMINI_1P5_PRO_PREVIEW_ENDPOINT,
@@ -96,6 +100,8 @@
96100
_GEMINI_2P5_PRO_ENDPOINT,
97101
_GEMINI_2P5_FLASH_ENDPOINT,
98102
_GEMINI_2P5_FLASH_LITE_ENDPOINT,
103+
_GEMINI_3P1_FLASH_LITE_ENDPOINT,
104+
_GEMINI_3P5_FLASH_ENDPOINT,
99105
)
100106

101107
_CLAUDE_3_SONNET_ENDPOINT = "claude-3-sonnet"
@@ -431,21 +437,22 @@ class GeminiTextGenerator(base.RetriableRemotePredictor):
431437
"""Gemini text generator LLM model.
432438
433439
.. note::
434-
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.
440+
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.
435441
436442
Args:
437-
model_name (str, Default to "gemini-2.0-flash-001"):
443+
model_name (str, Default to "gemini-3.5-flash"):
438444
The model for natural language tasks. Accepted values are
439445
"gemini-1.5-pro-preview-0514", "gemini-1.5-flash-preview-0514",
440446
"gemini-1.5-pro-001", "gemini-1.5-pro-002", "gemini-1.5-flash-001",
441447
"gemini-1.5-flash-002", "gemini-2.0-flash-exp",
442448
"gemini-2.0-flash-lite-001", "gemini-2.0-flash-001",
443-
"gemini-2.5-pro", "gemini-2.5-flash" and "gemini-2.5-flash-lite".
444-
If no setting is provided, "gemini-2.0-flash-001" will be used by
449+
"gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-lite",
450+
"gemini-3.1-flash-lite" and "gemini-3.5-flash".
451+
If no setting is provided, "gemini-3.5-flash" will be used by
445452
default and a warning will be issued.
446453
447454
.. note::
448-
"gemini-1.5-X" is going to be deprecated. Please use gemini-2.5-X instead. For example, "gemini-2.5-flash".
455+
"gemini-1.5-X" is going to be deprecated. Please use gemini-3.5-flash instead.
449456
"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
450457
Service Specific Terms(https://cloud.google.com/terms/service-terms#1). Pre-GA products and features are available "as is"
451458
and might have limited support. For more information, see the launch stage descriptions
@@ -478,6 +485,8 @@ def __init__(
478485
"gemini-2.5-pro",
479486
"gemini-2.5-flash",
480487
"gemini-2.5-flash-lite",
488+
"gemini-3.1-flash-lite",
489+
"gemini-3.5-flash",
481490
]
482491
] = None,
483492
session: Optional[bigframes.Session] = None,
@@ -496,7 +505,7 @@ def __init__(
496505
warnings.warn(msg, category=exceptions.PreviewWarning)
497506

498507
if model_name is None:
499-
model_name = "gemini-2.0-flash-001"
508+
model_name = "gemini-3.5-flash"
500509
msg = exceptions.format_message(_REMOVE_DEFAULT_MODEL_WARNING)
501510
warnings.warn(msg, category=FutureWarning, stacklevel=2)
502511

@@ -522,15 +531,21 @@ def _create_bqml_model(self):
522531
)
523532
)
524533
warnings.warn(msg)
525-
if self.model_name.startswith("gemini-1.5"):
534+
if self.model_name.startswith(
535+
("gemini-2.0", "gemini-1.5")
536+
) or self.model_name in (
537+
"gemini-2.5-flash",
538+
"gemini-2.5-flash-lite",
539+
"gemini-2.5-pro-preview-05-06",
540+
):
526541
msg = exceptions.format_message(
527542
_MODEL_DEPRECATE_WARNING.format(
528543
model_name=self.model_name,
529-
new_model_name="gemini-2.5-X",
544+
new_model_name="gemini-3.5-flash",
530545
link="https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator",
531546
)
532547
)
533-
warnings.warn(msg)
548+
warnings.warn(msg, category=exceptions.ApiDeprecationWarning)
534549

535550
options = {"endpoint": self.model_name}
536551

‎packages/bigframes/bigframes/ml/loader.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@
7070
llm._GEMINI_2P5_FLASH_ENDPOINT: llm.GeminiTextGenerator,
7171
llm._GEMINI_2P5_FLASH_LITE_ENDPOINT: llm.GeminiTextGenerator,
7272
llm._GEMINI_2P5_PRO_ENDPOINT: llm.GeminiTextGenerator,
73+
llm._GEMINI_3P1_FLASH_LITE_ENDPOINT: llm.GeminiTextGenerator,
74+
llm._GEMINI_3P5_FLASH_ENDPOINT: llm.GeminiTextGenerator,
7375
llm._CLAUDE_3_HAIKU_ENDPOINT: llm.Claude3TextGenerator,
7476
llm._CLAUDE_3_SONNET_ENDPOINT: llm.Claude3TextGenerator,
7577
llm._CLAUDE_3_5_SONNET_ENDPOINT: llm.Claude3TextGenerator,

‎packages/bigframes/bigframes/operations/semantics.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ def agg(
5757
>>> bpd.options.compute.semantic_ops_confirmation_threshold = 25
5858
5959
>>> import bigframes.ml.llm as llm
60-
>>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001") # doctest: +SKIP
60+
>>> model = llm.GeminiTextGenerator(model_name="gemini-3.5-flash") # doctest: +SKIP
6161
6262
>>> df = bpd.DataFrame(
6363
... {
@@ -330,7 +330,7 @@ def filter(self, instruction: str, model, ground_with_google_search: bool = Fals
330330
>>> bpd.options.compute.semantic_ops_confirmation_threshold = 25
331331
332332
>>> import bigframes.ml.llm as llm
333-
>>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001") # doctest: +SKIP
333+
>>> model = llm.GeminiTextGenerator(model_name="gemini-3.5-flash") # doctest: +SKIP
334334
335335
>>> df = bpd.DataFrame({"country": ["USA", "Germany"], "city": ["Seattle", "Berlin"]})
336336
>>> df.semantics.filter("{city} is the capital of {country}", model) # doctest: +SKIP
@@ -443,7 +443,7 @@ def map(
443443
>>> bpd.options.compute.semantic_ops_confirmation_threshold = 25
444444
445445
>>> import bigframes.ml.llm as llm
446-
>>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001") # doctest: +SKIP
446+
>>> model = llm.GeminiTextGenerator(model_name="gemini-3.5-flash") # doctest: +SKIP
447447
448448
>>> df = bpd.DataFrame({"ingredient_1": ["Burger Bun", "Soy Bean"], "ingredient_2": ["Beef Patty", "Bittern"]})
449449
>>> df.semantics.map("What is the food made from {ingredient_1} and {ingredient_2}? One word only.", output_column="food", model=model) # doctest: +SKIP
@@ -565,7 +565,7 @@ def join(
565565
>>> bpd.options.compute.semantic_ops_confirmation_threshold = 25
566566
567567
>>> import bigframes.ml.llm as llm
568-
>>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001") # doctest: +SKIP
568+
>>> model = llm.GeminiTextGenerator(model_name="gemini-3.5-flash") # doctest: +SKIP
569569
570570
>>> cities = bpd.DataFrame({'city': ['Seattle', 'Ottawa', 'Berlin', 'Shanghai', 'New Delhi']})
571571
>>> continents = bpd.DataFrame({'continent': ['North America', 'Africa', 'Asia']})
@@ -805,7 +805,7 @@ def top_k(
805805
>>> bpd.options.compute.semantic_ops_confirmation_threshold = 25
806806
807807
>>> import bigframes.ml.llm as llm
808-
>>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001") # doctest: +SKIP
808+
>>> model = llm.GeminiTextGenerator(model_name="gemini-3.5-flash") # doctest: +SKIP
809809
810810
>>> df = bpd.DataFrame(
811811
... {

0 commit comments

Comments
 (0)