16
16
import logging
17
17
import sys
18
18
import typing
19
+
19
20
import pydantic
20
21
import pytest
21
22
@@ -120,7 +121,9 @@ def divide_floats(a: float, b: float) -> float:
120
121
'tools' : [{
121
122
'retrieval' : {
122
123
'vertex_ai_search' : {
123
- 'datastore' : 'projects/vertex-sdk-dev/locations/global/collections/default_collection/dataStores/yvonne_1728691676574'
124
+ 'datastore' : (
125
+ 'projects/vertex-sdk-dev/locations/global/collections/default_collection/dataStores/yvonne_1728691676574'
126
+ )
124
127
}
125
128
}
126
129
}]
@@ -149,7 +152,7 @@ def divide_floats(a: float, b: float) -> float:
149
152
exception_if_mldev = 'retrieval' ,
150
153
exception_if_vertex = '400' ,
151
154
),
152
- pytest_helper .TestTableItem (
155
+ pytest_helper .TestTableItem (
153
156
name = 'test_vai_search_engine' ,
154
157
parameters = types ._GenerateContentParameters (
155
158
model = 'gemini-2.0-flash-001' ,
@@ -258,6 +261,31 @@ def divide_floats(a: float, b: float) -> float:
258
261
config = {'tools' : [{'code_execution' : {}}]},
259
262
),
260
263
),
264
+ pytest_helper .TestTableItem (
265
+ name = 'test_function_google_search_retrieval_with_long_lat' ,
266
+ parameters = types ._GenerateContentParameters (
267
+ model = 'gemini-1.5-flash' ,
268
+ contents = t .t_contents (None , 'what is the price of GOOG?' ),
269
+ config = types .GenerateContentConfig (
270
+ tools = [
271
+ types .Tool (
272
+ google_search_retrieval = types .GoogleSearchRetrieval (
273
+ dynamic_retrieval_config = types .DynamicRetrievalConfig (
274
+ mode = 'MODE_UNSPECIFIED'
275
+ )
276
+ )
277
+ ),
278
+ ],
279
+ tool_config = types .ToolConfig (
280
+ retrieval_config = types .RetrievalConfig (
281
+ lat_lng = types .LatLngDict (
282
+ latitude = 37.7749 , longitude = - 122.4194
283
+ )
284
+ )
285
+ ),
286
+ ),
287
+ ),
288
+ ),
261
289
]
262
290
263
291
@@ -464,7 +492,9 @@ async def test_disable_automatic_function_calling_stream_async(client):
464
492
465
493
466
494
@pytest .mark .asyncio
467
- async def test_automatic_function_calling_no_function_response_stream_async (client ):
495
+ async def test_automatic_function_calling_no_function_response_stream_async (
496
+ client ,
497
+ ):
468
498
response = await client .aio .models .generate_content_stream (
469
499
model = 'gemini-1.5-flash' ,
470
500
contents = 'what is the weather in Boston?' ,
@@ -658,7 +688,7 @@ def get_weather_pydantic_model(
658
688
contents = 'it is winter now, what is the weather in Boston?' ,
659
689
config = {
660
690
'tools' : [get_weather_pydantic_model ],
661
- 'automatic_function_calling' : {'ignore_call_history' : True }
691
+ 'automatic_function_calling' : {'ignore_call_history' : True },
662
692
},
663
693
)
664
694
@@ -689,8 +719,9 @@ def get_weather_from_list_of_cities(
689
719
response = client .models .generate_content (
690
720
model = 'gemini-1.5-flash' ,
691
721
contents = 'it is winter now, what is the weather in Boston and New York?' ,
692
- config = {'tools' : [get_weather_from_list_of_cities ],
693
- 'automatic_function_calling' : {'ignore_call_history' : True },
722
+ config = {
723
+ 'tools' : [get_weather_from_list_of_cities ],
724
+ 'automatic_function_calling' : {'ignore_call_history' : True },
694
725
},
695
726
)
696
727
@@ -736,30 +767,34 @@ def get_information(
736
767
),
737
768
config = {
738
769
'tools' : [get_information ],
739
- 'automatic_function_calling' : {'ignore_call_history' : True }
770
+ 'automatic_function_calling' : {'ignore_call_history' : True },
740
771
},
741
772
)
742
773
assert 'Sundae' in response .text
743
774
assert 'cat' in response .text
744
775
745
776
746
- def test_automatic_function_calling_with_parameterized_generic_union_type (client ):
777
+ def test_automatic_function_calling_with_parameterized_generic_union_type (
778
+ client ,
779
+ ):
747
780
def describe_cities (
748
781
country : str ,
749
782
cities : typing .Optional [list [str ]] = None ,
750
783
) -> str :
751
- " Given a country and an optional list of cities, describe the cities."
784
+ ' Given a country and an optional list of cities, describe the cities.'
752
785
if cities is None :
753
786
return 'There are no cities to describe.'
754
787
else :
755
- return f'The cities in { country } are: { ", " .join (cities )} and they are nice.'
788
+ return (
789
+ f'The cities in { country } are: { ", " .join (cities )} and they are nice.'
790
+ )
756
791
757
792
response = client .models .generate_content (
758
793
model = 'gemini-1.5-flash' ,
759
- contents = ( 'Can you describe the city of San Francisco?' ) ,
794
+ contents = 'Can you describe the city of San Francisco?' ,
760
795
config = {
761
796
'tools' : [describe_cities ],
762
- 'automatic_function_calling' : {'ignore_call_history' : True }
797
+ 'automatic_function_calling' : {'ignore_call_history' : True },
763
798
},
764
799
)
765
800
assert 'San Francisco' in response .text
@@ -794,7 +829,7 @@ def test_with_1_empty_tool(client):
794
829
contents = 'What is the price of GOOG?.' ,
795
830
config = {
796
831
'tools' : [{}, get_stock_price ],
797
- 'automatic_function_calling' : {'ignore_call_history' : True }
832
+ 'automatic_function_calling' : {'ignore_call_history' : True },
798
833
},
799
834
)
800
835
@@ -819,7 +854,9 @@ async def test_vai_search_stream_async(client):
819
854
'tools' : [{
820
855
'retrieval' : {
821
856
'vertex_ai_search' : {
822
- 'datastore' : 'projects/vertex-sdk-dev/locations/global/collections/default_collection/dataStores/yvonne_1728691676574'
857
+ 'datastore' : (
858
+ 'projects/vertex-sdk-dev/locations/global/collections/default_collection/dataStores/yvonne_1728691676574'
859
+ )
823
860
}
824
861
}
825
862
}]
@@ -835,7 +872,9 @@ async def test_vai_search_stream_async(client):
835
872
'tools' : [{
836
873
'retrieval' : {
837
874
'vertex_ai_search' : {
838
- 'datastore' : 'projects/vertex-sdk-dev/locations/global/collections/default_collection/dataStores/yvonne_1728691676574'
875
+ 'datastore' : (
876
+ 'projects/vertex-sdk-dev/locations/global/collections/default_collection/dataStores/yvonne_1728691676574'
877
+ )
839
878
}
840
879
}
841
880
}]
@@ -855,7 +894,7 @@ async def divide_integers(a: int, b: int) -> int:
855
894
contents = 'what is the result of 1000/2?' ,
856
895
config = {
857
896
'tools' : [divide_integers ],
858
- 'automatic_function_calling' : {'ignore_call_history' : True }
897
+ 'automatic_function_calling' : {'ignore_call_history' : True },
859
898
},
860
899
)
861
900
@@ -868,13 +907,13 @@ async def divide_integers(a: int, b: int) -> int:
868
907
return a // b
869
908
870
909
response = await client .aio .models .generate_content (
871
- model = 'gemini-1.5-flash' ,
872
- contents = 'what is the result of 1000/2?' ,
873
- config = {
874
- 'tools' : [divide_integers ],
875
- 'automatic_function_calling' : {'ignore_call_history' : True }
876
- },
877
- )
910
+ model = 'gemini-1.5-flash' ,
911
+ contents = 'what is the result of 1000/2?' ,
912
+ config = {
913
+ 'tools' : [divide_integers ],
914
+ 'automatic_function_calling' : {'ignore_call_history' : True },
915
+ },
916
+ )
878
917
879
918
assert '500' in response .text
880
919
@@ -889,7 +928,7 @@ def divide_integers(a: int, b: int) -> int:
889
928
contents = 'what is the result of 1000/2?' ,
890
929
config = {
891
930
'tools' : [divide_integers ],
892
- 'automatic_function_calling' : {'ignore_call_history' : True }
931
+ 'automatic_function_calling' : {'ignore_call_history' : True },
893
932
},
894
933
)
895
934
@@ -907,7 +946,12 @@ def mystery_function(a: int, b: int) -> int:
907
946
config = {'tools' : [divide_integers ]},
908
947
)
909
948
assert response .automatic_function_calling_history
910
- assert response .automatic_function_calling_history [- 1 ].parts [0 ].function_response .response ['error' ]
949
+ assert (
950
+ response .automatic_function_calling_history [- 1 ]
951
+ .parts [0 ]
952
+ .function_response .response ['error' ]
953
+ )
954
+
911
955
912
956
@pytest .mark .asyncio
913
957
async def test_automatic_function_calling_async_float_without_decimal (client ):
@@ -971,7 +1015,9 @@ async def get_current_weather_async(city: str) -> str:
971
1015
972
1016
973
1017
@pytest .mark .asyncio
974
- async def test_automatic_function_calling_async_with_async_function_stream (client ):
1018
+ async def test_automatic_function_calling_async_with_async_function_stream (
1019
+ client ,
1020
+ ):
975
1021
async def get_current_weather_async (city : str ) -> str :
976
1022
"""Returns the current weather in the city."""
977
1023
@@ -1182,8 +1228,9 @@ def test_code_execution_tool(client):
1182
1228
1183
1229
assert response .executable_code
1184
1230
assert (
1185
- 'prime' in response .code_execution_result .lower () or
1186
- '5117' in response .code_execution_result )
1231
+ 'prime' in response .code_execution_result .lower ()
1232
+ or '5117' in response .code_execution_result
1233
+ )
1187
1234
1188
1235
1189
1236
def test_afc_logs_to_logger_instance (client , caplog ):
0 commit comments