Skip to content

Commit 92c13dd

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: GenAI - add labels
PiperOrigin-RevId: 685980434
1 parent 914c63c commit 92c13dd

File tree

2 files changed

+91
-0
lines changed

2 files changed

+91
-0
lines changed

tests/unit/vertexai/test_generative_models.py

+28
Original file line numberDiff line numberDiff line change
@@ -698,6 +698,34 @@ def test_generate_content(self, generative_models: generative_models):
698698
)
699699
assert response3.text
700700

701+
model4 = generative_models.GenerativeModel("gemini-1.5-pro-preview-0409")
702+
response4 = model4.generate_content(
703+
"Why is sky blue? Respond in JSON.",
704+
generation_config=generative_models.GenerationConfig(
705+
temperature=0.2,
706+
top_p=0.9,
707+
top_k=20,
708+
candidate_count=1,
709+
max_output_tokens=200,
710+
stop_sequences=["\n\n\n"],
711+
response_mime_type="application/json",
712+
),
713+
safety_settings=[
714+
generative_models.SafetySetting(
715+
category=generative_models.SafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
716+
threshold=generative_models.SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
717+
method=generative_models.SafetySetting.HarmBlockMethod.SEVERITY,
718+
),
719+
generative_models.SafetySetting(
720+
category=generative_models.SafetySetting.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
721+
threshold=generative_models.SafetySetting.HarmBlockThreshold.BLOCK_ONLY_HIGH,
722+
method=generative_models.SafetySetting.HarmBlockMethod.PROBABILITY,
723+
),
724+
],
725+
labels={"label1": "value1", "label2": "value2"},
726+
)
727+
assert response4.text
728+
701729
@mock.patch.object(
702730
target=prediction_service.PredictionServiceClient,
703731
attribute="generate_content",

0 commit comments

Comments
 (0)