mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Update test_transformation.py
This commit is contained in:
Monesh Ram
GitHub
parent
de0124a7e0
commit
1cd346f083
1 changed files with 6 additions and 328 deletions
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue