scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-26 06:03:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

Merge 8e5ce6d66c into d7620c77a1

Browse source
This commit is contained in:
devin-ai-integration[bot] 2026-10-19 19:09:53 +00:00 • committed by GitHub
commit 76c3511b71
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 201 additions and 180 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/images/main.py
View file

@ -422,6 +422,7 @@ def image_generation(
logging_obj=litellm_logging_obj,
timeout=timeout,
client=client,
_is_async=aimg_generation is True,
)
elif custom_llm_provider == "black_forest_labs":
# Route to BFL-specific handler (polling required)

26
litellm/llms/base_llm/image_generation/transformation.py
View file

@ -100,6 +100,32 @@ class BaseImageGenerationConfig(ABC):
"ImageVariationConfig implements 'transform_response_image_variation' for image variation models"
)
async def async_transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict[str, object], # mutable-ok: mirrors the sync hook so the default passes the same objects
optional_params: dict[str, object], # mutable-ok: mirrors the sync hook so the default passes the same objects
litellm_params: dict[str, object], # mutable-ok: mirrors the sync hook so the default passes the same objects
encoding: "Tokenizer | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ImageResponse:
return self.transform_image_generation_response( # pyright: ignore[reportUnknownMemberType] # sync hook signature is untyped
model=model,
raw_response=raw_response,
model_response=model_response,
logging_obj=logging_obj,
request_data=request_data,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
api_key=api_key,
json_mode=json_mode,
)
def use_multipart_form_data(self) -> bool:
"""
Returns True if this provider requires multipart/form-data instead of JSON.

12
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -44,7 +44,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import (
DEFAULT_MAX_AGENTIC_LOOPS,
validated_max_agentic_loops,
)
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.asyncify import asyncify, run_async_function
from litellm.litellm_core_utils.audio_utils.subtitle_utils import (
SUBTITLE_RESPONSE_FORMATS,
synthesize_subtitle_document,
@ -7464,7 +7464,7 @@ class BaseLLMHTTPHandler:
else:
async_httpx_client = client
headers: Final = image_generation_provider_config.validate_environment(
headers: Final = await asyncify(image_generation_provider_config.validate_environment)( # pyright: ignore[reportUnknownArgumentType] # sync hook signature is untyped
api_key=api_key,
headers=image_generation_optional_request_params.get("extra_headers", {}) or {},
model=model,
@ -7529,13 +7529,15 @@ class BaseLLMHTTPHandler:
provider_config=image_generation_provider_config,
)
model_response: Final[ImageResponse] = image_generation_provider_config.transform_image_generation_response(
model_response: Final[
ImageResponse
] = await image_generation_provider_config.async_transform_image_generation_response(
model=model,
raw_response=response,
model_response=litellm.ImageResponse(),
logging_obj=logging_obj,
request_data=data,
optional_params=image_generation_optional_request_params,
request_data=data, # pyright: ignore[reportUnknownArgumentType] # handler params are untyped
optional_params=image_generation_optional_request_params, # pyright: ignore[reportUnknownArgumentType] # handler params are untyped
litellm_params=dict(litellm_params),
encoding=None,
)

182
tests/unit/images/test_main.py
View file

@ -1,6 +1,8 @@
import asyncio
import itertools
import json
from datetime import datetime
from collections.abc import Coroutine
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -13,6 +15,7 @@ from openai.types.image import Image
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.types.utils import CallTypes, StandardLoggingPayload
import os
@ -249,122 +252,59 @@ class TestAimlImageGeneration:
pytest.fail(f"An exception occurred - {str(e)}")
def _capturing_client(response_json: dict[str, object]) -> tuple[AsyncHTTPHandler, list[httpx.Request]]:
requests: Final[list[httpx.Request]] = []
def handler(request: httpx.Request) -> httpx.Response:
requests.append(request)
return httpx.Response(200, json=response_json)
return AsyncHTTPHandler(transport=httpx.MockTransport(handler)), requests
@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.
client, requests = _capturing_client(
{"created": 1703658209, "data": [{"url": "https://example.com/generated_image.png"}]}
)
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, MagicMock, patch
await litellm.aimage_generation(
prompt="A cute baby sea otter",
model="aiml/flux-pro/v1.1",
api_key="test-dynamic-api-key-12345",
client=client,
)
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"
assert len(requests) == 1
assert requests[0].headers["Authorization"] == "Bearer test-dynamic-api-key-12345"
assert requests[0].url.host == "api.aimlapi.com"
assert requests[0].url.path.endswith("/v1/images/generations")
body: Final = json.loads(requests[0].content)
assert body["prompt"] == "A cute baby sea otter"
assert body["model"] == "flux-pro/v1.1"
@pytest.mark.asyncio
async def test_aiml_openai_gpt_image_2_request_uses_openai_param_shape():
"""End-to-end check that ``aiml/openai/gpt-image-2`` keeps the upstream
OpenAI request shape (``size``/``n``/``response_format``) instead of
being remapped to the AI/ML flux schema (``image_size``/``num_images``/
``output_format``), and hits the correct upstream model name.
"""
import json as _json
from unittest.mock import MagicMock, patch
client, requests = _capturing_client(
{"created": 1703658209, "data": [{"url": "https://example.com/gpt-image-2.png"}]}
)
mock_aiml_response = {
"created": 1703658209,
"data": [{"url": "https://example.com/gpt-image-2.png"}],
}
await litellm.aimage_generation(
prompt="A T-Rex relaxing on a beach",
model="aiml/openai/gpt-image-2",
api_key="test-key-mocked-no-credits-needed",
size="1024x1536",
quality="high",
response_format="b64_json",
n=1,
client=client,
)
captured = {}
def capture_post_call(*args, **kwargs):
captured["url"] = kwargs.get("url") or (args[0] if args else None)
captured["headers"] = kwargs.get("headers", {})
captured["json"] = kwargs.get("json", {})
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
with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post:
mock_post.side_effect = capture_post_call
await litellm.aimage_generation(
prompt="A T-Rex relaxing on a beach",
model="aiml/openai/gpt-image-2",
api_key="test-key-mocked-no-credits-needed",
size="1024x1536",
quality="high",
response_format="b64_json",
n=1,
)
assert captured["url"] is not None
assert "api.aimlapi.com" in captured["url"]
assert "/v1/images/generations" in captured["url"]
body = captured["json"]
assert len(requests) == 1
assert requests[0].url.host == "api.aimlapi.com"
assert requests[0].url.path.endswith("/v1/images/generations")
body: Final = json.loads(requests[0].content)
assert body["model"] == "openai/gpt-image-2"
assert body["prompt"] == "A T-Rex relaxing on a beach"
assert body["size"] == "1024x1536"
@ -405,3 +345,33 @@ async def test_azure_image_generation_request_body():
call_args = mock_post.call_args
request_json = call_args.kwargs.get("json", {})
assert request_json == expected_body
@pytest.mark.asyncio
async def test_aimage_generation_runs_vertex_requests_concurrently_on_the_event_loop() -> None:
response_json: Final = {
"candidates": [{"content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": "aW1n"}}]}}],
}
arrivals: Final = itertools.count(1)
both_in_flight: Final = asyncio.Event()
async def handler(request: httpx.Request) -> httpx.Response:
if next(arrivals) == 2:
both_in_flight.set()
await both_in_flight.wait()
return httpx.Response(status_code=200, json=response_json)
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(handler))
def generate() -> Coroutine[object, object, litellm.ImageResponse]:
return litellm.aimage_generation( # pyright: ignore[reportUnknownMemberType] # aimage_generation kwargs are untyped
model="vertex_ai/gemini-2.5-flash-image",
prompt="a red circle",
api_base="http://localhost:1/generateContent",
vertex_location="us-central1",
client=client,
)
results: Final = await asyncio.wait_for(asyncio.gather(generate(), generate()), timeout=5)
assert [result.data[0].b64_json if result.data else None for result in results] == ["aW1n", "aW1n"]

57
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -64,6 +64,7 @@ from litellm.responses.streaming_iterator import (
)
from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse
from litellm.types.router import GenericLiteLLMParams
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse
from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe
@ -4902,6 +4903,62 @@ async def test_async_image_generation_handler_records_upstream_response_headers(
_assert_upstream_headers_recorded(response)
class _OffLoopImageGenerationConfig(_HeaderImageGenerationConfig):
def validate_environment(
self,
headers: dict[str, object],
model: str,
messages: list[AllMessageValues],
optional_params: dict[str, object],
litellm_params: dict[str, object],
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, object]:
try:
asyncio.get_running_loop()
except RuntimeError:
return {"Authorization": "Bearer from-worker-thread"}
raise AssertionError("validate_environment ran on the event loop")
def transform_image_generation_response(self, *args: object, **kwargs: object) -> ImageResponse:
raise AssertionError("sync transform_image_generation_response ran on the event loop")
async def async_transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: object,
request_data: dict[str, object],
optional_params: dict[str, object],
litellm_params: dict[str, object],
encoding: object,
api_key: str | None = None,
json_mode: bool | None = None,
) -> ImageResponse:
return ImageResponse(data=[ImageObject(b64_json=raw_response.json()["b64_json"])])
@pytest.mark.asyncio
async def test_async_image_generation_handler_keeps_provider_auth_and_polling_off_the_event_loop():
seen_authorization: Final[list[str | None]] = []
def handler(request: httpx.Request) -> httpx.Response:
seen_authorization.append(request.headers.get("Authorization"))
return httpx.Response(200, json={"b64_json": "abc"})
client: Final = AsyncHTTPHandler()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
response: Final = await BaseLLMHTTPHandler().async_image_generation_handler( # pyright: ignore[reportUnknownMemberType] # handler params are untyped
client=client,
**{**_image_generation_call_kwargs(), "image_generation_provider_config": _OffLoopImageGenerationConfig()}, # pyright: ignore[reportUnknownArgumentType] # call kwargs helper is untyped
)
assert response.data[0].b64_json == "abc" # pyright: ignore[reportUnknownMemberType] # ImageObject fields are partially untyped
assert seen_authorization == ["Bearer from-worker-thread"]
class _HeaderTextToSpeechConfig(BaseTextToSpeechConfig):
def get_supported_openai_params(self, model):
return []

103
tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation.py
View file

@ -1,8 +1,11 @@
from unittest.mock import MagicMock, patch
import json
from typing import Final
import httpx
import pytest
from litellm import aimage_generation
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@ -22,78 +25,40 @@ from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outco
],
)
@pytest.mark.asyncio
async def test_fal_ai_image_generation_basic(model, expected_endpoint):
"""
Test that fal_ai image generation constructs correct request body and URL.
async def test_fal_ai_image_generation_basic(model: str, expected_endpoint: str) -> None:
requests: Final[list[httpx.Request]] = []
Validates:
- Correct API endpoint URL construction
- Proper request body format with prompt
- Correct Authorization header format
"""
captured_url = None
captured_json_data = None
captured_headers = None
def capture_post_call(*args, **kwargs):
nonlocal captured_url, captured_json_data, captured_headers
captured_url = args[0] if args else kwargs.get("url")
captured_json_data = kwargs.get("json")
captured_headers = kwargs.get("headers")
# Mock response with fal.ai format
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.json.return_value = {
"images": [
{
"url": "https://example.com/generated-image.png",
"width": 1024,
"height": 768,
"content_type": "image/jpeg",
}
],
"seed": 42,
}
return mock_response
with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post:
mock_post.side_effect = capture_post_call
test_api_key = "test-fal-ai-key-12345"
test_prompt = "A cute baby sea otter"
response = await aimage_generation(
model=model,
prompt=test_prompt,
api_key=test_api_key,
def handler(request: httpx.Request) -> httpx.Response:
requests.append(request)
return httpx.Response(
200,
json={
"images": [
{
"url": "https://example.com/generated-image.png",
"width": 1024,
"height": 768,
"content_type": "image/jpeg",
}
],
"seed": 42,
},
)
# Validate response
assert response is not None
assert hasattr(response, "data")
assert response.data is not None
assert len(response.data) > 0
response: Final = await aimage_generation(
model=model,
prompt="A cute baby sea otter",
api_key="test-fal-ai-key-12345",
client=AsyncHTTPHandler(transport=httpx.MockTransport(handler)),
)
# Validate URL
assert captured_url is not None
assert "fal.run" in captured_url
assert expected_endpoint in captured_url
print(f"Validated URL: {captured_url}")
# Validate headers
assert captured_headers is not None
assert "Authorization" in captured_headers
assert captured_headers["Authorization"] == f"Key {test_api_key}"
print(f"Validated headers: {captured_headers}")
# Validate request body
assert captured_json_data is not None
assert captured_json_data["prompt"] == test_prompt
print(f"Validated request body: {captured_json_data}")
assert response.data
assert response.data[0].url == "https://example.com/generated-image.png"
assert len(requests) == 1
assert requests[0].url.host == "fal.run"
assert expected_endpoint in requests[0].url.path
assert requests[0].headers["Authorization"] == "Key test-fal-ai-key-12345"
assert json.loads(requests[0].content)["prompt"] == "A cute baby sea otter"
@pytest.fixture(autouse=True)