Merge branch 'main' into fix/gemini-imagen-model-name-validation

This commit is contained in:
Ifta Khairul Alam Adil 2025-08-27 21:52:47 +02:00
commit 75aeca14b2
14 changed files with 1076 additions and 980 deletions

View file

@ -570,6 +570,7 @@ router_settings:
| LITELLM_LICENSE | License key for LiteLLM usage
| LITELLM_LOCAL_MODEL_COST_MAP | Local configuration for model cost mapping in LiteLLM
| LITELLM_LOG | Enable detailed logging for LiteLLM
| LITELLM_LOG_FILE | File path to write LiteLLM logs to. When set, logs will be written to both console and the specified file
| LITELLM_MASTER_KEY | Master key for proxy authentication
| LITELLM_MODE | Operating mode for LiteLLM (e.g., production, development)
| LITELLM_RATE_LIMIT_WINDOW_SIZE | Rate limit window size for LiteLLM. Default is 60

View file

@ -1,7 +1,7 @@
import asyncio
import contextvars
from functools import partial
from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast, overload, List
from typing import Any, Coroutine, Dict, List, Literal, Optional, Union, cast, overload
import httpx
@ -347,6 +347,7 @@ def image_generation( # noqa: PLR0915
raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
return llm_http_handler.image_generation_handler(
api_key=api_key,
model=model,
prompt=prompt,
image_generation_provider_config=image_generation_config,

View file

@ -2681,6 +2681,7 @@ class BaseLLMHTTPHandler:
_is_async: bool = False,
fake_stream: bool = False,
litellm_metadata: Optional[Dict[str, Any]] = None,
api_key: Optional[str] = None,
) -> Union[
ImageResponse,
Coroutine[Any, Any, ImageResponse],
@ -2705,6 +2706,7 @@ class BaseLLMHTTPHandler:
client=client if isinstance(client, AsyncHTTPHandler) else None,
fake_stream=fake_stream,
litellm_metadata=litellm_metadata,
api_key=api_key,
)
if client is None or not isinstance(client, HTTPHandler):
@ -2715,7 +2717,7 @@ class BaseLLMHTTPHandler:
sync_httpx_client = client
headers = image_generation_provider_config.validate_environment(
api_key=litellm_params.get("api_key", None),
api_key=api_key,
headers=image_generation_optional_request_params.get("extra_headers", {})
or {},
model=model,
@ -2798,6 +2800,7 @@ class BaseLLMHTTPHandler:
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
fake_stream: bool = False,
litellm_metadata: Optional[Dict[str, Any]] = None,
api_key: Optional[str] = None,
) -> ImageResponse:
"""
Async version of the image generation handler.
@ -2812,7 +2815,7 @@ class BaseLLMHTTPHandler:
async_httpx_client = client
headers = image_generation_provider_config.validate_environment(
api_key=litellm_params.get("api_key", None),
api_key=api_key,
headers=image_generation_optional_request_params.get("extra_headers", {})
or {},
model=model,

View file

@ -35,6 +35,7 @@ from litellm.types.llms.openai import (
ChatCompletionFileObject,
ChatCompletionImageObject,
ChatCompletionTextObject,
ChatCompletionUserMessage,
)
from litellm.types.llms.vertex_ai import *
from litellm.types.llms.vertex_ai import (
@ -475,6 +476,13 @@ async def async_transform_request_body(
optional_params=optional_params,
)
def _default_user_message_when_system_message_passed() -> ChatCompletionUserMessage:
"""
Returns a default user message when a "system" message is passed in gemini fails.
This adds a blank user message to the messages list, to ensure that gemini doesn't fail the request.
"""
return ChatCompletionUserMessage(content=".", role="user")
def _transform_system_message(
supports_system_message: bool, messages: List[AllMessageValues]
@ -510,6 +518,13 @@ def _transform_system_message(
messages.pop(idx)
if len(system_content_blocks) > 0:
#########################################################
# If no messages are passed in, add a blank user message
# Relevant Issue - https://github.com/BerriAI/litellm/issues/13769
#########################################################
if len(messages) == 0:
messages.append(_default_user_message_when_system_message_passed())
#########################################################
return SystemInstructions(parts=system_content_blocks), messages
return None, messages

File diff suppressed because it is too large Load diff

View file

@ -2,7 +2,7 @@
Handler for transforming responses api requests to litellm.completion requests
"""
from typing import Any, Coroutine, Optional, Union
from typing import Any, Coroutine, Dict, Optional, Union
import litellm
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
@ -30,6 +30,7 @@ class LiteLLMCompletionTransformationHandler:
custom_llm_provider: Optional[str] = None,
_is_async: bool = False,
stream: Optional[bool] = None,
extra_headers: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Union[
ResponsesAPIResponse,
@ -45,6 +46,7 @@ class LiteLLMCompletionTransformationHandler:
responses_api_request=responses_api_request,
custom_llm_provider=custom_llm_provider,
stream=stream,
extra_headers=extra_headers,
**kwargs,
)
)

View file

@ -99,6 +99,7 @@ class LiteLLMCompletionResponsesConfig:
responses_api_request: ResponsesAPIOptionalRequestParams,
custom_llm_provider: Optional[str] = None,
stream: Optional[bool] = None,
extra_headers: Optional[Dict[str, Any]] = None,
**kwargs,
) -> dict:
"""
@ -126,6 +127,7 @@ class LiteLLMCompletionResponsesConfig:
"web_search_options": web_search_options,
# litellm specific params
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
}
# Responses API `Completed` events require usage, we pass `stream_options` to litellm.completion to include usage

View file

@ -455,6 +455,7 @@ def responses(
custom_llm_provider=custom_llm_provider,
_is_async=_is_async,
stream=stream,
extra_headers=extra_headers,
**kwargs,
)

View file

@ -651,14 +651,7 @@ async def test_image_edit_array_handling():
image=TEST_IMAGES,
)
# Test 3: Empty list (should fail validation)
with pytest.raises(Exception):
await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=[],
)
# Both valid calls should succeed
ImageResponse.model_validate(result1)
ImageResponse.model_validate(result2)
@ -667,117 +660,3 @@ async def test_image_edit_array_handling():
assert mock_post.call_count == 2
@pytest.mark.asyncio
async def test_openai_transformation_handles_multiple_images():
"""Test that OpenAI transformation correctly handles multiple images in request"""
from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
from litellm.types.router import GenericLiteLLMParams
config = OpenAIImageEditConfig()
# Test with multiple images
prompt = "Edit these images"
images = [b"fake_image_1", b"fake_image_2", b"fake_image_3"]
litellm_params = GenericLiteLLMParams(api_key="test_key")
data, files = config.transform_image_edit_request(
model="gpt-image-1",
prompt=prompt,
image=images,
image_edit_optional_request_params={"n": 1},
litellm_params=litellm_params,
headers={}
)
# Check that data contains the prompt and parameters
assert data["prompt"] == prompt
assert data["model"] == "gpt-image-1"
assert data["n"] == 1
# Check that files contains all images with correct field names
assert len(files) == len(images)
for i, file_entry in enumerate(files):
assert file_entry[0] == "image[]" # OpenAI uses image[] for multiple files
assert file_entry[1][1] == images[i] # Image data
assert file_entry[1][2] == "image/png" # Content type
print(f"Successfully processed {len(images)} images in transformation")
@pytest.mark.asyncio
async def test_multiple_image_edit_parameter_validation():
"""Test parameter validation with multiple images"""
from litellm import aimage_edit
# Mock response
mock_response = {
"created": 1589478378,
"data": [
{
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
}
]
}
class MockResponse:
def __init__(self, json_data, status_code):
self._json_data = json_data
self.status_code = status_code
self.text = json.dumps(json_data)
def json(self):
return self._json_data
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post:
mock_post.return_value = MockResponse(mock_response, 200)
# Test with valid parameters
result = await aimage_edit(
prompt="Test prompt",
model="gpt-image-1",
image=TEST_IMAGES,
n=1,
size="1024x1024",
response_format="b64_json"
)
ImageResponse.model_validate(result)
# Verify the request was made with correct parameters
mock_post.assert_called_once()
call_args = mock_post.call_args
# Check that the request contains the expected data
if 'data' in call_args.kwargs:
form_data = call_args.kwargs['data']
assert 'model' in form_data
assert 'prompt' in form_data
assert 'n' in form_data
assert form_data['n'] == 1 # Could be int or string depending on implementation print("Parameter validation passed for multiple image edit")
@pytest.mark.asyncio
async def test_multiple_image_edit_error_handling():
"""Test error handling with multiple images"""
from litellm import aimage_edit
# Test with None image (should raise error)
with pytest.raises(Exception):
await aimage_edit(
prompt="Test prompt",
model="gpt-image-1",
image=None,
)
# Test with invalid model (should raise error)
with pytest.raises(Exception):
await aimage_edit(
prompt="Test prompt",
model="invalid-model",
image=TEST_IMAGES,
)
print("Error handling tests passed for multiple image edit")

View file

@ -330,3 +330,78 @@ async def test_gpt_image_1_with_input_fidelity():
assert captured_kwargs["quality"] == "medium"
assert captured_kwargs["size"] == "1024x1024"
@pytest.mark.asyncio
async def test_aiml_image_generation_with_dynamic_api_key():
"""
Test that when api_key is passed as a dynamic parameter to aimage_generation,
it gets properly used for AIML provider authentication instead of falling back
to environment variables.
This test validates the fix for ensuring dynamic API keys are respected
when making image generation requests to the AIML provider.
"""
from unittest.mock import AsyncMock, patch, MagicMock
import httpx
# Mock AIML response
mock_aiml_response = {
"created": 1703658209,
"data": [
{
"url": "https://example.com/generated_image.png"
}
]
}
# Track captured arguments
captured_headers = None
captured_url = None
captured_json_data = None
def capture_post_call(*args, **kwargs):
nonlocal captured_headers, captured_url, captured_json_data
captured_url = kwargs.get('url') or (args[0] if args else None)
captured_headers = kwargs.get('headers', {})
captured_json_data = kwargs.get('json', {})
# Create a mock response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = mock_aiml_response
mock_response.text = json.dumps(mock_aiml_response)
return mock_response
# Mock the HTTP client that actually makes the request (sync version for image generation)
with patch('litellm.llms.custom_httpx.http_handler.HTTPHandler.post') as mock_post:
mock_post.side_effect = capture_post_call
# Test with dynamic api_key
test_api_key = "test-dynamic-api-key-12345"
response = await litellm.aimage_generation(
prompt="A cute baby sea otter",
model="aiml/flux-pro/v1.1",
api_key=test_api_key, # This should be used instead of env vars
)
# Validate the response (mocked response processing might not populate data correctly)
assert response is not None
# The most important validations: API key and endpoint usage
# These prove that the dynamic API key was properly used
assert captured_headers is not None
assert "Authorization" in captured_headers
assert captured_headers["Authorization"] == f"Bearer {test_api_key}"
print("TESTCAPTURED HEADERS", captured_headers)
# Validate the correct AIML endpoint was called
assert captured_url is not None
assert "api.aimlapi.com" in captured_url
assert "/v1/images/generations" in captured_url
# Validate the request data
assert captured_json_data is not None
assert captured_json_data["prompt"] == "A cute baby sea otter"
assert captured_json_data["model"] == "flux-pro/v1.1"

View file

@ -119,6 +119,28 @@ class BaseLLMChatTest(ABC):
pytest.skip("Model is overloaded")
assert response.choices[0].message.content is not None
def test_system_message_with_no_user_message(self):
"""
Test that the system message is translated correctly for non-OpenAI providers.
"""
base_completion_call_args = self.get_base_completion_call_args()
messages = [
{
"role": "system",
"content": "Be a good bot!",
},
]
try:
response = self.completion_function(
**base_completion_call_args,
messages=messages,
)
assert response is not None
except litellm.InternalServerError:
pytest.skip("Model is overloaded")
assert response.choices[0].message.content is not None
def test_content_list_handling(self):
"""Check if content list is supported by LLM API"""

View file

@ -1469,3 +1469,37 @@ def test_vertex_parallel_tool_calls_false_single_tool():
parallel_tool_calls=False,
)
assert "tools" in optional_params
from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body
def test_system_prompt_only_adds_blank_user_message():
"""
Test that the system prompt only adds a blank user message when a system message is passed in.
Relevant Issue - https://github.com/BerriAI/litellm/issues/13769
"""
SYSTEM_INSTRUCTION = "System instructions for the model"
data = _transform_request_body(
messages=[{"role": "system", "content": SYSTEM_INSTRUCTION}],
model="gemini-2.5-flash",
optional_params={},
custom_llm_provider="vertex_ai",
litellm_params={},
cached_content=None,
)
print("Final data: ", data)
# validate that a blank user message is added when a system message is passed in
assert len(data["contents"]) == 1
first_content = data["contents"][0]
assert first_content["role"] == "user"
assert len(first_content["parts"]) == 1
#########################################################
# system message was passed in
#########################################################
assert len(data["system_instruction"]) == 1
assert data["system_instruction"]["parts"][0]["text"] == SYSTEM_INSTRUCTION

View file

@ -541,7 +541,8 @@ class TestFunctionCallTransformation:
result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
model="gemini/gemini-2.0-flash",
input=test_input,
responses_api_request=responses_api_request
responses_api_request=responses_api_request,
extra_headers={"X-Test-Header": "test-value"}
)
assert "messages" in result
@ -563,6 +564,8 @@ class TestFunctionCallTransformation:
tool_msg = messages[2]
assert tool_msg["role"] == "tool"
assert result["extra_headers"] == {"X-Test-Header": "test-value"}
def test_function_call_without_call_id_fallback_to_id(self):
"""Test that function_call items can use 'id' field when 'call_id' is missing"""
function_call_item = {

View file

@ -847,6 +847,7 @@ async def test_supports_tool_choice():
or "o3" in model_name
or "mistral" in model_name
or "oci" in model_name
or "openrouter" in model_name
):
continue