From 3b03f9adf3c3ca1c96d8a952d18cd429b05d5c51 Mon Sep 17 00:00:00 2001 From: Zhenyi Qi Date: Wed, 12 Jun 2024 19:57:18 -0700 Subject: [PATCH] chore: use the gemini-1.5-pro-001 for caching system test. This model is the officially supported model after the feature is launched. PiperOrigin-RevId: 642822908 --- tests/system/vertexai/test_generative_models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/system/vertexai/test_generative_models.py b/tests/system/vertexai/test_generative_models.py index 4fb9c877e0..abe042cdb1 100644 --- a/tests/system/vertexai/test_generative_models.py +++ b/tests/system/vertexai/test_generative_models.py @@ -34,7 +34,7 @@ GEMINI_MODEL_NAME = "gemini-1.0-pro-002" GEMINI_VISION_MODEL_NAME = "gemini-1.0-pro-vision" GEMINI_15_MODEL_NAME = "gemini-1.5-pro-preview-0409" -GEMINI_15_0514_MODEL_NAME = "gemini-1.5-pro-preview-0514" +GEMINI_15_PRO_MODEL_NAME = "gemini-1.5-pro-001" # A dummy function for function calling @@ -102,7 +102,7 @@ def setup_method(self): def test_generate_content_with_cached_content_from_text(self): cached_content = caching.CachedContent.create( - model_name=GEMINI_15_0514_MODEL_NAME, + model_name=GEMINI_15_PRO_MODEL_NAME, system_instruction="Please answer all the questions like a pirate.", contents=[ Content.from_dict(