fix: correct class name in system test

PiperOrigin-RevId: 579293600
googleapis · Nov 3, 2023 · b822b57 · b822b57
1 parent 1473e19
commit b822b57
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/system/aiplatform/test_language_models.py b/tests/system/aiplatform/test_language_models.py
@@ -50,7 +50,7 @@ def test_text_generation(self):
         aiplatform.init(project=e2e_base._PROJECT, location=e2e_base._LOCATION)
 
         model = TextGenerationModel.from_pretrained("google/text-bison@001")
-        grounding_source = language_models.WebSearchGroundingSource()
+        grounding_source = language_models.GroundingSource.WebSearch()
         assert model.predict(
             "What is the best recipe for banana bread? Recipe:",
             max_output_tokens=128,
@@ -78,7 +78,7 @@ async def test_text_generation_model_predict_async(self):
         aiplatform.init(project=e2e_base._PROJECT, location=e2e_base._LOCATION)
 
         model = TextGenerationModel.from_pretrained("google/text-bison@001")
-        grounding_source = language_models.WebSearchGroundingSource()
+        grounding_source = language_models.GroundingSource.WebSearch()
         response = await model.predict_async(
             "What is the best recipe for banana bread? Recipe:",
             max_output_tokens=128,