From 1021d5761a291f2327fc10216e938826e53dbcc4 Mon Sep 17 00:00:00 2001
From: Garrett Wu <6505921+GarrettWu@users.noreply.github.com>
Date: Wed, 11 Dec 2024 12:43:47 -0800
Subject: [PATCH] feat: add Gemini 2.0 preview text model support (#1209)

* feat: add Gemini 2.0 preview text model support

* update docs
---
 bigframes/ml/llm.py    | 20 ++++++++++++++++++--
 bigframes/ml/loader.py |  1 +
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py
index 427e99583d..dca7e555f6 100644
--- a/bigframes/ml/llm.py
+++ b/bigframes/ml/llm.py
@@ -63,6 +63,7 @@
 _GEMINI_1P5_PRO_002_ENDPOINT = "gemini-1.5-pro-002"
 _GEMINI_1P5_FLASH_001_ENDPOINT = "gemini-1.5-flash-001"
 _GEMINI_1P5_FLASH_002_ENDPOINT = "gemini-1.5-flash-002"
+_GEMINI_2_FLASH_EXP_ENDPOINT = "gemini-2.0-flash-exp"
 _GEMINI_ENDPOINTS = (
     _GEMINI_PRO_ENDPOINT,
     _GEMINI_1P5_PRO_PREVIEW_ENDPOINT,
@@ -71,6 +72,12 @@
     _GEMINI_1P5_PRO_002_ENDPOINT,
     _GEMINI_1P5_FLASH_001_ENDPOINT,
     _GEMINI_1P5_FLASH_002_ENDPOINT,
+    _GEMINI_2_FLASH_EXP_ENDPOINT,
+)
+_GEMINI_PREVIEW_ENDPOINTS = (
+    _GEMINI_1P5_PRO_PREVIEW_ENDPOINT,
+    _GEMINI_1P5_PRO_FLASH_PREVIEW_ENDPOINT,
+    _GEMINI_2_FLASH_EXP_ENDPOINT,
 )
 
 _CLAUDE_3_SONNET_ENDPOINT = "claude-3-sonnet"
@@ -757,10 +764,10 @@ class GeminiTextGenerator(base.BaseEstimator):
 
     Args:
         model_name (str, Default to "gemini-pro"):
-            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".
+            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", "gemini-1.5-flash-002" and "gemini-2.0-flash-exp". Default to "gemini-pro".
 
         .. note::
-            "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
+            "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
             Service Specific Terms(https://cloud.google.com/terms/service-terms#1). Pre-GA products and features are available "as is"
             and might have limited support. For more information, see the launch stage descriptions
             (https://cloud.google.com/products#product-launch-stages).
@@ -786,11 +793,20 @@ def __init__(
             "gemini-1.5-pro-002",
             "gemini-1.5-flash-001",
             "gemini-1.5-flash-002",
+            "gemini-2.0-flash-exp",
         ] = "gemini-pro",
         session: Optional[bigframes.Session] = None,
         connection_name: Optional[str] = None,
         max_iterations: int = 300,
     ):
+        if model_name in _GEMINI_PREVIEW_ENDPOINTS:
+            warnings.warn(
+                f"""Model {model_name} is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the
+            Service Specific Terms(https://cloud.google.com/terms/service-terms#1). Pre-GA products and features are available "as is"
+            and might have limited support. For more information, see the launch stage descriptions
+            (https://cloud.google.com/products#product-launch-stages).""",
+                category=exceptions.PreviewWarning,
+            )
         self.model_name = model_name
         self.session = session or bpd.get_global_session()
         self.max_iterations = max_iterations
diff --git a/bigframes/ml/loader.py b/bigframes/ml/loader.py
index 6a14fb3451..1cf8dc8a53 100644
--- a/bigframes/ml/loader.py
+++ b/bigframes/ml/loader.py
@@ -67,6 +67,7 @@
         llm._GEMINI_1P5_PRO_002_ENDPOINT: llm.GeminiTextGenerator,
         llm._GEMINI_1P5_FLASH_001_ENDPOINT: llm.GeminiTextGenerator,
         llm._GEMINI_1P5_FLASH_002_ENDPOINT: llm.GeminiTextGenerator,
+        llm._GEMINI_2_FLASH_EXP_ENDPOINT: llm.GeminiTextGenerator,
         llm._CLAUDE_3_HAIKU_ENDPOINT: llm.Claude3TextGenerator,
         llm._CLAUDE_3_SONNET_ENDPOINT: llm.Claude3TextGenerator,
         llm._CLAUDE_3_5_SONNET_ENDPOINT: llm.Claude3TextGenerator,