From 77d6f397c214555915f1c602db2ed7213db69bd3 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Tue, 2 Jun 2026 22:50:46 +0000 Subject: [PATCH 1/5] feat: migrate gemini default and deprecations to gemini-2.5-pro --- packages/bigframes/bigframes/ml/llm.py | 12 ++++++---- .../tests/system/large/bigquery/test_ai.py | 2 +- .../tests/system/large/operations/conftest.py | 2 +- .../tests/system/small/bigquery/test_ai.py | 18 +++++++------- .../test_ai_classify_with_params/out.sql | 2 +- .../test_ai_ops/test_ai_generate/out.sql | 2 +- .../test_ai_ops/test_ai_generate_bool/out.sql | 2 +- .../out.sql | 2 +- .../test_ai_generate_double/out.sql | 2 +- .../out.sql | 2 +- .../test_ai_ops/test_ai_generate_int/out.sql | 2 +- .../out.sql | 2 +- .../out.sql | 2 +- .../out.sql | 2 +- .../test_ai_if_with_endpoint/out.sql | 2 +- .../out.sql | 2 +- .../sqlglot/expressions/test_ai_ops.py | 24 +++++++++---------- 17 files changed, 43 insertions(+), 39 deletions(-) diff --git a/packages/bigframes/bigframes/ml/llm.py b/packages/bigframes/bigframes/ml/llm.py index 3887453a2239..d626f177169d 100644 --- a/packages/bigframes/bigframes/ml/llm.py +++ b/packages/bigframes/bigframes/ml/llm.py @@ -496,7 +496,7 @@ def __init__( warnings.warn(msg, category=exceptions.PreviewWarning) if model_name is None: - model_name = "gemini-2.0-flash-001" + model_name = "gemini-2.5-pro" msg = exceptions.format_message(_REMOVE_DEFAULT_MODEL_WARNING) warnings.warn(msg, category=FutureWarning, stacklevel=2) @@ -522,15 +522,19 @@ def _create_bqml_model(self): ) ) warnings.warn(msg) - if self.model_name.startswith("gemini-1.5"): + if ( + self.model_name.startswith("gemini-2.0") + or self.model_name.startswith("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-2.5-pro", 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} diff --git a/packages/bigframes/tests/system/large/bigquery/test_ai.py b/packages/bigframes/tests/system/large/bigquery/test_ai.py index 504fe5aa3897..c74469026910 100644 --- a/packages/bigframes/tests/system/large/bigquery/test_ai.py +++ b/packages/bigframes/tests/system/large/bigquery/test_ai.py @@ -33,7 +33,7 @@ def text_model(bq_connection, dataset_id): model_name = f"{dataset_id}.text_model" return ml.create_model( model_name=model_name, - options={"endpoint": "gemini-2.5-flash"}, + options={"endpoint": "gemini-2.5-pro"}, connection_name=bq_connection, ) diff --git a/packages/bigframes/tests/system/large/operations/conftest.py b/packages/bigframes/tests/system/large/operations/conftest.py index 0122c860ca74..e4c63a359275 100644 --- a/packages/bigframes/tests/system/large/operations/conftest.py +++ b/packages/bigframes/tests/system/large/operations/conftest.py @@ -22,7 +22,7 @@ def gemini_flash_model(session, bq_connection) -> llm.GeminiTextGenerator: return llm.GeminiTextGenerator( session=session, connection_name=bq_connection, - model_name="gemini-2.5-flash", + model_name="gemini-2.5-pro", ) diff --git a/packages/bigframes/tests/system/small/bigquery/test_ai.py b/packages/bigframes/tests/system/small/bigquery/test_ai.py index f3c94edd1969..dd713e26dfa3 100644 --- a/packages/bigframes/tests/system/small/bigquery/test_ai.py +++ b/packages/bigframes/tests/system/small/bigquery/test_ai.py @@ -60,7 +60,7 @@ def test_ai_function_pandas_input(session): s2 = bpd.Series(["fruit", "tree"], session=session) prompt = (s1, " is a ", s2) - result = bbq.ai.generate_bool(prompt, endpoint="gemini-2.5-flash") + result = bbq.ai.generate_bool(prompt, endpoint="gemini-2.5-pro") assert _contains_no_nulls(result) assert result.dtype == pd.ArrowDtype( @@ -81,7 +81,7 @@ def test_ai_function_string_input(session): mock_get_session.return_value = session prompt = "Is apple a fruit?" - result = bbq.ai.generate_bool(prompt, endpoint="gemini-2.5-flash") + result = bbq.ai.generate_bool(prompt, endpoint="gemini-2.5-pro") assert _contains_no_nulls(result) assert result.dtype == pd.ArrowDtype( @@ -102,7 +102,7 @@ def test_ai_function_compile_model_params(session): model_params = {"generation_config": {"thinking_config": {"thinking_budget": 0}}} result = bbq.ai.generate_bool( - prompt, endpoint="gemini-2.5-flash", model_params=model_params + prompt, endpoint="gemini-2.5-pro", model_params=model_params ) assert _contains_no_nulls(result) @@ -121,7 +121,7 @@ def test_ai_generate(session): country = bpd.Series(["Japan", "Canada"], session=session) prompt = ("What's the capital city of ", country, "? one word only") - result = bbq.ai.generate(prompt, endpoint="gemini-2.5-flash") + result = bbq.ai.generate(prompt, endpoint="gemini-2.5-pro") assert _contains_no_nulls(result) assert result.dtype == pd.ArrowDtype( @@ -141,7 +141,7 @@ def test_ai_generate_with_output_schema(session): result = bbq.ai.generate( prompt, - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", output_schema={"population": "INT64", "is_in_north_america": "bool"}, ) @@ -165,7 +165,7 @@ def test_ai_generate_with_invalid_output_schema_raise_error(session): with pytest.raises(ValueError): bbq.ai.generate( prompt, - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", output_schema={"population": "INT64", "is_in_north_america": "JSON"}, ) @@ -175,7 +175,7 @@ def test_ai_generate_bool(session): s2 = bpd.Series(["fruit", "tree"], session=session) prompt = (s1, " is a ", s2) - result = bbq.ai.generate_bool(prompt, endpoint="gemini-2.5-flash") + result = bbq.ai.generate_bool(prompt, endpoint="gemini-2.5-pro") assert _contains_no_nulls(result) assert result.dtype == pd.ArrowDtype( @@ -216,7 +216,7 @@ def test_ai_generate_int(session): s = bpd.Series(["Cat"], session=session) prompt = ("How many legs does a ", s, " have?") - result = bbq.ai.generate_int(prompt, endpoint="gemini-2.5-flash") + result = bbq.ai.generate_int(prompt, endpoint="gemini-2.5-pro") assert _contains_no_nulls(result) assert result.dtype == pd.ArrowDtype( @@ -259,7 +259,7 @@ def test_ai_generate_double(session): s = bpd.Series(["Cat"], session=session) prompt = ("How many legs does a ", s, " have?") - result = bbq.ai.generate_double(prompt, endpoint="gemini-2.5-flash") + result = bbq.ai.generate_double(prompt, endpoint="gemini-2.5-pro") assert _contains_no_nulls(result) assert result.dtype == pd.ArrowDtype( diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_params/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_params/out.sql index 30542740a2dc..4120ce5b4174 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_params/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_params/out.sql @@ -3,7 +3,7 @@ SELECT input => STRUCT(`string_col`), categories => ['greeting', 'rejection'], examples => [('hi', 'greeting'), ('bye', 'rejection')], - endpoint => 'gemini-2.5-flash', + endpoint => 'gemini-2.5-pro', max_error_ratio => 0.1 ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate/out.sql index 622782fa7d65..605fd7d36ba0 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate/out.sql @@ -1,7 +1,7 @@ SELECT AI.GENERATE( prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), - endpoint => 'gemini-2.5-flash', + endpoint => 'gemini-2.5-pro', request_type => 'SHARED' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool/out.sql index a71bce037a5d..b54dec796cb6 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_BOOL( prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), - endpoint => 'gemini-2.5-flash' + endpoint => 'gemini-2.5-pro' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_connection_id/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_connection_id/out.sql index db1ec378aaf9..b0438c280b9a 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_connection_id/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_connection_id/out.sql @@ -2,6 +2,6 @@ SELECT AI.GENERATE_BOOL( prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), connection_id => 'bigframes-dev.us.bigframes-default-connection', - endpoint => 'gemini-2.5-flash' + endpoint => 'gemini-2.5-pro' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double/out.sql index 1cef75687988..46a583f2a9c3 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_DOUBLE( prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), - endpoint => 'gemini-2.5-flash' + endpoint => 'gemini-2.5-pro' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_connection_id/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_connection_id/out.sql index d0088721e386..01661dc34c54 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_connection_id/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_connection_id/out.sql @@ -2,6 +2,6 @@ SELECT AI.GENERATE_DOUBLE( prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), connection_id => 'bigframes-dev.us.bigframes-default-connection', - endpoint => 'gemini-2.5-flash' + endpoint => 'gemini-2.5-pro' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int/out.sql index 9ef143c8b9e4..d1ed914c4cb0 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int/out.sql @@ -1,6 +1,6 @@ SELECT AI.GENERATE_INT( prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), - endpoint => 'gemini-2.5-flash' + endpoint => 'gemini-2.5-pro' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_connection_id/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_connection_id/out.sql index 3fa3e8cc05e1..4276ab3d1243 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_connection_id/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_int_with_connection_id/out.sql @@ -2,6 +2,6 @@ SELECT AI.GENERATE_INT( prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), connection_id => 'bigframes-dev.us.bigframes-default-connection', - endpoint => 'gemini-2.5-flash' + endpoint => 'gemini-2.5-pro' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_connection_id/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_connection_id/out.sql index 14604cfc8dfd..c735bba03247 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_connection_id/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_connection_id/out.sql @@ -2,6 +2,6 @@ SELECT AI.GENERATE( prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), connection_id => 'bigframes-dev.us.bigframes-default-connection', - endpoint => 'gemini-2.5-flash' + endpoint => 'gemini-2.5-pro' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_output_schema/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_output_schema/out.sql index 31c179e7b01a..2377c9194ef6 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_output_schema/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_with_output_schema/out.sql @@ -1,7 +1,7 @@ SELECT AI.GENERATE( prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), - endpoint => 'gemini-2.5-flash', + endpoint => 'gemini-2.5-pro', output_schema => 'x INT64, y FLOAT64' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if_with_endpoint/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if_with_endpoint/out.sql index 4dd910528a41..6eb15118efd7 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if_with_endpoint/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if_with_endpoint/out.sql @@ -1,6 +1,6 @@ SELECT AI.IF( prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), - endpoint => 'gemini-2.5-flash' + endpoint => 'gemini-2.5-pro' ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score_with_endpoint_and_max_error_ratio/out.sql b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score_with_endpoint_and_max_error_ratio/out.sql index a802e5a396bf..48f919493ba1 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score_with_endpoint_and_max_error_ratio/out.sql +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score_with_endpoint_and_max_error_ratio/out.sql @@ -1,7 +1,7 @@ SELECT AI.SCORE( prompt => STRUCT(`string_col`, ' is the same as ', `string_col`), - endpoint => 'gemini-2.5-flash', + endpoint => 'gemini-2.5-pro', max_error_ratio => 0.5 ) AS `result` FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` \ No newline at end of file diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_ai_ops.py b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_ai_ops.py index 57c524908607..dfa71c9c4e6f 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_ai_ops.py +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_ai_ops.py @@ -30,7 +30,7 @@ def test_ai_generate(scalar_types_df: dataframe.DataFrame, snapshot): op = ops.AIGenerate( prompt_context=(None, " is the same as ", None), - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", request_type="SHARED", ) @@ -47,7 +47,7 @@ def test_ai_generate_with_connection_id(scalar_types_df: dataframe.DataFrame, sn op = ops.AIGenerate( prompt_context=(None, " is the same as ", None), connection_id=CONNECTION_ID, - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", ) sql = utils._apply_ops_to_sql( @@ -63,7 +63,7 @@ def test_ai_generate_with_output_schema(scalar_types_df: dataframe.DataFrame, sn op = ops.AIGenerate( prompt_context=(None, " is the same as ", None), connection_id=None, - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", output_schema="x INT64, y FLOAT64", ) @@ -94,7 +94,7 @@ def test_ai_generate_bool(scalar_types_df: dataframe.DataFrame, snapshot): op = ops.AIGenerateBool( prompt_context=(None, " is the same as ", None), - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", ) sql = utils._apply_ops_to_sql( @@ -112,7 +112,7 @@ def test_ai_generate_bool_with_connection_id( op = ops.AIGenerateBool( prompt_context=(None, " is the same as ", None), connection_id=CONNECTION_ID, - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", ) sql = utils._apply_ops_to_sql( @@ -145,7 +145,7 @@ def test_ai_generate_int(scalar_types_df: dataframe.DataFrame, snapshot): op = ops.AIGenerateInt( # The prompt does not make semantic sense but we only care about syntax correctness. prompt_context=(None, " is the same as ", None), - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", ) sql = utils._apply_ops_to_sql( @@ -164,7 +164,7 @@ def test_ai_generate_int_with_connection_id( # The prompt does not make semantic sense but we only care about syntax correctness. prompt_context=(None, " is the same as ", None), connection_id=CONNECTION_ID, - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", ) sql = utils._apply_ops_to_sql( @@ -198,7 +198,7 @@ def test_ai_generate_double(scalar_types_df: dataframe.DataFrame, snapshot): op = ops.AIGenerateDouble( # The prompt does not make semantic sense but we only care about syntax correctness. prompt_context=(None, " is the same as ", None), - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", ) sql = utils._apply_ops_to_sql( @@ -217,7 +217,7 @@ def test_ai_generate_double_with_connection_id( # The prompt does not make semantic sense but we only care about syntax correctness. prompt_context=(None, " is the same as ", None), connection_id=CONNECTION_ID, - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", ) sql = utils._apply_ops_to_sql( @@ -322,7 +322,7 @@ def test_ai_if_with_endpoint(scalar_types_df: dataframe.DataFrame, snapshot): op = ops.AIIf( prompt_context=(None, " is the same as ", None), - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", ) sql = utils._apply_ops_to_sql( @@ -354,7 +354,7 @@ def test_ai_classify_with_params(scalar_types_df: dataframe.DataFrame, snapshot) prompt_context=(None,), categories=("greeting", "rejection"), examples=(("hi", "greeting"), ("bye", "rejection")), - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", max_error_ratio=0.1, ) @@ -421,7 +421,7 @@ def test_ai_score_with_endpoint_and_max_error_ratio( op = ops.AIScore( prompt_context=(None, " is the same as ", None), - endpoint="gemini-2.5-flash", + endpoint="gemini-2.5-pro", max_error_ratio=0.5, ) From 16ca2604af0129a9d67e687d7840ff740a1f3c63 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Wed, 3 Jun 2026 18:46:32 +0000 Subject: [PATCH 2/5] ruff format --- packages/bigframes/bigframes/ml/llm.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/packages/bigframes/bigframes/ml/llm.py b/packages/bigframes/bigframes/ml/llm.py index d626f177169d..2810bbfee491 100644 --- a/packages/bigframes/bigframes/ml/llm.py +++ b/packages/bigframes/bigframes/ml/llm.py @@ -525,7 +525,12 @@ def _create_bqml_model(self): if ( self.model_name.startswith("gemini-2.0") or self.model_name.startswith("gemini-1.5") - or self.model_name in ("gemini-2.5-flash", "gemini-2.5-flash-lite", "gemini-2.5-pro-preview-05-06") + 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( From df218c1df4331c0cc7936d973a7d8e8750c17fb5 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Wed, 3 Jun 2026 18:58:37 +0000 Subject: [PATCH 3/5] simply logic --- packages/bigframes/bigframes/ml/llm.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/bigframes/bigframes/ml/llm.py b/packages/bigframes/bigframes/ml/llm.py index 2810bbfee491..26a1b4834dd9 100644 --- a/packages/bigframes/bigframes/ml/llm.py +++ b/packages/bigframes/bigframes/ml/llm.py @@ -523,8 +523,7 @@ def _create_bqml_model(self): ) warnings.warn(msg) if ( - self.model_name.startswith("gemini-2.0") - or self.model_name.startswith("gemini-1.5") + self.model_name.startswith(("gemini-2.0", "gemini-1.5")) or self.model_name in ( "gemini-2.5-flash", From 2e0757382343e7f6febc756c05e9124e7bf10d2a Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Wed, 3 Jun 2026 19:40:10 +0000 Subject: [PATCH 4/5] fix: gemini checks and nameless col to_frame bug --- packages/bigframes/bigframes/ml/llm.py | 14 ++-- packages/bigframes/bigframes/series.py | 10 ++- .../core/test_dataframe_accessor.py | 74 +++++++++---------- 3 files changed, 49 insertions(+), 49 deletions(-) diff --git a/packages/bigframes/bigframes/ml/llm.py b/packages/bigframes/bigframes/ml/llm.py index 26a1b4834dd9..55c38193ff07 100644 --- a/packages/bigframes/bigframes/ml/llm.py +++ b/packages/bigframes/bigframes/ml/llm.py @@ -522,14 +522,12 @@ def _create_bqml_model(self): ) ) warnings.warn(msg) - 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", - ) + 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( diff --git a/packages/bigframes/bigframes/series.py b/packages/bigframes/bigframes/series.py index d4e704591b01..6df6c333c121 100644 --- a/packages/bigframes/bigframes/series.py +++ b/packages/bigframes/bigframes/series.py @@ -2258,11 +2258,13 @@ def mask(self, cond, other=None) -> Series: return self.where(~cond, other) def to_frame(self, name: blocks.Label = None) -> bigframes.dataframe.DataFrame: - provided_name = name if name else self.name + provided_name = name if name is not None else self.name # To be consistent with Pandas, it assigns 0 as the column name if missing. 0 is the first element of RangeIndex. - block = self._block.with_column_labels( - [provided_name] if provided_name else [0] - ) + if provided_name is None or pandas.isna(provided_name): + column_names = [0] + else: + column_names = [provided_name] + block = self._block.with_column_labels(column_names) return bigframes.dataframe.DataFrame(block) def to_csv( diff --git a/packages/bigframes/tests/unit/extensions/core/test_dataframe_accessor.py b/packages/bigframes/tests/unit/extensions/core/test_dataframe_accessor.py index 914a448700f4..7ab4f5176980 100644 --- a/packages/bigframes/tests/unit/extensions/core/test_dataframe_accessor.py +++ b/packages/bigframes/tests/unit/extensions/core/test_dataframe_accessor.py @@ -96,16 +96,16 @@ def mock_generate(prompt, **kwargs): output_schema={"res": "STRING"}, ) - assert result == ( - df["text_input"], - { - "connection_id": "conn", - "endpoint": "endpoint", - "request_type": "dedicated", - "model_params": {"temp": 0.5}, - "output_schema": {"res": "STRING"}, - }, - ) + assert isinstance(result, tuple) + assert len(result) == 2 + pd.testing.assert_series_equal(result[0], df["text_input"]) + assert result[1] == { + "connection_id": "conn", + "endpoint": "endpoint", + "request_type": "dedicated", + "model_params": {"temp": 0.5}, + "output_schema": {"res": "STRING"}, + } def test_bigframes_ai_generate(scalar_types_df: bpd.DataFrame, monkeypatch): @@ -147,15 +147,15 @@ def mock_generate_bool(prompt, **kwargs): model_params={"temp": 0.5}, ) - assert result == ( - df["text_input"], - { - "connection_id": "conn", - "endpoint": "endpoint", - "request_type": "dedicated", - "model_params": {"temp": 0.5}, - }, - ) + assert isinstance(result, tuple) + assert len(result) == 2 + pd.testing.assert_series_equal(result[0], df["text_input"]) + assert result[1] == { + "connection_id": "conn", + "endpoint": "endpoint", + "request_type": "dedicated", + "model_params": {"temp": 0.5}, + } def test_bigframes_ai_generate_bool(scalar_types_df: bpd.DataFrame, monkeypatch): @@ -196,15 +196,15 @@ def mock_generate_int(prompt, **kwargs): model_params={"temp": 0.5}, ) - assert result == ( - df["text_input"], - { - "connection_id": "conn", - "endpoint": "endpoint", - "request_type": "dedicated", - "model_params": {"temp": 0.5}, - }, - ) + assert isinstance(result, tuple) + assert len(result) == 2 + pd.testing.assert_series_equal(result[0], df["text_input"]) + assert result[1] == { + "connection_id": "conn", + "endpoint": "endpoint", + "request_type": "dedicated", + "model_params": {"temp": 0.5}, + } def test_bigframes_ai_generate_int(scalar_types_df: bpd.DataFrame, monkeypatch): @@ -245,15 +245,15 @@ def mock_generate_double(prompt, **kwargs): model_params={"temp": 0.5}, ) - assert result == ( - df["text_input"], - { - "connection_id": "conn", - "endpoint": "endpoint", - "request_type": "dedicated", - "model_params": {"temp": 0.5}, - }, - ) + assert isinstance(result, tuple) + assert len(result) == 2 + pd.testing.assert_series_equal(result[0], df["text_input"]) + assert result[1] == { + "connection_id": "conn", + "endpoint": "endpoint", + "request_type": "dedicated", + "model_params": {"temp": 0.5}, + } def test_bigframes_ai_generate_double(scalar_types_df: bpd.DataFrame, monkeypatch): From a37898a7b04976c94ae7ab41bfce766fa4501da6 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Wed, 3 Jun 2026 23:25:30 +0000 Subject: [PATCH 5/5] revert: nameless column to_frame fix --- packages/bigframes/bigframes/series.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/packages/bigframes/bigframes/series.py b/packages/bigframes/bigframes/series.py index 73a819cdee91..ebf32ac7850d 100644 --- a/packages/bigframes/bigframes/series.py +++ b/packages/bigframes/bigframes/series.py @@ -2280,13 +2280,11 @@ def mask(self, cond, other=None) -> Series: return self.where(~cond, other) def to_frame(self, name: blocks.Label = None) -> bigframes.dataframe.DataFrame: - provided_name = name if name is not None else self.name + provided_name = name if name else self.name # To be consistent with Pandas, it assigns 0 as the column name if missing. 0 is the first element of RangeIndex. - if provided_name is None or pandas.isna(provided_name): - column_names = [0] - else: - column_names = [provided_name] - block = self._block.with_column_labels(column_names) + block = self._block.with_column_labels( + [provided_name] if provided_name else [0] + ) return bigframes.dataframe.DataFrame(block) def to_csv(