chore: Deflake test_generative_models system tests by allowing empty …

…text with finish reason

PiperOrigin-RevId: 636328795

tests/system/aiplatform/test_language_models.py

@@ -58,15 +58,16 @@ def test_text_generation(self, api_transport):
 
         model = TextGenerationModel.from_pretrained("google/text-bison@001")
         grounding_source = language_models.GroundingSource.WebSearch()
-        assert model.predict(
-            "What is the best recipe for banana bread? Recipe:",
+        response = model.predict(
+            "What is the best recipe for cupcakes? Recipe:",
             max_output_tokens=128,
             temperature=0.0,
             top_p=1.0,
             top_k=5,
             stop_sequences=["# %%"],
             grounding_source=grounding_source,
-        ).text
+        )
+        assert response.text or response.is_blocked
 
     @pytest.mark.parametrize("api_transport", ["grpc", "rest"])
     def test_text_generation_preview_count_tokens(self, api_transport):
@@ -97,15 +98,15 @@ async def test_text_generation_model_predict_async(self, api_transport):
         model = TextGenerationModel.from_pretrained("google/text-bison@001")
         grounding_source = language_models.GroundingSource.WebSearch()
         response = await model.predict_async(
-            "What is the best recipe for banana bread? Recipe:",
+            "What is the best recipe for cupcakes? Recipe:",
             max_output_tokens=128,
             temperature=0.0,
             top_p=1.0,
             top_k=5,
             stop_sequences=["# %%"],
             grounding_source=grounding_source,
         )
-        assert response.text
+        assert response.text or response.is_blocked
 
     @pytest.mark.parametrize("api_transport", ["grpc", "rest"])
     def test_text_generation_streaming(self, api_transport):
@@ -118,13 +119,13 @@ def test_text_generation_streaming(self, api_transport):
         model = TextGenerationModel.from_pretrained("google/text-bison@001")
 
         for response in model.predict_streaming(
-            "What is the best recipe for banana bread? Recipe:",
+            "What is the best recipe for cupcakes? Recipe:",
             max_output_tokens=128,
             temperature=0.0,
             top_p=1.0,
             top_k=5,
         ):
-            assert response.text
+            assert response.text or response.is_blocked
 
     @pytest.mark.parametrize("api_transport", ["grpc", "rest"])
     def test_preview_text_embedding_top_level_from_pretrained(self, api_transport):
@@ -138,14 +139,15 @@ def test_preview_text_embedding_top_level_from_pretrained(self, api_transport):
             foundation_model_name="google/text-bison@001"
         )
 
-        assert model.predict(
-            "What is the best recipe for banana bread? Recipe:",
+        response = model.predict(
+            "What is the best recipe for cupcakes? Recipe:",
             max_output_tokens=128,
             temperature=0.0,
             top_p=1.0,
             top_k=5,
             stop_sequences=["# %%"],
-        ).text
+        )
+        assert response.text or response.is_blocked
 
         assert isinstance(model, preview_language_models.TextGenerationModel)
 
@@ -430,13 +432,13 @@ def test_tuning(self, shared_state, api_transport):
 
         # Testing the new model returned by the `tuning_job.get_tuned_model` method
         response1 = tuned_model1.predict(
-            "What is the best recipe for banana bread? Recipe:",
+            "What is the best recipe for cupcakes? Recipe:",
             max_output_tokens=128,
             temperature=0.0,
             top_p=1.0,
             top_k=5,
         )
-        assert response1.text
+        assert response1.text or response1.is_blocked
 
         # Testing listing and getting tuned models
         tuned_model_names = model.list_tuned_model_names()
@@ -446,13 +448,13 @@ def test_tuning(self, shared_state, api_transport):
         tuned_model = TextGenerationModel.get_tuned_model(tuned_model_name)
 
         tuned_model_response = tuned_model.predict(
-            "What is the best recipe for banana bread? Recipe:",
+            "What is the best recipe for cupcakes? Recipe:",
             max_output_tokens=128,
             temperature=0.0,
             top_p=1.0,
             top_k=5,
         )
-        assert tuned_model_response.text
+        assert tuned_model_response.text or tuned_model_response.is_blocked
 
     @pytest.mark.parametrize("api_transport", ["grpc", "rest"])
     def test_batch_prediction_for_text_generation(self, api_transport):

tests/system/vertexai/test_generative_models.py

@@ -112,7 +112,11 @@ def test_generate_content_streaming(self):
             generation_config=generative_models.GenerationConfig(temperature=0),
         )
         for chunk in stream:
-            assert chunk.text
+            assert (
+                chunk.text
+                or chunk.candidates[0].finish_reason
+                is generative_models.FinishReason.STOP
+            )
 
     @pytest.mark.asyncio
     async def test_generate_content_streaming_async(self):
@@ -123,7 +127,11 @@ async def test_generate_content_streaming_async(self):
             generation_config=generative_models.GenerationConfig(temperature=0),
         )
         async for chunk in async_stream:
-            assert chunk.text
+            assert (
+                chunk.text
+                or chunk.candidates[0].finish_reason
+                is generative_models.FinishReason.STOP
+            )
 
     def test_generate_content_with_parameters(self):
         model = generative_models.GenerativeModel(
@@ -241,7 +249,11 @@ def test_grounding_google_search_retriever(self):
             tools=[google_search_retriever_tool],
             generation_config=generative_models.GenerationConfig(temperature=0),
         )
-        assert response.text
+        assert (
+            response.candidates[0].finish_reason
+            is generative_models.FinishReason.RECITATION
+            or response.text
+        )
 
     # Chat