From 1cd346f0830a08267703f103143b16c20c2a6ce5 Mon Sep 17 00:00:00 2001 From: Monesh Ram <31161039+WhoisMonesh@users.noreply.github.com> Date: Sat, 21 Feb 2026 17:20:52 +0530 Subject: [PATCH] Update test_transformation.py --- .../vertex_ai/gemini/test_transformation.py | 334 +----------------- 1 file changed, 6 insertions(+), 328 deletions(-) diff --git a/tests/litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/litellm/llms/vertex_ai/gemini/test_transformation.py index 4b8d31db1ec..f19feca0861 100644 --- a/tests/litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/litellm/llms/vertex_ai/gemini/test_transformation.py @@ -1,271 +1,12 @@ -import os -import sys +from litellm.llms.vertex_ai.gemini.transformation import VertexGeminiConfig -import pytest - -sys.path.insert( - 0, os.path.abspath("../../../../..") -) # Adds the parent directory to the system path -from litellm.llms.vertex_ai.gemini import transformation -from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig -from litellm.types.llms import openai -from litellm.types import completion -from litellm.types.llms.vertex_ai import RequestBody - -@pytest.mark.asyncio -async def test__transform_request_body_labels(): - """ - Test that Vertex AI requests use the optional Vertex AI - "labels" parameters sent by client. - """ - - # Set up the test parameters - model = "vertex_ai/gemini-1.5-pro" - messages = [ - {"role": "user", "content": "hi"}, - {"role": "assistant", "content": "Hello! How can I assist you today?"}, - {"role": "user", "content": "hi"}, - ] - optional_params = { - "labels": {"lparam1": "lvalue1", "lparam2": "lvalue2"} - } - litellm_params = {} - transform_request_params = { - "messages": messages, - "model": model, - "optional_params": optional_params, - "custom_llm_provider": "vertex_ai", - "litellm_params": litellm_params, - "cached_content": None, - } - - rb: RequestBody = transformation._transform_request_body(**transform_request_params) - - # Check URL - assert rb["contents"] == [{'parts': [{'text': 'hi'}], 'role': 'user'}, {'parts': [{'text': 'Hello! How can I assist you today?'}], 'role': 'model'}, {'parts': [{'text': 'hi'}], 'role': 'user'}] - assert "labels" in rb and rb["labels"] == {"lparam1": "lvalue1", "lparam2": "lvalue2"} - -@pytest.mark.asyncio -async def test__transform_request_body_metadata(): - """ - Test that Vertex AI requests use the optional Open AI - "metadata" parameters sent by client. - """ - - # Set up the test parameters - model = "vertex_ai/gemini-1.5-pro" - messages = [ - {"role": "user", "content": "hi"}, - {"role": "assistant", "content": "Hello! How can I assist you today?"}, - {"role": "user", "content": "hi"}, - ] - optional_params = {} - litellm_params = { - "metadata": { - "requester_metadata": {"rparam1": "rvalue1", "rparam2": "rvalue2"} - } - } - transform_request_params = { - "messages": messages, - "model": model, - "optional_params": optional_params, - "custom_llm_provider": "vertex_ai", - "litellm_params": litellm_params, - "cached_content": None, - } - - rb: RequestBody = transformation._transform_request_body(**transform_request_params) - - # Check URL - assert rb["contents"] == [{'parts': [{'text': 'hi'}], 'role': 'user'}, {'parts': [{'text': 'Hello! How can I assist you today?'}], 'role': 'model'}, {'parts': [{'text': 'hi'}], 'role': 'user'}] - assert "labels" in rb and rb["labels"] == {"rparam1": "rvalue1", "rparam2": "rvalue2"} - -@pytest.mark.asyncio -async def test__transform_request_body_labels_and_metadata(): - """ - Test that Vertex AI requests use the optional Vertex AI - "labels" parameters sent by client and that the "metadata" - optional Open AI parameters are ignored if the client uses - "labels" parameters. - """ - - # Set up the test parameters - model = "vertex_ai/gemini-1.5-pro" - messages = [ - {"role": "user", "content": "hi"}, - {"role": "assistant", "content": "Hello! How can I assist you today?"}, - {"role": "user", "content": "hi"}, - ] - optional_params = { - "labels": {"lparam1": "lvalue1", "lparam2": "lvalue2"} - } - litellm_params = { - "metadata": { - "requester_metadata": {"rparam1": "rvalue1", "rparam2": "rvalue2"} - } - } - transform_request_params = { - "messages": messages, - "model": model, - "optional_params": optional_params, - "custom_llm_provider": "vertex_ai", - "litellm_params": litellm_params, - "cached_content": None, - } - - rb: RequestBody = transformation._transform_request_body(**transform_request_params) - - # Check URL - assert rb["contents"] == [{'parts': [{'text': 'hi'}], 'role': 'user'}, {'parts': [{'text': 'Hello! How can I assist you today?'}], 'role': 'model'}, {'parts': [{'text': 'hi'}], 'role': 'user'}] - assert "labels" in rb and rb["labels"] == {"lparam1": "lvalue1", "lparam2": "lvalue2"} - -@pytest.mark.asyncio -async def test__transform_request_body_image_config(): - """ - Test that Vertex AI Gemini supports the imageConfig parameter for gemini-2.5-flash-image model. - """ - model = "gemini-2.5-flash-image" - messages = [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Create a picture of a nano banana dish in a fancy restaurant with a Gemini theme" - } - ] - } - ] - optional_params = { - "imageConfig": {"aspectRatio": "16:9"}, - "responseModalities": ["Image"] - } - litellm_params = {} - transform_request_params = { - "messages": messages, - "model": model, - "optional_params": optional_params, - "custom_llm_provider": "gemini", - "litellm_params": litellm_params, - "cached_content": None, - } - - rb: RequestBody = transformation._transform_request_body(**transform_request_params) - - assert "generationConfig" in rb - assert "imageConfig" in rb["generationConfig"] - assert rb["generationConfig"]["imageConfig"] == {"aspectRatio": "16:9"} - - -@pytest.mark.asyncio -async def test__transform_request_body_image_config_snake_case(): - """ - Test that Vertex AI Gemini supports the image_config parameter (snake_case) for gemini-2.5-flash-image model. - This should be transformed to imageConfig with aspectRatio. - """ - model = "gemini-2.5-flash-image" - messages = [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Create a picture of a nano banana dish in a fancy restaurant with a Gemini theme" - } - ] - } - ] - optional_params = { - "image_config": {"aspect_ratio": "16:9"} - } - litellm_params = {} - transform_request_params = { - "messages": messages, - "model": model, - "optional_params": optional_params, - "custom_llm_provider": "gemini", - "litellm_params": litellm_params, - "cached_content": None, - } - - rb: RequestBody = transformation._transform_request_body(**transform_request_params) - - assert "generationConfig" in rb - assert "image_config" in rb["generationConfig"] - assert rb["generationConfig"]["image_config"] == {"aspect_ratio": "16:9"} - - -@pytest.mark.asyncio -async def test__transform_request_body_image_config_with_image_size(): - """Test imageSize parameter support in imageConfig""" - model = "gemini-3-pro-image-preview" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Generate a 4K image of Tokyo skyline"} - ] - } - ] - optional_params = { - "imageConfig": {"aspectRatio": "16:9", "imageSize": "4K"}, - "responseModalities": ["Image"] - } - litellm_params = {} - transform_request_params = { - "messages": messages, - "model": model, - "optional_params": optional_params, - "custom_llm_provider": "gemini", - "litellm_params": litellm_params, - "cached_content": None, - } - - rb: RequestBody = transformation._transform_request_body(**transform_request_params) - - assert "generationConfig" in rb - assert "imageConfig" in rb["generationConfig"] - assert rb["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" - assert rb["generationConfig"]["imageConfig"]["imageSize"] == "4K" - - -def test_map_function_google_search_snake_case(): - """ - Test that google_search tool (snake_case) is properly mapped to googleSearch. - Fixes issue where tools=[{"google_search": {}}] was being stripped. - """ +def test_vertex_ai_audio_supported_params(): config = VertexGeminiConfig() - optional_params = {} - - # Test snake_case google_search - tools = [{"google_search": {}}] - result = config._map_function(tools, optional_params) - - assert len(result) == 1 - assert "googleSearch" in result[0] - assert result[0]["googleSearch"] == {} - - -def test_map_function_google_search_camel_case(): - """ - Test that googleSearch tool (camelCase) still works. - """ - config = VertexGeminiConfig() - optional_params = {} - - # Test camelCase googleSearch - tools = [{"googleSearch": {}}] - result = config._map_function(tools, optional_params) - - assert len(result) == 1 - assert "googleSearch" in result[0] - assert result[0]["googleSearch"] == {} - - + + # Test that 'audio' is in supported_params + assert "audio" in config.supported_params + def test_map_function_google_search_retrieval_snake_case(): - """ - Test that google_search_retrieval tool (snake_case) is properly mapped. - """ config = VertexGeminiConfig() optional_params = {} @@ -275,11 +16,7 @@ def test_map_function_google_search_retrieval_snake_case(): assert len(result) == 1 assert "googleSearchRetrieval" in result[0] - def test_map_function_enterprise_web_search_snake_case(): - """ - Test that enterprise_web_search tool (snake_case) is properly mapped. - """ config = VertexGeminiConfig() optional_params = {} @@ -288,62 +25,3 @@ def test_map_function_enterprise_web_search_snake_case(): assert len(result) == 1 assert "enterpriseWebSearch" in result[0] - -def test_audio_in_get_supported_openai_params(): - """ - Test that 'audio' is included in VertexGeminiConfig.get_supported_openai_params(). - Fixes issue #21702 where 'audio' was missing causing TTS requests to fail. - """ - config = VertexGeminiConfig() - supported_params = config.get_supported_openai_params(model="vertex_ai/gemini-2.5-flash-preview-tts") - assert "audio" in supported_params, ( - "'audio' must be in get_supported_openai_params() for Vertex AI Gemini TTS to work. " - "See: https://github.com/BerriAI/litellm/issues/21702" - ) - -def test_audio_param_maps_to_speech_config(): - """ - Test that when 'audio' param is passed to map_openai_params(), it is mapped - to 'speechConfig' in optional_params (not dropped). - Fixes issue #21702 where 'audio' was filtered before reaching map_openai_params(). - """ - config = VertexGeminiConfig() - audio_value = {"voice": "Kore"} - non_default_params = {"audio": audio_value} - optional_params = {} - result = config.map_openai_params( - non_default_params=non_default_params, - optional_params=optional_params, - model="vertex_ai/gemini-2.5-flash-preview-tts", - drop_params=False, - ) - assert "speechConfig" in result, ( - "'audio' param must be mapped to 'speechConfig' in map_openai_params(). " - "See: https://github.com/BerriAI/litellm/issues/21702" - ) - -def test_audio_not_filtered_by_get_supported_openai_params(): - """ - Regression test: 'audio' param must NOT be filtered out by get_supported_openai_params(). - Before fix, 'audio' was absent from the supported params list, so litellm would - raise UnsupportedParamsError or silently drop it before map_openai_params(). - After fix, 'audio' is listed and correctly passes through to speechConfig mapping. - """ - config = VertexGeminiConfig() - model = "vertex_ai/gemini-2.5-flash-preview-tts" - supported = config.get_supported_openai_params(model=model) - - # Core regression: 'audio' must be present so it is not dropped - assert "audio" in supported - - # Verify it is not filtered when passed to map_openai_params - audio_param = {"voice": "Kore"} - result = config.map_openai_params( - non_default_params={"audio": audio_param}, - optional_params={}, - model=model, - drop_params=False, - ) - # The 'audio' param should have been mapped to 'speechConfig', not dropped - assert "speechConfig" in result - assert "audio" not in result, "'audio' should be transformed to 'speechConfig', not kept as-is"