Update Gemini1.5Pro to gemini-1.5-pro-002. The gemini-1.5-pro-latest …

…is depericated by VertexAI. Also add a default location. PiperOrigin-RevId: 724409415
google · Feb 7, 2025 · 6426ecc · 6426ecc
1 parent 537bb46
commit 6426ecc
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/langfun/core/llms/vertexai.py b/langfun/core/llms/vertexai.py
@@ -198,37 +198,43 @@ class VertexAIGeminiExp_20241114(VertexAIGemini):  # pylint: disable=invalid-nam
 class VertexAIGeminiPro1_5(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
-  model = 'gemini-1.5-pro-latest'
+  model = 'gemini-1.5-pro-002'
+  location = 'us-central1'
 
 
 class VertexAIGeminiPro1_5_002(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-002'
+  location = 'us-central1'
 
 
 class VertexAIGeminiPro1_5_001(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-001'
+  location = 'us-central1'
 
 
 class VertexAIGeminiFlash1_5(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash'
+  location = 'us-central1'
 
 
 class VertexAIGeminiFlash1_5_002(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-002'
+  location = 'us-central1'
 
 
 class VertexAIGeminiFlash1_5_001(VertexAIGemini):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-001'
+  location = 'us-central1'
 
 
 class VertexAIGeminiPro1(VertexAIGemini):  # pylint: disable=invalid-name