Commit 5770ff6
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -440,15 +440,15 @@ class GeminiTextGenerator(base.RetriableRemotePredictor): | |||
| 440 | 440 | 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. | |
| 441 | 441 | ||
| 442 | 442 | Args: | |
| 443 | - model_name (str, Default to "gemini-2.0-flash-001"): | ||
| 443 | + model_name (str, Default to "gemini-2.5-flash"): | ||
| 444 | 444 | The model for natural language tasks. Accepted values are | |
| 445 | 445 | "gemini-1.5-pro-preview-0514", "gemini-1.5-flash-preview-0514", | |
| 446 | 446 | "gemini-1.5-pro-001", "gemini-1.5-pro-002", "gemini-1.5-flash-001", | |
| 447 | 447 | "gemini-1.5-flash-002", "gemini-2.0-flash-exp", | |
| 448 | 448 | "gemini-2.0-flash-lite-001", "gemini-2.0-flash-001", | |
| 449 | 449 | "gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-lite", | |
| 450 | 450 | "gemini-3.1-flash-lite" and "gemini-3.5-flash". | |
| 451 | - If no setting is provided, "gemini-2.0-flash-001" will be used by | ||
| 451 | + If no setting is provided, "gemini-2.5-flash" will be used by | ||
| 452 | 452 | default and a warning will be issued. | |
| 453 | 453 | ||
| 454 | 454 | .. note:: | |
@@ -505,7 +505,7 @@ def __init__( | |||
| 505 | 505 | warnings.warn(msg, category=exceptions.PreviewWarning) | |
| 506 | 506 | ||
| 507 | 507 | if model_name is None: | |
| 508 | - model_name = "gemini-2.0-flash-001" | ||
| 508 | + model_name = "gemini-2.5-flash" | ||
| 509 | 509 | msg = exceptions.format_message(_REMOVE_DEFAULT_MODEL_WARNING) | |
| 510 | 510 | warnings.warn(msg, category=FutureWarning, stacklevel=2) | |
| 511 | 511 | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -0,0 +1,50 @@ | |||
| 1 | + # Copyright 2026 Google LLC | ||
| 2 | + # | ||
| 3 | + # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | + # you may not use this file except in compliance with the License. | ||
| 5 | + # You may obtain a copy of the License at | ||
| 6 | + # | ||
| 7 | + # http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | + # | ||
| 9 | + # Unless required by applicable law or agreed to in writing, software | ||
| 10 | + # distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | + # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | + # See the License for the specific language governing permissions and | ||
| 13 | + # limitations under the License. | ||
| 14 | + | ||
| 15 | + from unittest import mock | ||
| 16 | + | ||
| 17 | + import pytest | ||
| 18 | + from google.cloud import bigquery | ||
| 19 | + | ||
| 20 | + import bigframes.session | ||
| 21 | + from bigframes.ml import llm | ||
| 22 | + | ||
| 23 | + | ||
| 24 | + def test_gemini_text_generator_default_model(): | ||
| 25 | + mock_session = mock.create_autospec(spec=bigframes.session.Session) | ||
| 26 | + mock_session._create_bq_connection.return_value = ( | ||
| 27 | + "projects/test-project/locations/us-central1/connections/test-conn" | ||
| 28 | + ) | ||
| 29 | + mock_session._anonymous_dataset = bigquery.DatasetReference( | ||
| 30 | + "test-project", "test_dataset" | ||
| 31 | + ) | ||
| 32 | + mock_job = mock.MagicMock() | ||
| 33 | + mock_job.destination.project = "test-project" | ||
| 34 | + mock_job.destination.dataset_id = "test_dataset" | ||
| 35 | + mock_job.destination.table_id = "test_model" | ||
| 36 | + mock_session._start_query_ml_ddl.return_value = (None, mock_job) | ||
| 37 | + mock_session.bqclient.get_model.return_value = mock.MagicMock(spec=bigquery.Model) | ||
| 38 | + | ||
| 39 | + with pytest.warns( | ||
| 40 | + FutureWarning, match="default model will be removed in BigFrames 3.0" | ||
| 41 | + ): | ||
| 42 | + model = llm.GeminiTextGenerator( | ||
| 43 | + session=mock_session, | ||
| 44 | + connection_name="test-conn", | ||
| 45 | + ) | ||
| 46 | + | ||
| 47 | + assert model.model_name == "gemini-2.5-flash" | ||
| 48 | + mock_session._start_query_ml_ddl.assert_called_once() | ||
| 49 | + generated_sql = mock_session._start_query_ml_ddl.call_args[0][0] | ||
| 50 | + assert "gemini-2.5-flash" in generated_sql | ||
0 commit comments