chore: Parameterize generative model end-to-end test with prod and staging API endpoints, setup and fetch staging endpoint environment variable

PiperOrigin-RevId: 661454748

Author: speedstorm1

.kokoro/build.sh

@@ -33,6 +33,9 @@ export GOOGLE_APPLICATION_CREDENTIALS=${KOKORO_GFILE_DIR}/service-account.json
 # Setup project id.
 export PROJECT_ID=$(cat "${KOKORO_GFILE_DIR}/project-id.json")
 
+# Setup staging endpoint.
+export STAGING_ENDPOINT=$(cat "${KOKORO_KEYSTORE_DIR}/73713_vertexai-staging-endpoint")
+
 # Remove old nox
 python3 -m pip uninstall --yes --quiet nox-automation
 

.kokoro/continuous/common.cfg

@@ -16,6 +16,16 @@ gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/google-cloud-python"
 # Use the trampoline script to run in docker.
 build_file: "python-aiplatform/.kokoro/trampoline.sh"
 
+# Fetch vertexai staging endpoint
+before_action {
+  fetch_keystore {
+    keystore_resource {
+      keystore_config_id: 73713
+      keyname: "vertexai-staging-endpoint"
+    }
+  }
+}
+
 # Configure the docker image for kokoro-trampoline.
 env_vars: {
     key: "TRAMPOLINE_IMAGE"

tests/system/vertexai/test_generative_models.py

@@ -18,6 +18,7 @@
 """System tests for generative models."""
 
 import json
+import os
 import pytest
 
 # Google imports
@@ -36,6 +37,9 @@
 GEMINI_15_MODEL_NAME = "gemini-1.5-pro-preview-0409"
 GEMINI_15_PRO_MODEL_NAME = "gemini-1.5-pro-001"
 
+STAGING_API_ENDPOINT = os.getenv("STAGING_ENDPOINT")
+PROD_API_ENDPOINT = None
+
 
 # A dummy function for function calling
 def get_current_weather(location: str, unit: str = "centigrade"):
@@ -84,12 +88,14 @@ def get_current_weather(location: str, unit: str = "centigrade"):
 }
 
 
+@pytest.mark.parametrize("api_endpoint", [STAGING_API_ENDPOINT, PROD_API_ENDPOINT])
 class TestGenerativeModels(e2e_base.TestEndToEnd):
     """System tests for generative models."""
 
     _temp_prefix = "temp_generative_models_test_"
 
-    def setup_method(self):
+    @pytest.fixture(scope="function", autouse=True)
+    def setup_method(self, api_endpoint):
         super().setup_method()
         credentials, _ = auth.default(
             scopes=["https://www.googleapis.com/auth/cloud-platform"]
@@ -98,9 +104,10 @@ def setup_method(self):
             project=e2e_base._PROJECT,
             location=e2e_base._LOCATION,
             credentials=credentials,
+            api_endpoint=api_endpoint,
         )
 
-    def test_generate_content_with_cached_content_from_text(self):
+    def test_generate_content_with_cached_content_from_text(self, api_endpoint):
         cached_content = caching.CachedContent.create(
             model_name=GEMINI_15_PRO_MODEL_NAME,
             system_instruction="Please answer all the questions like a pirate.",
@@ -138,7 +145,7 @@ def test_generate_content_with_cached_content_from_text(self):
         finally:
             cached_content.delete()
 
-    def test_generate_content_from_text(self):
+    def test_generate_content_from_text(self, api_endpoint):
         model = generative_models.GenerativeModel(GEMINI_MODEL_NAME)
         response = model.generate_content(
             "Why is sky blue?",
@@ -147,15 +154,15 @@ def test_generate_content_from_text(self):
         assert response.text
 
     @pytest.mark.asyncio
-    async def test_generate_content_async(self):
+    async def test_generate_content_async(self, api_endpoint):
         model = generative_models.GenerativeModel(GEMINI_MODEL_NAME)
         response = await model.generate_content_async(
             "Why is sky blue?",
             generation_config=generative_models.GenerationConfig(temperature=0),
         )
         assert response.text
 
-    def test_generate_content_streaming(self):
+    def test_generate_content_streaming(self, api_endpoint):
         model = generative_models.GenerativeModel(GEMINI_MODEL_NAME)
         stream = model.generate_content(
             "Why is sky blue?",
@@ -170,7 +177,7 @@ def test_generate_content_streaming(self):
             )
 
     @pytest.mark.asyncio
-    async def test_generate_content_streaming_async(self):
+    async def test_generate_content_streaming_async(self, api_endpoint):
         model = generative_models.GenerativeModel(GEMINI_MODEL_NAME)
         async_stream = await model.generate_content_async(
             "Why is sky blue?",
@@ -184,7 +191,7 @@ async def test_generate_content_streaming_async(self):
                 is generative_models.FinishReason.STOP
             )
 
-    def test_generate_content_with_parameters(self):
+    def test_generate_content_with_parameters(self, api_endpoint):
         model = generative_models.GenerativeModel(
             GEMINI_MODEL_NAME,
             system_instruction=[
@@ -211,7 +218,7 @@ def test_generate_content_with_parameters(self):
         )
         assert response.text
 
-    def test_generate_content_with_gemini_15_parameters(self):
+    def test_generate_content_with_gemini_15_parameters(self, api_endpoint):
         model = generative_models.GenerativeModel(GEMINI_15_MODEL_NAME)
         response = model.generate_content(
             contents="Why is sky blue? Respond in JSON Format.",
@@ -237,7 +244,7 @@ def test_generate_content_with_gemini_15_parameters(self):
         assert response.text
         assert json.loads(response.text)
 
-    def test_generate_content_from_list_of_content_dict(self):
+    def test_generate_content_from_list_of_content_dict(self, api_endpoint):
         model = generative_models.GenerativeModel(GEMINI_MODEL_NAME)
         response = model.generate_content(
             contents=[{"role": "user", "parts": [{"text": "Why is sky blue?"}]}],
@@ -248,7 +255,7 @@ def test_generate_content_from_list_of_content_dict(self):
     @pytest.mark.skip(
         reason="Breaking change in the gemini-pro-vision model. See b/315803556#comment3"
     )
-    def test_generate_content_from_remote_image(self):
+    def test_generate_content_from_remote_image(self, api_endpoint):
         vision_model = generative_models.GenerativeModel(GEMINI_VISION_MODEL_NAME)
         image_part = generative_models.Part.from_uri(
             uri="gs://download.tensorflow.org/example_images/320px-Felis_catus-cat_on_snow.jpg",
@@ -261,7 +268,7 @@ def test_generate_content_from_remote_image(self):
         assert response.text
         assert "cat" in response.text
 
-    def test_generate_content_from_text_and_remote_image(self):
+    def test_generate_content_from_text_and_remote_image(self, api_endpoint):
         vision_model = generative_models.GenerativeModel(GEMINI_VISION_MODEL_NAME)
         image_part = generative_models.Part.from_uri(
             uri="gs://download.tensorflow.org/example_images/320px-Felis_catus-cat_on_snow.jpg",
@@ -274,7 +281,7 @@ def test_generate_content_from_text_and_remote_image(self):
         assert response.text
         assert "cat" in response.text
 
-    def test_generate_content_from_text_and_remote_video(self):
+    def test_generate_content_from_text_and_remote_video(self, api_endpoint):
         vision_model = generative_models.GenerativeModel(GEMINI_VISION_MODEL_NAME)
         video_part = generative_models.Part.from_uri(
             uri="gs://cloud-samples-data/video/animals.mp4",
@@ -287,7 +294,7 @@ def test_generate_content_from_text_and_remote_video(self):
         assert response.text
         assert "Zootopia" in response.text
 
-    def test_grounding_google_search_retriever(self):
+    def test_grounding_google_search_retriever(self, api_endpoint):
         model = preview_generative_models.GenerativeModel(GEMINI_MODEL_NAME)
         google_search_retriever_tool = (
             preview_generative_models.Tool.from_google_search_retrieval(
@@ -309,7 +316,7 @@ def test_grounding_google_search_retriever(self):
 
     # Chat
 
-    def test_send_message_from_text(self):
+    def test_send_message_from_text(self, api_endpoint):
         model = generative_models.GenerativeModel(GEMINI_MODEL_NAME)
         chat = model.start_chat()
         response1 = chat.send_message(
@@ -326,7 +333,7 @@ def test_send_message_from_text(self):
         assert response2.text
         assert len(chat.history) == 4
 
-    def test_chat_function_calling(self):
+    def test_chat_function_calling(self, api_endpoint):
         get_current_weather_func = generative_models.FunctionDeclaration(
             name="get_current_weather",
             description="Get the current weather in a given location",
@@ -360,7 +367,7 @@ def test_chat_function_calling(self):
         )
         assert response2.text
 
-    def test_generate_content_function_calling(self):
+    def test_generate_content_function_calling(self, api_endpoint):
         get_current_weather_func = generative_models.FunctionDeclaration(
             name="get_current_weather",
             description="Get the current weather in a given location",
@@ -440,7 +447,7 @@ def test_generate_content_function_calling(self):
 
         assert summary
 
-    def test_chat_automatic_function_calling(self):
+    def test_chat_automatic_function_calling(self, api_endpoint):
         get_current_weather_func = generative_models.FunctionDeclaration.from_func(
             get_current_weather
         )
@@ -471,7 +478,7 @@ def test_chat_automatic_function_calling(self):
         assert chat.history[-2].parts[0].function_response
         assert chat.history[-2].parts[0].function_response.name == "get_current_weather"
 
-    def test_additional_request_metadata(self):
+    def test_additional_request_metadata(self, api_endpoint):
         aiplatform.init(request_metadata=[("foo", "bar")])
         model = generative_models.GenerativeModel(GEMINI_MODEL_NAME)
         response = model.generate_content(
@@ -480,7 +487,7 @@ def test_additional_request_metadata(self):
         )
         assert response
 
-    def test_compute_tokens_from_text(self):
+    def test_compute_tokens_from_text(self, api_endpoint):
         model = generative_models.GenerativeModel(GEMINI_MODEL_NAME)
         response = model.compute_tokens(["Why is sky blue?", "Explain it like I'm 5."])
         assert len(response.tokens_info) == 2