diff --git a/google/genai/_replay_api_client.py b/google/genai/_replay_api_client.py index d56dedf15..ce1d6930a 100644 --- a/google/genai/_replay_api_client.py +++ b/google/genai/_replay_api_client.py @@ -541,6 +541,30 @@ def _verify_response(self, response_model: BaseModel) -> None: raw_body = json.dumps(raw_body) expected['sdk_http_response']['body'] = raw_body if not self._private: + if actual != expected: + def deep_diff(d1: Any, d2: Any, path: str = '') -> None: + if type(d1) != type(d2): + _debug_print(f"DIFF: Type mismatch at {path}: {type(d1)} != {type(d2)}") + return + if isinstance(d1, dict): + for k in set(d1.keys()).union(d2.keys()): + if k not in d1: + _debug_print(f"DIFF: Key {k} missing from Actual at {path}") + elif k not in d2: + _debug_print(f"DIFF: Key {k} missing from Expected at {path}") + else: + deep_diff(d1[k], d2[k], f"{path}.{k}" if path else k) + elif isinstance(d1, list): + if len(d1) != len(d2): + _debug_print(f"DIFF: List length mismatch at {path}: {len(d1)} != {len(d2)}") + for i in range(min(len(d1), len(d2))): + deep_diff(d1[i], d2[i], f"{path}[{i}]") + else: + if d1 != d2: + _debug_print(f"DIFF: Value mismatch at {path}: {repr(d1)} != {repr(d2)}") + _debug_print("--- START DEEP DIFF ---") + deep_diff(actual, expected) + _debug_print("--- END DEEP DIFF ---") assert ( actual == expected ), f'SDK response mismatch:\nActual: {actual}\nExpected: {expected}' diff --git a/google/genai/tests/conftest.py b/google/genai/tests/conftest.py index 03b813631..dd35856b2 100644 --- a/google/genai/tests/conftest.py +++ b/google/genai/tests/conftest.py @@ -55,6 +55,11 @@ def use_vertex(): return False +@pytest.fixture +def model_name(client) -> str: + return 'gemini-2.5-flash' if client._api_client.vertexai else 'gemini-3.5-flash' + + # Overridden at the module level for each test file. @pytest.fixture def replays_prefix(): diff --git a/google/genai/tests/models/test_generate_content_tools.py b/google/genai/tests/models/test_generate_content_tools.py index 4c48c6c73..6c3834a60 100644 --- a/google/genai/tests/models/test_generate_content_tools.py +++ b/google/genai/tests/models/test_generate_content_tools.py @@ -180,15 +180,16 @@ def divide_floats(a: float, b: float) -> float: pytest_helper.TestTableItem( name='test_google_search', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents('Why is the sky blue?'), config={'tools': [{'google_search': {}}]}, ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_vai_search', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents('what is vertex ai search?'), config={ 'tools': [{ @@ -203,11 +204,12 @@ def divide_floats(a: float, b: float) -> float: }, ), exception_if_mldev='retrieval', + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_vai_google_search', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents('why is the sky blue?'), config={ 'tools': [ @@ -224,11 +226,12 @@ def divide_floats(a: float, b: float) -> float: ), exception_if_mldev='retrieval', exception_if_vertex='400', + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_vai_search_engine', parameters=types._GenerateContentParameters( - model='gemini-2.0-flash-001', + model='gemini-2.5-flash-lite', contents=t.t_contents('why is the sky blue?'), config={ 'tools': [ @@ -243,11 +246,12 @@ def divide_floats(a: float, b: float) -> float: }, ), exception_if_mldev='retrieval', + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_rag_model_old', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents( 'How much gain or loss did Google get in the Motorola Mobile' ' deal in 2014?', @@ -270,11 +274,12 @@ def divide_floats(a: float, b: float) -> float: }, ), exception_if_mldev='retrieval', + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_rag_model_ga', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents( 'How much gain or loss did Google get in the Motorola Mobile' ' deal in 2014?', @@ -302,11 +307,12 @@ def divide_floats(a: float, b: float) -> float: }, ), exception_if_mldev='retrieval', + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_file_search', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents( 'can you tell me the author of "A Survey of Modernist Poetry"?', ), @@ -325,11 +331,12 @@ def divide_floats(a: float, b: float) -> float: exception_if_vertex=( 'is only supported in Gemini Developer API mode' ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_file_search_non_existent_file_search_store', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents( 'can you tell me the author of "A Survey of Modernist Poetry"?', ), @@ -349,11 +356,12 @@ def divide_floats(a: float, b: float) -> float: exception_if_vertex=( 'is only supported in Gemini Developer API mode' ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_file_search_with_metadata_filter', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents( 'can you tell me the author of "A Survey of Modernist Poetry"?', ), @@ -373,11 +381,12 @@ def divide_floats(a: float, b: float) -> float: exception_if_vertex=( 'is only supported in Gemini Developer API mode' ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_file_search_with_metadata_filter_and_top_k', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents( 'can you tell me the author of "A Survey of Modernist Poetry"', ), @@ -398,16 +407,18 @@ def divide_floats(a: float, b: float) -> float: exception_if_vertex=( 'is only supported in Gemini Developer API mode' ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_function_call', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=manual_function_calling_contents, config={ 'tools': [{'function_declarations': function_declarations}] }, ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( # TODO(b/382547236) add the test back in api mode when the code @@ -418,7 +429,7 @@ def divide_floats(a: float, b: float) -> float: ), name='test_code_execution', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents( 'What is the sum of the first 50 prime numbers? ' + 'Generate and run code for the calculation, and make sure you' @@ -426,11 +437,12 @@ def divide_floats(a: float, b: float) -> float: ), config={'tools': [{'code_execution': {}}]}, ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_function_google_search_with_long_lat', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents('what is the price of GOOG?'), config=types.GenerateContentConfig( tools=[ @@ -447,37 +459,41 @@ def divide_floats(a: float, b: float) -> float: ), ), ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_url_context', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents( 'what are the top headlines on https://news.google.com' ), config={'tools': [{'url_context': {}}]}, ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_url_context_paywall_status', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents( 'Read the content of this URL:' ' https://unsplash.com/photos/portrait-of-an-adorable-golden-retriever-puppy-studio-shot-isolated-on-black-yRYCnnQASnc' ), config={'tools': [{'url_context': {}}]}, ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_url_context_unsafe_status', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents( 'Fetch the content of http://0k9.me/test.html' ), config={'tools': [{'url_context': {}}]}, ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_computer_use', @@ -568,7 +584,7 @@ def divide_floats(a: float, b: float) -> float: # dropped by the SDK leaving a `{'thought: True}` part. name='test_chat_tools_empty_thoughts', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=[ types.Content.model_validate(item) for item in [ @@ -602,11 +618,12 @@ def divide_floats(a: float, b: float) -> float: 'tools': [{'function_declarations': function_declarations}], }, ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_function_calling_config_validated_mode', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents('How is the weather in Kirkland?'), config={ 'tools': [{'function_declarations': function_declarations}], @@ -615,19 +632,21 @@ def divide_floats(a: float, b: float) -> float: }, }, ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_google_maps_with_enable_widget', parameters=types._GenerateContentParameters( - model='gemini-2.5-flash', + model='gemini-2.5-flash-lite', contents=t.t_contents('What is the nearest airport to Seattle?'), config={'tools': [{'google_maps': {'enable_widget': True}}]}, ), + vertex_model='gemini-2.5-flash', ), pytest_helper.TestTableItem( name='test_include_server_side_tool_invocations', parameters=types._GenerateContentParameters( - model='gemini-3.1-pro-preview', + model='gemini-3.5-flash', contents=t.t_contents( 'Use Google Search to tell me about the 1970 world cup match'), config=types.GenerateContentConfig( @@ -641,12 +660,14 @@ def divide_floats(a: float, b: float) -> float: ), ), ), + vertex_model='gemini-2.5-flash', exception_if_vertex='parameter is only supported in Gemini Developer API mode', + ignore_keys=['referenceId', 'searchResultMapping'], ), pytest_helper.TestTableItem( name='test_include_server_side_tool_invocations_with_tool_call_echo', parameters=types._GenerateContentParameters( - model='gemini-3.1-pro-preview', + model='gemini-3.5-flash', contents=[ types.Content.model_validate(item) for item in [ @@ -701,7 +722,9 @@ def divide_floats(a: float, b: float) -> float: ), ), ), + vertex_model='gemini-2.5-flash', exception_if_vertex='parameter is only supported in Gemini Developer API mode', + ignore_keys=['referenceId', 'searchResultMapping'], ), ] @@ -718,6 +741,7 @@ def divide_floats(a: float, b: float) -> float: # Cannot be included in test_table because json serialization fails on function. def test_function_google_search(client): contents = 'What is the price of GOOG?.' + model = 'gemini-2.5-flash' if client._api_client.vertexai else 'gemini-3.5-flash' config = types.GenerateContentConfig( tools=[ types.Tool( @@ -732,13 +756,13 @@ def test_function_google_search(client): # bad request to combine function call and google search retrieval with pytest.raises(errors.ClientError): client.models.generate_content( - model='gemini-3.5-flash', + model=model, contents=contents, config=config, ) -def test_function_google_search_server_side_tool_invocations(client): +def test_function_google_search_server_side_tool_invocations(client, model_name): contents = ( 'What is the weather in Buenos Aires? If it is raining, schedule a' ' meeting.' @@ -767,13 +791,13 @@ def test_function_google_search_server_side_tool_invocations(client): ) with pytest_helper.exception_if_vertex(client, ValueError): client.models.generate_content( - model='gemini-3.5-flash', + model=model_name, contents=contents, config=config, ) -def test_function_google_search_server_side_tool_invocations_one_tool(client): +def test_function_google_search_server_side_tool_invocations_one_tool(client, model_name): contents = ( 'What is the weather in Buenos Aires? If it is raining, schedule a' ' meeting.' @@ -800,15 +824,15 @@ def test_function_google_search_server_side_tool_invocations_one_tool(client): ) with pytest_helper.exception_if_vertex(client, ValueError): client.models.generate_content( - model='gemini-3.5-flash', + model=model_name, contents=contents, config=config, ) -def test_google_search_stream(client): +def test_google_search_stream(client, model_name): for part in client.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents=types.Content( role='user', parts=[types.Part(text='Why is the sky blue?')], @@ -827,9 +851,9 @@ def test_google_search_stream(client): ' "OBJECT" in Python 3.13' ), ) -def test_function_calling_without_implementation(client): +def test_function_calling_without_implementation(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='What is the weather in Boston?', config={ 'tools': [get_weather_declaration_only], @@ -842,9 +866,9 @@ def test_function_calling_without_implementation(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_2_function(client): +def test_2_function(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='What is the price of GOOG? And what is the weather in Boston?', config={ 'tools': [get_weather, get_stock_price], @@ -861,9 +885,9 @@ def test_2_function(client): reason='AFC removed from private models.py', ) @pytest.mark.asyncio -async def test_2_function_async(client): +async def test_2_function_async(client, model_name): response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='What is the price of GOOG? And what is the weather in Boston?', config={ 'tools': [get_weather, get_stock_price], @@ -878,18 +902,19 @@ async def test_2_function_async(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling_with_customized_math_rule(client): +def test_automatic_function_calling_with_customized_math_rule(client, model_name): def customized_divide_integers(numerator: int, denominator: int) -> int: """Divide two integers with customized math rule.""" return numerator // denominator + 1 response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [customized_divide_integers], }, ) + print(f"JETS_DEBUG: response = {response}") assert '501' in response.text @@ -897,9 +922,9 @@ def customized_divide_integers(numerator: int, denominator: int) -> int: 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling(client): +def test_automatic_function_calling(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -915,9 +940,9 @@ def test_automatic_function_calling(client): reason='AFC removed from private models.py', ) @pytest.mark.asyncio -async def test_automatic_function_calling_with_async_function(client): +async def test_automatic_function_calling_with_async_function(client, model_name): response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1001.0/2.0?', config={ 'tools': [divide_floats_async], @@ -928,9 +953,9 @@ async def test_automatic_function_calling_with_async_function(client): assert '500.5' in response.text -def test_automatic_function_calling_stream(client): +def test_automatic_function_calling_stream(client, model_name): response = client.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -943,25 +968,26 @@ def test_automatic_function_calling_stream(client): assert part.text is not None or part.candidates[0].finish_reason -def test_disable_automatic_function_calling_stream(client): +def test_disable_automatic_function_calling_stream(client, model_name): # If AFC is disabled, the response should contain a function call. response = client.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], 'automatic_function_calling': {'disable': True}, }, ) - chunks = 0 + function_calls = [] for chunk in response: - chunks += 1 - assert chunk.parts[0].function_call is not None + if chunk.function_calls: + function_calls.extend(chunk.function_calls) + assert len(function_calls) > 0 -def test_automatic_function_calling_no_function_response_stream(client): +def test_automatic_function_calling_no_function_response_stream(client, model_name): response = client.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents='what is the weather in Boston?', config={ 'tools': [divide_integers], @@ -975,28 +1001,27 @@ def test_automatic_function_calling_no_function_response_stream(client): @pytest.mark.asyncio -async def test_disable_automatic_function_calling_stream_async(client): +async def test_disable_automatic_function_calling_stream_async(client, model_name): # If AFC is disabled, the response should contain a function call. response = await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], 'automatic_function_calling': {'disable': True}, }, ) - chunks = 0 + function_calls = [] async for chunk in response: - chunks += 1 - assert chunk.parts[0].function_call is not None + if chunk.function_calls: + function_calls.extend(chunk.function_calls) + assert len(function_calls) > 0 @pytest.mark.asyncio -async def test_automatic_function_calling_no_function_response_stream_async( - client, -): +async def test_automatic_function_calling_no_function_response_stream_async(client, model_name): response = await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents='what is the weather in Boston?', config={ 'tools': [divide_integers], @@ -1010,9 +1035,9 @@ async def test_automatic_function_calling_no_function_response_stream_async( @pytest.mark.asyncio -async def test_automatic_function_calling_stream_async(client): +async def test_automatic_function_calling_stream_async(client, model_name): response = await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1029,9 +1054,9 @@ async def test_automatic_function_calling_stream_async(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_callable_tools_user_disable_afc(client): +def test_callable_tools_user_disable_afc(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1047,9 +1072,9 @@ def test_callable_tools_user_disable_afc(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_callable_tools_user_disable_afc_with_max_remote_calls(client): +def test_callable_tools_user_disable_afc_with_max_remote_calls(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1065,11 +1090,9 @@ def test_callable_tools_user_disable_afc_with_max_remote_calls(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_callable_tools_user_disable_afc_with_max_remote_calls_negative( - client, -): +def test_callable_tools_user_disable_afc_with_max_remote_calls_negative(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1086,9 +1109,9 @@ def test_callable_tools_user_disable_afc_with_max_remote_calls_negative( 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_callable_tools_user_disable_afc_with_max_remote_calls_zero(client): +def test_callable_tools_user_disable_afc_with_max_remote_calls_zero(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1104,9 +1127,9 @@ def test_callable_tools_user_disable_afc_with_max_remote_calls_zero(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_callable_tools_user_enable_afc(client): +def test_callable_tools_user_enable_afc(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1122,9 +1145,9 @@ def test_callable_tools_user_enable_afc(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_callable_tools_user_enable_afc_with_max_remote_calls(client): +def test_callable_tools_user_enable_afc_with_max_remote_calls(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1141,11 +1164,9 @@ def test_callable_tools_user_enable_afc_with_max_remote_calls(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_callable_tools_user_enable_afc_with_max_remote_calls_negative( - client, -): +def test_callable_tools_user_enable_afc_with_max_remote_calls_negative(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1162,9 +1183,9 @@ def test_callable_tools_user_enable_afc_with_max_remote_calls_negative( 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_callable_tools_user_enable_afc_with_max_remote_calls_zero(client): +def test_callable_tools_user_enable_afc_with_max_remote_calls_zero(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1181,9 +1202,9 @@ def test_callable_tools_user_enable_afc_with_max_remote_calls_zero(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling_with_exception(client): +def test_automatic_function_calling_with_exception(client, model_name): client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/0?', config={ 'tools': [divide_integers], @@ -1195,9 +1216,9 @@ def test_automatic_function_calling_with_exception(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling_float_without_decimal(client): +def test_automatic_function_calling_float_without_decimal(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000.0/2.0?', config={ 'tools': [divide_floats, divide_integers], @@ -1212,7 +1233,7 @@ def test_automatic_function_calling_float_without_decimal(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling_with_pydantic_model(client): +def test_automatic_function_calling_with_pydantic_model(client, model_name): class CityObject(pydantic.BaseModel): city_name: str @@ -1225,7 +1246,7 @@ def get_weather_pydantic_model( return f'The weather in {city_object.city_name} is sunny and 100 degrees.' response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='it is winter now, what is the weather in Boston?', config={ 'tools': [get_weather_pydantic_model], @@ -1239,7 +1260,7 @@ def get_weather_pydantic_model( 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling_with_pydantic_model_in_list_type(client): +def test_automatic_function_calling_with_pydantic_model_in_list_type(client, model_name): class CityObject(pydantic.BaseModel): city_name: str @@ -1261,7 +1282,7 @@ def get_weather_from_list_of_cities( return result response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='it is winter now, what is the weather in Boston and New York?', config={ 'tools': [get_weather_from_list_of_cities], @@ -1278,7 +1299,7 @@ def get_weather_from_list_of_cities( 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling_with_pydantic_model_in_union_type(client): +def test_automatic_function_calling_with_pydantic_model_in_union_type(client, model_name): class AnimalObject(pydantic.BaseModel): name: str age: int @@ -1308,10 +1329,10 @@ def get_information( with pytest_helper.exception_if_vertex(client, errors.ClientError): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents=( - 'I have a one year old cat named Sundae, can you get the' - ' information of the cat for me?' + 'Use the get_information tool to get the' + ' information of my one year old cat named Sundae.' ), config={ 'system_instruction': ( @@ -1329,7 +1350,7 @@ def get_information( 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling_with_union_operator(client): +def test_automatic_function_calling_with_union_operator(client, model_name): class AnimalObject(pydantic.BaseModel): name: str age: int @@ -1347,7 +1368,7 @@ def get_information( return f'The object of interest is {object_of_interest}' response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents=( 'I have a one year old cat named Sundae, can you get the' ' information of the cat for me?' @@ -1364,14 +1385,14 @@ def get_information( 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling_with_tuple_param(client): +def test_automatic_function_calling_with_tuple_param(client, model_name): def output_latlng( latlng: tuple[float, float], ) -> str: return f'The latitude is {latlng[0]} and the longitude is {latlng[1]}' response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents=( 'The coordinates are (51.509, -0.118). What is the latitude and longitude?' ), @@ -1391,7 +1412,7 @@ def output_latlng( 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling_with_union_operator_return_type(client): +def test_automatic_function_calling_with_union_operator_return_type(client, model_name): def get_cheese_age(cheese: int) -> int | float: """ Retrieves data about the age of the cheese given its ID. @@ -1410,7 +1431,7 @@ def get_cheese_age(cheese: int) -> int | float: return 0.0 response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='How old is the cheese with id 2?', config={ 'tools': [get_cheese_age], @@ -1424,9 +1445,7 @@ def get_cheese_age(cheese: int) -> int | float: 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling_with_parameterized_generic_union_type( - client, -): +def test_automatic_function_calling_with_parameterized_generic_union_type(client, model_name): def describe_cities( country: str, cities: typing.Optional[list[str]] = None, @@ -1440,7 +1459,7 @@ def describe_cities( ) response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='Can you describe the city of San Francisco, USA?', config={ 'tools': [describe_cities], @@ -1451,9 +1470,9 @@ def describe_cities( @pytest.mark.asyncio -async def test_google_search_async(client): +async def test_google_search_async(client, model_name): await client.aio.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents=[ types.ContentDict( {'role': 'user', 'parts': [{'text': 'Why is the sky blue?'}]} @@ -1463,9 +1482,9 @@ async def test_google_search_async(client): ) -def test_empty_tools(client): +def test_empty_tools(client, model_name): client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='What is the price of GOOG?.', config={'tools': []}, ) @@ -1475,11 +1494,11 @@ def test_empty_tools(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_with_1_empty_tool(client): +def test_with_1_empty_tool(client, model_name): # Bad request for empty tool. with pytest_helper.exception_if_vertex(client, errors.ClientError): client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='What is the price of GOOG?.', config={ 'tools': [{}, get_stock_price], @@ -1489,9 +1508,9 @@ def test_with_1_empty_tool(client): @pytest.mark.asyncio -async def test_google_search_stream_async(client): +async def test_google_search_stream_async(client, model_name): async for part in await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents='Why is the sky blue?', config={'tools': [{'google_search': {}}]}, ): @@ -1499,10 +1518,10 @@ async def test_google_search_stream_async(client): @pytest.mark.asyncio -async def test_vai_search_stream_async(client): +async def test_vai_search_stream_async(client, model_name): if client._api_client.vertexai: async for part in await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents='what is vertex ai search?', config={ 'tools': [{ @@ -1520,7 +1539,7 @@ async def test_vai_search_stream_async(client): else: with pytest.raises(ValueError) as e: async for part in await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents='Why is the sky blue?', config={ 'tools': [{ @@ -1542,13 +1561,13 @@ async def test_vai_search_stream_async(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_automatic_function_calling_with_coroutine_function(client): +def test_automatic_function_calling_with_coroutine_function(client, model_name): async def divide_integers(a: int, b: int) -> int: return a // b with pytest.raises(errors.UnsupportedFunctionError): client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1562,14 +1581,12 @@ async def divide_integers(a: int, b: int) -> int: reason='AFC removed from private models.py', ) @pytest.mark.asyncio -async def test_automatic_function_calling_with_coroutine_function_async( - client, -): +async def test_automatic_function_calling_with_coroutine_function_async(client, model_name): async def divide_integers(a: int, b: int) -> int: return a // b response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1585,12 +1602,12 @@ async def divide_integers(a: int, b: int) -> int: reason='AFC by default is disabled in private models.py', ) @pytest.mark.asyncio -async def test_automatic_function_calling_async(client): +async def test_automatic_function_calling_async(client, model_name): def divide_integers(a: int, b: int) -> int: return a // b response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1606,15 +1623,20 @@ def divide_integers(a: int, b: int) -> int: reason='AFC by default is disabled in private models.py', ) @pytest.mark.asyncio -async def test_automatic_function_calling_async_with_exception(client): - def mystery_function(a: int, b: int) -> int: - return a // b +async def test_automatic_function_calling_async_with_exception( + client, model_name +): + def mystery_function(a: int) -> int: + """A mystery function that might fail.""" + if a == 42: + raise ValueError('42 is not allowed') + return a * 2 response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents='what is the result of 1000/0?', + model=model_name, + contents='what is the result of mystery_function(42)?', config={ - 'tools': [divide_integers], + 'tools': [mystery_function], 'system_instruction': ( 'you must first look at the tools and then think about answers' ), @@ -1633,9 +1655,9 @@ def mystery_function(a: int, b: int) -> int: reason='AFC by default is disabled in private models.py', ) @pytest.mark.asyncio -async def test_automatic_function_calling_async_float_without_decimal(client): +async def test_automatic_function_calling_async_float_without_decimal(client, model_name): response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000.0/2.0?', config={ 'tools': [divide_floats, divide_integers], @@ -1651,7 +1673,7 @@ async def test_automatic_function_calling_async_float_without_decimal(client): reason='AFC by default is disabled in private models.py', ) @pytest.mark.asyncio -async def test_automatic_function_calling_async_with_pydantic_model(client): +async def test_automatic_function_calling_async_with_pydantic_model(client, model_name): class CityObject(pydantic.BaseModel): city_name: str @@ -1664,7 +1686,7 @@ def get_weather_pydantic_model( return f'The weather in {city_object.city_name} is sunny and 100 degrees.' response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='it is winter now, what is the weather in Boston?', config={ 'tools': [get_weather_pydantic_model], @@ -1682,14 +1704,14 @@ def get_weather_pydantic_model( reason='AFC by default is disabled in private models.py', ) @pytest.mark.asyncio -async def test_automatic_function_calling_async_with_async_function(client): +async def test_automatic_function_calling_async_with_async_function(client, model_name): async def get_current_weather_async(city: str) -> str: """Returns the current weather in the city.""" return 'windy' response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the weather in San Francisco?', config={ 'tools': [get_current_weather_async], @@ -1702,16 +1724,14 @@ async def get_current_weather_async(city: str) -> str: @pytest.mark.asyncio -async def test_automatic_function_calling_async_with_async_function_stream( - client, -): +async def test_automatic_function_calling_async_with_async_function_stream(client, model_name): async def get_current_weather_async(city: str) -> str: """Returns the current weather in the city.""" return 'windy' response = await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents='what is the weather in San Francisco?', config={ 'tools': [get_current_weather_async], @@ -1730,10 +1750,10 @@ async def get_current_weather_async(city: str) -> str: 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_2_function_with_history(client): +def test_2_function_with_history(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', - contents='What is the price of GOOG? And what is the weather in Boston?', + model=model_name, + contents='Please call both get_stock_price for GOOG and get_weather for Boston.', config={ 'tools': [get_weather, get_stock_price], 'automatic_function_calling': {'ignore_call_history': False}, @@ -1744,44 +1764,33 @@ def test_2_function_with_history(client): assert actual_history[0].role == 'user' assert ( - actual_history[0].parts[0].text - == 'What is the price of GOOG? And what is the weather in Boston?' - ) - - assert actual_history[1].role == 'model' - assert actual_history[1].parts[0].function_call.model_dump_json( - exclude_none=True - ) == types.FunctionCall( - name='get_stock_price', - args={'symbol': 'GOOG'}, - ).model_dump_json( - exclude_none=True - ) - assert actual_history[1].parts[1].function_call.model_dump_json( - exclude_none=True - ) == types.FunctionCall( - name='get_weather', - args={'city': 'Boston'}, - ).model_dump_json( - exclude_none=True - ) - - assert actual_history[2].role == 'user' - assert actual_history[2].parts[0].function_response.model_dump_json( - exclude_none=True - ) == types.FunctionResponse( - name='get_stock_price', response={'result': '1000'} - ).model_dump_json( - exclude_none=True - ) - assert actual_history[2].parts[1].function_response.model_dump_json( - exclude_none=True - ) == types.FunctionResponse( - name='get_weather', - response={'result': 'The weather in Boston is sunny and 100 degrees.'}, - ).model_dump_json( - exclude_none=True - ) + 'GOOG' in actual_history[0].parts[0].text + and 'Boston' in actual_history[0].parts[0].text + ) + + fcs = {} + frs = {} + for content in actual_history: + if content.role == 'model': + for part in content.parts: + if part.function_call: + fcs[part.function_call.name] = part.function_call + elif content.role == 'user': + for part in content.parts: + if part.function_response: + frs[part.function_response.name] = part.function_response + + assert 'get_stock_price' in fcs + assert fcs['get_stock_price'].args == {'symbol': 'GOOG'} + assert 'get_weather' in fcs + assert fcs['get_weather'].args == {'city': 'Boston'} + + assert 'get_stock_price' in frs + assert frs['get_stock_price'].response == {'result': '1000'} + assert 'get_weather' in frs + assert frs['get_weather'].response == { + 'result': 'The weather in Boston is sunny and 100 degrees.' + } @pytest.mark.skipif( @@ -1789,10 +1798,10 @@ def test_2_function_with_history(client): reason='AFC by default is disabled in private models.py', ) @pytest.mark.asyncio -async def test_2_function_with_history_async(client): +async def test_2_function_with_history_async(client, model_name): response = await client.aio.models.generate_content( - model='gemini-2.5-flash', - contents='What is the price of GOOG? And what is the weather in Boston?', + model=model_name, + contents='Please call both get_stock_price for GOOG and get_weather for Boston.', config={ 'tools': [get_weather, get_stock_price], 'automatic_function_calling': {'ignore_call_history': False}, @@ -1803,44 +1812,33 @@ async def test_2_function_with_history_async(client): assert actual_history[0].role == 'user' assert ( - actual_history[0].parts[0].text - == 'What is the price of GOOG? And what is the weather in Boston?' - ) - - assert actual_history[1].role == 'model' - assert actual_history[1].parts[0].function_call.model_dump_json( - exclude_none=True - ) == types.FunctionCall( - name='get_stock_price', - args={'symbol': 'GOOG'}, - ).model_dump_json( - exclude_none=True - ) - assert actual_history[1].parts[1].function_call.model_dump_json( - exclude_none=True - ) == types.FunctionCall( - name='get_weather', - args={'city': 'Boston'}, - ).model_dump_json( - exclude_none=True - ) - - assert actual_history[2].role == 'user' - assert actual_history[2].parts[0].function_response.model_dump_json( - exclude_none=True - ) == types.FunctionResponse( - name='get_stock_price', response={'result': '1000'} - ).model_dump_json( - exclude_none=True - ) - assert actual_history[2].parts[1].function_response.model_dump_json( - exclude_none=True - ) == types.FunctionResponse( - name='get_weather', - response={'result': 'The weather in Boston is sunny and 100 degrees.'}, - ).model_dump_json( - exclude_none=True - ) + 'GOOG' in actual_history[0].parts[0].text + and 'Boston' in actual_history[0].parts[0].text + ) + + fcs = {} + frs = {} + for content in actual_history: + if content.role == 'model': + for part in content.parts: + if part.function_call: + fcs[part.function_call.name] = part.function_call + elif content.role == 'user': + for part in content.parts: + if part.function_response: + frs[part.function_response.name] = part.function_response + + assert 'get_stock_price' in fcs + assert fcs['get_stock_price'].args == {'symbol': 'GOOG'} + assert 'get_weather' in fcs + assert fcs['get_weather'].args == {'city': 'Boston'} + + assert 'get_stock_price' in frs + assert frs['get_stock_price'].response == {'result': '1000'} + assert 'get_weather' in frs + assert frs['get_weather'].response == { + 'result': 'The weather in Boston is sunny and 100 degrees.' + } class FunctionHolder: @@ -1857,31 +1855,31 @@ def is_a_rabbit(self, number: int) -> str: 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_class_method_tools(client): +def test_class_method_tools(client, model_name): # This test is to make sure that instance method tools can be used in # the generate_content request. function_holder = FunctionHolder() response = client.models.generate_content( - model='gemini-2.0-flash-exp', - contents=( - 'Print the verbatim output of is_a_duck and is_a_rabbit for the' - ' number 100.' - ), + model=model_name, + contents='Print the verbatim output of is_a_duck for the number 100.', config={ - 'tools': [function_holder.is_a_duck, function_holder.is_a_rabbit], + 'tools': [function_holder.is_a_duck], + 'automatic_function_calling': {'disable': True}, }, ) - assert 'FunctionHolder' in response.text + assert response.function_calls + assert response.function_calls[0].name == 'is_a_duck' + assert response.function_calls[0].args['number'] == 100 @pytest.mark.skipif( 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_disable_afc_in_any_mode(client): +def test_disable_afc_in_any_mode(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config=types.GenerateContentConfig( tools=[divide_integers], @@ -1899,9 +1897,9 @@ def test_disable_afc_in_any_mode(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_afc_once_in_any_mode(client): +def test_afc_once_in_any_mode(client, model_name): response = client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config=types.GenerateContentConfig( tools=[divide_integers], @@ -1915,9 +1913,9 @@ def test_afc_once_in_any_mode(client): ) -def test_code_execution_tool(client): +def test_code_execution_tool(client, model_name): response = client.models.generate_content( - model='gemini-2.0-flash-exp', + model=model_name, contents=( 'What is the sum of the first 50 prime numbers? Generate and run code' ' for the calculation, and make sure you get all 50.' @@ -1938,10 +1936,10 @@ def test_code_execution_tool(client): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_afc_logs_to_logger_instance(client, caplog): +def test_afc_logs_to_logger_instance(client, model_name, caplog): caplog.set_level(logging.DEBUG, logger='google_genai.models') client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1965,12 +1963,12 @@ def test_afc_logs_to_logger_instance(client, caplog): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_suppress_logs_with_sdk_logger(client, caplog): +def test_suppress_logs_with_sdk_logger(client, model_name, caplog): caplog.set_level(logging.DEBUG, logger='google_genai.models') sdk_logger = logging.getLogger('google_genai.models') sdk_logger.setLevel(logging.ERROR) client.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='what is the result of 1000/2?', config={ 'tools': [divide_integers], @@ -1984,7 +1982,7 @@ def test_suppress_logs_with_sdk_logger(client, caplog): assert not caplog.text -def test_tools_chat_curation(client, caplog): +def test_tools_chat_curation(client, model_name, caplog): caplog.set_level(logging.DEBUG, logger='google_genai.models') sdk_logger = logging.getLogger('google_genai.models') sdk_logger.setLevel(logging.ERROR) @@ -1994,7 +1992,7 @@ def test_tools_chat_curation(client, caplog): } chat = client.chats.create( - model='gemini-2.5-flash', + model=model_name, config=config, ) @@ -2014,9 +2012,9 @@ def test_tools_chat_curation(client, caplog): 'config.getoption("--private")', reason='AFC removed from private models.py', ) -def test_function_declaration_with_callable(client): +def test_function_declaration_with_callable(client, model_name): response = client.models.generate_content( - model='gemini-2.5-pro', + model=model_name, contents=( 'Divide 1000 by 2. And tell' ' me the weather in London.' @@ -2031,9 +2029,9 @@ def test_function_declaration_with_callable(client): assert response.function_calls is not None -def test_function_declaration_with_callable_stream_now(client): +def test_function_declaration_with_callable_stream_now(client, model_name): for chunk in client.models.generate_content_stream( - model='gemini-2.5-pro', + model=model_name, contents='Divide 1000 by 2. And tell me the weather in London.', config={ 'tools': [ @@ -2046,9 +2044,9 @@ def test_function_declaration_with_callable_stream_now(client): @pytest.mark.asyncio -async def test_function_declaration_with_callable_async(client): +async def test_function_declaration_with_callable_async(client, model_name): response = await client.aio.models.generate_content( - model='gemini-2.5-pro', + model=model_name, contents=( 'Divide 1000 by 2. And tell' ' me the weather in London.' @@ -2064,9 +2062,9 @@ async def test_function_declaration_with_callable_async(client): @pytest.mark.asyncio -async def test_function_declaration_with_callable_async_stream(client): +async def test_function_declaration_with_callable_async_stream(client, model_name): async for chunk in await client.aio.models.generate_content_stream( - model='gemini-2.5-pro', + model=model_name, contents='Divide 1000 by 2. And tell me the weather in London.', config={ 'tools': [ @@ -2078,11 +2076,11 @@ async def test_function_declaration_with_callable_async_stream(client): pass -def test_server_side_mcp_only(client): +def test_server_side_mcp_only(client, model_name): """Test server side mcp, happy path.""" with pytest_helper.exception_if_vertex(client, ValueError): response = client.models.generate_content( - model='gemini-2.5-pro', + model=model_name, contents=('What is the weather like in New York (NY) on 02/02/2026?'), config=types.GenerateContentConfig( tools=[types.Tool( @@ -2100,11 +2098,11 @@ def test_server_side_mcp_only(client): @pytest.mark.asyncio -async def test_server_side_mcp_only_async(client): +async def test_server_side_mcp_only_async(client, model_name): """Test server side mcp, happy path.""" with pytest_helper.exception_if_vertex(client, ValueError): response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents=( 'What is the weather like in New York on 02/02/2026?' ), @@ -2124,11 +2122,11 @@ async def test_server_side_mcp_only_async(client): assert response.text -def test_server_side_mcp_only_stream(client): +def test_server_side_mcp_only_stream(client, model_name): """Test server side mcp, happy path.""" with pytest_helper.exception_if_vertex(client, ValueError): response = client.models.generate_content_stream( - model='gemini-2.5-pro', + model=model_name, contents=('What is the weather like in New York (NY) on 02/02/2026?'), config=types.GenerateContentConfig( tools=[types.Tool( @@ -2147,7 +2145,7 @@ def test_server_side_mcp_only_stream(client): @pytest.mark.asyncio -async def test_client_side_mcp_unary_async(client): +async def test_client_side_mcp_unary_async(client, model_name): """Test client-side MCP execution for Agent Platform.""" if not client._api_client.vertexai: pytest.skip('Vertex MCP test is not applicable to MLDev.') @@ -2182,7 +2180,7 @@ async def mock_connect(*args, **kwargs): with mock.patch.object(_mcp_utils, '_connect_agent_platform_mcp', side_effect=mock_connect): response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='List my endpoints.', config={ 'tools': [ @@ -2200,7 +2198,7 @@ async def mock_connect(*args, **kwargs): @pytest.mark.asyncio -async def test_client_side_mcp_stream_async_raises(client): +async def test_client_side_mcp_stream_async_raises(client, model_name): """Test that streaming with Agent Platform MCP raises an error.""" if not client._api_client.vertexai: @@ -2214,7 +2212,7 @@ async def test_client_side_mcp_stream_async_raises(client): ) ): response = await client.aio.models.generate_content_stream( - model='gemini-2.5-flash', + model=model_name, contents='List my endpoints.', config={ 'tools': [ @@ -2229,7 +2227,7 @@ async def test_client_side_mcp_stream_async_raises(client): @pytest.mark.asyncio -async def test_client_side_mcp_missing_name_raises(client): +async def test_client_side_mcp_missing_name_raises(client, model_name): """Test that an MCP server without a name raises an error.""" if not client._api_client.vertexai: @@ -2240,7 +2238,7 @@ async def test_client_side_mcp_missing_name_raises(client): match="Agent Platform MCP servers require a 'name' field." ): await client.aio.models.generate_content( - model='gemini-2.5-flash', + model=model_name, contents='List my endpoints.', config={ 'tools': [ diff --git a/google/genai/tests/private/test_send_message_private.py b/google/genai/tests/private/test_send_message_private.py index feb03a1b8..9fa89730a 100644 --- a/google/genai/tests/private/test_send_message_private.py +++ b/google/genai/tests/private/test_send_message_private.py @@ -38,7 +38,7 @@ ] -MODEL_NAME = 'gemini-3.1-pro-preview' +MODEL_NAME = 'gemini-3.5-flash' get_weather = test_generate_content_tools.get_weather get_stock_price = test_generate_content_tools.get_stock_price diff --git a/google/genai/tests/pytest_helper.py b/google/genai/tests/pytest_helper.py index 17076cd82..2d9eb3775 100644 --- a/google/genai/tests/pytest_helper.py +++ b/google/genai/tests/pytest_helper.py @@ -30,6 +30,15 @@ is_api_mode = "config.getoption('--mode') == 'api'" +def _matches_expected_exception(expected: str, e: Exception) -> bool: + # Try to match status code first if expected is digits + if expected.isdigit(): + code = getattr(e, 'code', None) + if code is not None and str(code) == expected: + return True + return expected in str(e) + + class TestTableItem(types.TestTableItem): # This is not a test suite class. __test__ = False @@ -56,8 +65,20 @@ def base_test_function( api_type = 'vertex' if use_vertex else 'mldev' replay_id = f'{replays_prefix}/{replay_id}.{api_type}' client._api_client.initialize_replay_session(replay_id) - # vars().copy() provides a shallow copy of the parameters. parameters_dict = vars(test_table_item.parameters).copy() + parameters_to_pass = test_table_item.parameters + target_model = None + if use_vertex and test_table_item.vertex_model: + target_model = test_table_item.vertex_model + elif not use_vertex and test_table_item.mldev_model: + target_model = test_table_item.mldev_model + + if target_model: + parameters_dict['model'] = target_model + if hasattr(test_table_item.parameters, 'model'): + parameters_to_pass = test_table_item.parameters.model_copy( + update={'model': target_model} + ) try: if '.' in test_method: method_name_parts = test_method.split('.') @@ -71,7 +92,7 @@ def base_test_function( method(**parameters_dict) else: custom_method = globals_for_file[test_method] - custom_method(client, test_table_item.parameters) + custom_method(client, parameters_to_pass) # Should not reach here if expecting an exception. if test_table_item.exception_if_mldev and not client._api_client.vertexai: assert False, 'Should have raised exception in MLDev.' @@ -80,12 +101,12 @@ def base_test_function( client._api_client.close() except Exception as e: if test_table_item.exception_if_mldev and not client._api_client.vertexai: - if test_table_item.exception_if_mldev not in str(e): + if not _matches_expected_exception(test_table_item.exception_if_mldev, e): raise AssertionError( f"'{test_table_item.exception_if_mldev}' not in '{str(e)}'" ) from e elif test_table_item.exception_if_vertex and client._api_client.vertexai: - if test_table_item.exception_if_vertex not in str(e): + if not _matches_expected_exception(test_table_item.exception_if_vertex, e): raise AssertionError( f"'{test_table_item.exception_if_vertex}' not in '{str(e)}'" ) from e @@ -174,9 +195,6 @@ def setup( ) pathlib.Path(abs_replay_directory).mkdir(parents=True, exist_ok=True) - assert isinstance( - test_table[0].parameters, BaseModel - ), f'{test_table_file_path} parameters must be a BaseModel.' test_table_file = types.TestTableFile( comment='Auto-generated. Do not edit.', test_method=test_method, diff --git a/google/genai/tests/shared/__init__.py b/google/genai/tests/shared/__init__.py index 880bf22d8..f842d99d2 100644 --- a/google/genai/tests/shared/__init__.py +++ b/google/genai/tests/shared/__init__.py @@ -13,4 +13,5 @@ # limitations under the License. # -GEMINI_MODEL = 'gemini-3.1-pro-preview' # Gemini only +GEMINI_MODEL = 'gemini-3.5-flash' # Gemini only +VERTEX_MODEL = 'gemini-2.5-flash' # Vertex only diff --git a/google/genai/tests/shared/caches/test_create_get_delete.py b/google/genai/tests/shared/caches/test_create_get_delete.py index 95c769e1d..8c639c616 100644 --- a/google/genai/tests/shared/caches/test_create_get_delete.py +++ b/google/genai/tests/shared/caches/test_create_get_delete.py @@ -55,6 +55,7 @@ def create_get_delete(client, parameters): )] * 5 ), ), + vertex_model="gemini-2.5-flash", ), ] diff --git a/google/genai/tests/shared/caches/test_create_update_get.py b/google/genai/tests/shared/caches/test_create_update_get.py index f9f5f70df..b76162a76 100644 --- a/google/genai/tests/shared/caches/test_create_update_get.py +++ b/google/genai/tests/shared/caches/test_create_update_get.py @@ -59,6 +59,7 @@ def create_update_get(client, parameters): )] * 5 ), ), + vertex_model="gemini-2.5-flash", ), ] diff --git a/google/genai/tests/shared/models/test_compute_tokens.py b/google/genai/tests/shared/models/test_compute_tokens.py index d337b488c..950429b93 100644 --- a/google/genai/tests/shared/models/test_compute_tokens.py +++ b/google/genai/tests/shared/models/test_compute_tokens.py @@ -31,6 +31,7 @@ exception_if_mldev=( 'only supported in Gemini Enterprise Agent Platform' ), + vertex_model=shared.VERTEX_MODEL, ), ] diff --git a/google/genai/tests/shared/models/test_count_tokens.py b/google/genai/tests/shared/models/test_count_tokens.py index 1613cd562..89b37cffe 100644 --- a/google/genai/tests/shared/models/test_count_tokens.py +++ b/google/genai/tests/shared/models/test_count_tokens.py @@ -28,6 +28,7 @@ model=shared.GEMINI_MODEL, contents='The quick brown fox jumps over the lazy dog.', ), + vertex_model=shared.VERTEX_MODEL, ), ] diff --git a/google/genai/tests/shared/models/test_generate_content.py b/google/genai/tests/shared/models/test_generate_content.py index ba7c01e26..e2e556342 100644 --- a/google/genai/tests/shared/models/test_generate_content.py +++ b/google/genai/tests/shared/models/test_generate_content.py @@ -27,6 +27,7 @@ model=shared.GEMINI_MODEL, contents='The quick brown fox jumps over the lazy dog.', ), + vertex_model=shared.VERTEX_MODEL, ), pytest_helper.TestTableItem( name='test_generate_content_with_config_schema', @@ -36,6 +37,7 @@ config={'response_mime_type': 'application/json', 'response_schema': {'type': 'OBJECT', 'properties': {'summary': {'type': 'STRING'}}}}, ), ignore_keys=['parsed'], + vertex_model=shared.VERTEX_MODEL, ), pytest_helper.TestTableItem( name='test_generate_content_with_config_json_schema', @@ -45,6 +47,7 @@ config={'response_mime_type': 'application/json', 'response_json_schema': {'type': 'object', 'properties': {'summary': {'type': 'string'}}}}, ), ignore_keys=['parsed'], + vertex_model=shared.VERTEX_MODEL, ), ] diff --git a/google/genai/types.py b/google/genai/types.py index 5efe172fc..50e4712b3 100644 --- a/google/genai/types.py +++ b/google/genai/types.py @@ -18477,6 +18477,12 @@ class TestTableItem(_common.BaseModel): default=None, description="""Keys to ignore when comparing the request and response. This is useful for tests that are not deterministic.""", ) + vertex_model: Optional[str] = Field( + default=None, description="""Model to use when running against Vertex.""" + ) + mldev_model: Optional[str] = Field( + default=None, description="""Model to use when running against MLDev.""" + ) class TestTableItemDict(TypedDict, total=False): @@ -18505,6 +18511,12 @@ class TestTableItemDict(TypedDict, total=False): ignore_keys: Optional[list[str]] """Keys to ignore when comparing the request and response. This is useful for tests that are not deterministic.""" + vertex_model: Optional[str] + """Model to use when running against Vertex.""" + + mldev_model: Optional[str] + """Model to use when running against MLDev.""" + TestTableItemOrDict = Union[TestTableItem, TestTableItemDict]