@@ -698,6 +698,34 @@ def test_generate_content(self, generative_models: generative_models):
698
698
)
699
699
assert response3 .text
700
700
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
+
701
729
@mock .patch .object (
702
730
target = prediction_service .PredictionServiceClient ,
703
731
attribute = "generate_content" ,
0 commit comments