feat: GenAI - add labels

PiperOrigin-RevId: 685980434
googleapis · Oct 15, 2024 · 92c13dd · 92c13dd
1 parent 914c63c
commit 92c13dd
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 0 deletions.
diff --git a/tests/unit/vertexai/test_generative_models.py b/tests/unit/vertexai/test_generative_models.py
@@ -698,6 +698,34 @@ def test_generate_content(self, generative_models: generative_models):
         )
         assert response3.text
 
+        model4 = generative_models.GenerativeModel("gemini-1.5-pro-preview-0409")
+        response4 = model4.generate_content(
+            "Why is sky blue? Respond in JSON.",
+            generation_config=generative_models.GenerationConfig(
+                temperature=0.2,
+                top_p=0.9,
+                top_k=20,
+                candidate_count=1,
+                max_output_tokens=200,
+                stop_sequences=["\n\n\n"],
+                response_mime_type="application/json",
+            ),
+            safety_settings=[
+                generative_models.SafetySetting(
+                    category=generative_models.SafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+                    threshold=generative_models.SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+                    method=generative_models.SafetySetting.HarmBlockMethod.SEVERITY,
+                ),
+                generative_models.SafetySetting(
+                    category=generative_models.SafetySetting.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+                    threshold=generative_models.SafetySetting.HarmBlockThreshold.BLOCK_ONLY_HIGH,
+                    method=generative_models.SafetySetting.HarmBlockMethod.PROBABILITY,
+                ),
+            ],
+            labels={"label1": "value1", "label2": "value2"},
+        )
+        assert response4.text
+
     @mock.patch.object(
         target=prediction_service.PredictionServiceClient,
         attribute="generate_content",