From f0803d60c4302d4d1ca788e07fdb1901aa7fef28 Mon Sep 17 00:00:00 2001
From: Jason Schleifer <jason.schleifer@gmail.com>
Date: Wed, 12 Feb 2025 05:02:51 +1300
Subject: [PATCH] updated google prompt models to have gemini-2.0-flash,
 gemini-2.0-flash-lite-preview-02-05 updated google embedding model to
 text-embedding-004

---
 nodes/drivers/gtUIGoogleEmbeddingDriver.py | 2 +-
 nodes/drivers/gtUIGooglePromptDriver.py    | 9 ++++++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/nodes/drivers/gtUIGoogleEmbeddingDriver.py b/nodes/drivers/gtUIGoogleEmbeddingDriver.py
index ccc20be..8beca32 100644
--- a/nodes/drivers/gtUIGoogleEmbeddingDriver.py
+++ b/nodes/drivers/gtUIGoogleEmbeddingDriver.py
@@ -3,7 +3,7 @@ from griptape.drivers.embedding.google import GoogleEmbeddingDriver
 from .gtUIBaseEmbeddingDriver import gtUIBaseEmbeddingDriver
 
 models = [
-    "models/embedding-001",
+    "text-embedding-004",
 ]
 task_types = [
     "RETRIEVAL_QUERY",
diff --git a/nodes/drivers/gtUIGooglePromptDriver.py b/nodes/drivers/gtUIGooglePromptDriver.py
index 6d85b18..40ef104 100644
--- a/nodes/drivers/gtUIGooglePromptDriver.py
+++ b/nodes/drivers/gtUIGooglePromptDriver.py
@@ -2,7 +2,14 @@ from griptape.drivers.prompt.google import GooglePromptDriver
 
 from .gtUIBasePromptDriver import gtUIBasePromptDriver
 
-models = ["gemini-1.5-pro", "gemini-1.5-flash", "gemini-1.0-pro", "gemini-pro"]
+models = [
+    "gemini-2.0-flash",
+    "gemini-2.0-flash-lite-preview-02-05",
+    "gemini-1.5-pro",
+    "gemini-1.5-flash",
+    "gemini-1.0-pro",
+    "gemini-pro",
+]
 
 DEFAULT_API_KEY_ENV_VAR = "GOOGLE_API_KEY"