fix(bigframes): update GeminiTextGenerator default model to gemini-2.5-flash - #18060
fix(bigframes): update GeminiTextGenerator default model to gemini-2.5-flash#18060shuoweil wants to merge 4 commits into
Conversation
There was a problem hiding this comment.
Code Review
This pull request updates the default model for GeminiTextGenerator from 'gemini-2.0-flash-001' to 'gemini-2.5-flash' in the bigframes library. It also introduces unit tests to verify the default model assignment and error handling for unsupported models. There are no review comments, and I have no additional feedback to provide.
|
Switching to draft since tests are failing. Please feel free to move it back to |
| def test_gemini_text_generator_unsupported_model_error(): | ||
| # Create a mock session | ||
| mock_session = mock.create_autospec(spec=bigframes.session.Session) | ||
|
|
||
| # Mock _create_bq_connection to return a dummy connection | ||
| mock_session._create_bq_connection.return_value = ( | ||
| "projects/test-project/locations/us-central1/connections/test-conn" | ||
| ) | ||
|
|
||
| # Mock _anonymous_dataset which is used to create the temporary model reference | ||
| mock_session._anonymous_dataset = bigquery.DatasetReference( | ||
| "test-project", "test_dataset" | ||
| ) | ||
|
|
||
| # Mock _start_query_ml_ddl to raise BadRequest (simulating BQML failure) | ||
| error_message = ( | ||
| "Unsupported endpoint: Publisher model " | ||
| "projects/296675019294/locations/us-central1/publishers/google/models/gemini-3.5-flash " | ||
| "was not found or your project does not have access to it." | ||
| ) | ||
| bq_error = google.api_core.exceptions.BadRequest(error_message) | ||
| mock_session._start_query_ml_ddl.side_effect = bq_error | ||
|
|
||
| # Attempting to create the model should raise the BadRequest exception | ||
| with pytest.raises(google.api_core.exceptions.BadRequest) as exc_info: | ||
| llm.GeminiTextGenerator( | ||
| model_name="gemini-3.5-flash", | ||
| session=mock_session, | ||
| connection_name="test-conn", | ||
| ) | ||
|
|
||
| assert error_message in str(exc_info.value) | ||
|
|
||
| # Verify that the session's DDL execution method was called | ||
| mock_session._start_query_ml_ddl.assert_called_once() | ||
| generated_sql = mock_session._start_query_ml_ddl.call_args[0][0] | ||
| assert "CREATE OR REPLACE MODEL" in generated_sql | ||
| assert "gemini-3.5-flash" in generated_sql | ||
| assert "test-conn" in generated_sql |
There was a problem hiding this comment.
I think this model validation logic is not handled by our code, right? If so, then we should probably not add test coverage for it.
There was a problem hiding this comment.
Good point, removed the redundant error-handling test case and kept only the default model resolution test.
| log_adapter.add_api_method("dataframe-max", session=session) | ||
| for _ in range(52): | ||
| df.head() | ||
| log_adapter.add_api_method("dataframe-head", session=session) |
There was a problem hiding this comment.
Hmmm why do we need this? It's not related to the default model change, right?
There was a problem hiding this comment.
Reverted these changes to keep this PR focused purely on the default model update. I start a new branch to fix this.
| job_config.labels = cur_labels | ||
|
|
||
| df.max() | ||
| log_adapter.add_api_method("dataframe-max", session=session) |
There was a problem hiding this comment.
Similar question here too: does the default model change break this test?
There was a problem hiding this comment.
Reverted these changes to keep this PR focused purely on the default model update. I start a new branch to fix it.
1e9393b to
89d87b2
Compare
Fixes #<544873054> 🦕