From d218b29484d687b069df2f067e362432bccdd98f Mon Sep 17 00:00:00 2001
From: ghchinoy <ghchinoy@gmail.com>
Date: Sat, 25 May 2024 11:30:03 -0600
Subject: [PATCH] feat: adds GA models

---
 langfun/core/llms/vertexai.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/langfun/core/llms/vertexai.py b/langfun/core/llms/vertexai.py
index cf557714..e065c3bc 100644
--- a/langfun/core/llms/vertexai.py
+++ b/langfun/core/llms/vertexai.py
@@ -26,7 +26,9 @@
 SUPPORTED_MODELS_AND_SETTINGS = {
     'gemini-1.5-pro-preview-0514': pg.Dict(api='gemini', rpm=5),
     'gemini-1.5-pro-preview-0409': pg.Dict(api='gemini', rpm=5),
+    'gemini-1.5-pro-001': pg.Dict(api='gemini', rpm=5),
     'gemini-1.5-flash-preview-0514': pg.Dict(api='gemini', rpm=5),
+    'gemini-1.5-flash-001': pg.Dict(api='gemini', rpm=5),
     'gemini-1.0-pro': pg.Dict(api='gemini', rpm=300),
     'gemini-1.0-pro-vision': pg.Dict(api='gemini', rpm=100),
     # PaLM APIs.