Skip to content

Commit 7094c85

Browse files
authored
feat: add gemini-1.5-pro-002 and gemini-1.5-flash-002 to known Gemini model list. (#1105)
1 parent bf93e80 commit 7094c85

File tree

3 files changed

+17
-1
lines changed

3 files changed

+17
-1
lines changed

bigframes/ml/llm.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,13 +58,17 @@
5858
_GEMINI_1P5_PRO_PREVIEW_ENDPOINT = "gemini-1.5-pro-preview-0514"
5959
_GEMINI_1P5_PRO_FLASH_PREVIEW_ENDPOINT = "gemini-1.5-flash-preview-0514"
6060
_GEMINI_1P5_PRO_001_ENDPOINT = "gemini-1.5-pro-001"
61+
_GEMINI_1P5_PRO_002_ENDPOINT = "gemini-1.5-pro-002"
6162
_GEMINI_1P5_FLASH_001_ENDPOINT = "gemini-1.5-flash-001"
63+
_GEMINI_1P5_FLASH_002_ENDPOINT = "gemini-1.5-flash-002"
6264
_GEMINI_ENDPOINTS = (
6365
_GEMINI_PRO_ENDPOINT,
6466
_GEMINI_1P5_PRO_PREVIEW_ENDPOINT,
6567
_GEMINI_1P5_PRO_FLASH_PREVIEW_ENDPOINT,
6668
_GEMINI_1P5_PRO_001_ENDPOINT,
69+
_GEMINI_1P5_PRO_002_ENDPOINT,
6770
_GEMINI_1P5_FLASH_001_ENDPOINT,
71+
_GEMINI_1P5_FLASH_002_ENDPOINT,
6872
)
6973

7074
_CLAUDE_3_SONNET_ENDPOINT = "claude-3-sonnet"
@@ -749,7 +753,7 @@ class GeminiTextGenerator(base.BaseEstimator):
749753
750754
Args:
751755
model_name (str, Default to "gemini-pro"):
752-
The model for natural language tasks. Accepted values are "gemini-pro", "gemini-1.5-pro-preview-0514", "gemini-1.5-flash-preview-0514", "gemini-1.5-pro-001" and "gemini-1.5-flash-001". Default to "gemini-pro".
756+
The model for natural language tasks. Accepted values are "gemini-pro", "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" and "gemini-1.5-flash-002". Default to "gemini-pro".
753757
754758
.. note::
755759
"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
@@ -775,7 +779,9 @@ def __init__(
775779
"gemini-1.5-pro-preview-0514",
776780
"gemini-1.5-flash-preview-0514",
777781
"gemini-1.5-pro-001",
782+
"gemini-1.5-pro-002",
778783
"gemini-1.5-flash-001",
784+
"gemini-1.5-flash-002",
779785
] = "gemini-pro",
780786
session: Optional[bigframes.Session] = None,
781787
connection_name: Optional[str] = None,

bigframes/ml/loader.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,9 @@
6464
llm._GEMINI_1P5_PRO_PREVIEW_ENDPOINT: llm.GeminiTextGenerator,
6565
llm._GEMINI_1P5_PRO_FLASH_PREVIEW_ENDPOINT: llm.GeminiTextGenerator,
6666
llm._GEMINI_1P5_PRO_001_ENDPOINT: llm.GeminiTextGenerator,
67+
llm._GEMINI_1P5_PRO_002_ENDPOINT: llm.GeminiTextGenerator,
6768
llm._GEMINI_1P5_FLASH_001_ENDPOINT: llm.GeminiTextGenerator,
69+
llm._GEMINI_1P5_FLASH_002_ENDPOINT: llm.GeminiTextGenerator,
6870
llm._CLAUDE_3_HAIKU_ENDPOINT: llm.Claude3TextGenerator,
6971
llm._CLAUDE_3_SONNET_ENDPOINT: llm.Claude3TextGenerator,
7072
llm._CLAUDE_3_5_SONNET_ENDPOINT: llm.Claude3TextGenerator,

tests/system/small/ml/test_llm.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -264,7 +264,9 @@ def test_text_embedding_generator_multi_cols_predict_success(
264264
"gemini-1.5-pro-preview-0514",
265265
"gemini-1.5-flash-preview-0514",
266266
"gemini-1.5-pro-001",
267+
"gemini-1.5-pro-002",
267268
"gemini-1.5-flash-001",
269+
"gemini-1.5-flash-002",
268270
),
269271
)
270272
def test_create_load_gemini_text_generator_model(
@@ -292,7 +294,9 @@ def test_create_load_gemini_text_generator_model(
292294
"gemini-1.5-pro-preview-0514",
293295
"gemini-1.5-flash-preview-0514",
294296
"gemini-1.5-pro-001",
297+
"gemini-1.5-pro-002",
295298
"gemini-1.5-flash-001",
299+
"gemini-1.5-flash-002",
296300
),
297301
)
298302
@pytest.mark.flaky(retries=2)
@@ -315,7 +319,9 @@ def test_gemini_text_generator_predict_default_params_success(
315319
"gemini-1.5-pro-preview-0514",
316320
"gemini-1.5-flash-preview-0514",
317321
"gemini-1.5-pro-001",
322+
"gemini-1.5-pro-002",
318323
"gemini-1.5-flash-001",
324+
"gemini-1.5-flash-002",
319325
),
320326
)
321327
@pytest.mark.flaky(retries=2)
@@ -340,7 +346,9 @@ def test_gemini_text_generator_predict_with_params_success(
340346
"gemini-1.5-pro-preview-0514",
341347
"gemini-1.5-flash-preview-0514",
342348
"gemini-1.5-pro-001",
349+
"gemini-1.5-pro-002",
343350
"gemini-1.5-flash-001",
351+
"gemini-1.5-flash-002",
344352
),
345353
)
346354
@pytest.mark.flaky(retries=2)

0 commit comments

Comments
 (0)