scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 178

refactor(bfl): separate HTTP logic into dedicated handlers

Browse source 
- Create handler.py for image generation and image edit
- Move polling logic from transformation to handlers
- Handlers use _get_httpx_client() / get_async_httpx_client()
- Transformation files now only transform request/response data
- Follows Bedrock pattern for provider-specific handlers

Addresses feedback: transformation files should not make HTTP requests
This commit is contained in:
Chesars 2026-01-19 12:07:33 -03:00
parent a14a79619a
commit 727fd76841
9 changed files with 1076 additions and 889 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
litellm/images/main.py
View file

@ -50,6 +50,10 @@ from litellm.main import (
openai_image_variations,
)
# BFL handlers
from litellm.llms.black_forest_labs.image_edit.handler import bfl_image_edit
from litellm.llms.black_forest_labs.image_generation.handler import bfl_image_generation
###########################################
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
@ -405,7 +409,6 @@ def image_generation( # noqa: PLR0915
litellm.LlmProviders.RUNWAYML,
litellm.LlmProviders.VERTEX_AI,
litellm.LlmProviders.OPENROUTER,
litellm.LlmProviders.BLACK_FOREST_LABS,
):
if image_generation_config is None:
raise ValueError(
@ -428,6 +431,22 @@ def image_generation( # noqa: PLR0915
timeout=timeout,
client=client,
)
elif custom_llm_provider == "black_forest_labs":
# Route to BFL-specific handler (polling required)
if model is None:
raise Exception("Model needs to be set for black_forest_labs")
return bfl_image_generation.image_generation(
model=model,
prompt=prompt,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params_dict,
logging_obj=litellm_logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=client,
aimg_generation=aimg_generation,
)
elif custom_llm_provider == "azure_ai":
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
@ -921,6 +940,23 @@ def image_edit( # noqa: PLR0915
_is_async=_is_async,
client=kwargs.get("client"),
)
elif custom_llm_provider == "black_forest_labs":
# Route to BFL-specific handler (polling required)
if model is None:
raise Exception("Model needs to be set for black_forest_labs")
image_edit_request_params.update(non_default_params)
return bfl_image_edit.image_edit(
model=model,
image=images,
prompt=prompt,
image_edit_optional_request_params=image_edit_request_params,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
timeout=timeout or DEFAULT_REQUEST_TIMEOUT,
extra_headers=extra_headers,
client=kwargs.get("client"),
aimage_edit=_is_async,
)
# Call the handler with _is_async flag instead of directly calling the async handler
return base_llm_http_handler.image_edit_handler(
model=model,

7
litellm/llms/black_forest_labs/image_edit/__init__.py
View file

@ -1,3 +1,8 @@
from .handler import BlackForestLabsImageEdit, bfl_image_edit
from .transformation import BlackForestLabsImageEditConfig
__all__ = ["BlackForestLabsImageEditConfig"]
__all__ = [
"BlackForestLabsImageEditConfig",
"BlackForestLabsImageEdit",
"bfl_image_edit",
]

427
litellm/llms/black_forest_labs/image_edit/handler.py Normal file
View file

@ -0,0 +1,427 @@
"""
Black Forest Labs Image Edit Handler
Handles image edit requests for Black Forest Labs models.
BFL uses an async polling pattern - the initial request returns a task ID,
then we poll until the result is ready.
"""
import asyncio
import time
from typing import Any, Dict, List, Optional, Union
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
_get_httpx_client,
get_async_httpx_client,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import FileTypes, ImageResponse
from ..common_utils import (
DEFAULT_MAX_POLLING_TIME,
DEFAULT_POLLING_INTERVAL,
BlackForestLabsError,
)
from .transformation import BlackForestLabsImageEditConfig
class BlackForestLabsImageEdit:
"""
Black Forest Labs Image Edit handler.
Handles the HTTP requests and polling logic, delegating data transformation
to the BlackForestLabsImageEditConfig class.
"""
def __init__(self):
self.config = BlackForestLabsImageEditConfig()
def image_edit(
self,
model: str,
image: Union[FileTypes, List[FileTypes]],
prompt: Optional[str],
image_edit_optional_request_params: Dict,
litellm_params: Union[GenericLiteLLMParams, Dict],
logging_obj: LiteLLMLoggingObj,
timeout: Optional[Union[float, httpx.Timeout]],
extra_headers: Optional[Dict[str, Any]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
aimage_edit: bool = False,
) -> Union[ImageResponse, Any]:
"""
Main entry point for image edit requests.
Args:
model: The model to use (e.g., "black_forest_labs/flux-kontext-pro")
image: The image(s) to edit
prompt: The edit instruction
image_edit_optional_request_params: Optional parameters for the request
litellm_params: LiteLLM parameters including api_key, api_base
logging_obj: Logging object
timeout: Request timeout
extra_headers: Additional headers
client: HTTP client to use
aimage_edit: If True, return async coroutine
Returns:
ImageResponse or coroutine if aimage_edit=True
"""
# Handle litellm_params as dict or object
if isinstance(litellm_params, dict):
api_key = litellm_params.get("api_key")
api_base = litellm_params.get("api_base")
litellm_params_dict = litellm_params
else:
api_key = litellm_params.api_key
api_base = litellm_params.api_base
litellm_params_dict = dict(litellm_params)
if aimage_edit:
return self.async_image_edit(
model=model,
image=image,
prompt=prompt,
image_edit_optional_request_params=image_edit_optional_request_params,
litellm_params=litellm_params,
logging_obj=logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=client if isinstance(client, AsyncHTTPHandler) else None,
)
# Sync version
if client is None or not isinstance(client, HTTPHandler):
sync_client = _get_httpx_client()
else:
sync_client = client
# Validate environment and get headers
headers = self.config.validate_environment(
api_key=api_key,
headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
model=model,
)
if extra_headers:
headers.update(extra_headers)
# Get complete URL
complete_url = self.config.get_complete_url(
model=model,
api_base=api_base,
litellm_params=litellm_params_dict,
)
# Transform request
# Handle image list vs single image
image_input = image[0] if isinstance(image, list) and image else image
data, _ = self.config.transform_image_edit_request(
model=model,
prompt=prompt or "",
image=image_input,
image_edit_optional_request_params=image_edit_optional_request_params,
litellm_params=litellm_params_dict,
headers=headers,
)
# Logging
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": complete_url,
"headers": headers,
},
)
# Make initial request
try:
response = sync_client.post(
url=complete_url,
headers=headers,
json=data,
timeout=timeout,
)
except Exception as e:
raise BlackForestLabsError(
status_code=500,
message=f"Request failed: {str(e)}",
)
# Poll for result
final_response = self._poll_for_result_sync(
initial_response=response,
headers=headers,
sync_client=sync_client,
)
# Transform response
return self.config.transform_image_edit_response(
model=model,
raw_response=final_response,
logging_obj=logging_obj,
)
async def async_image_edit(
self,
model: str,
image: Union[FileTypes, List[FileTypes]],
prompt: Optional[str],
image_edit_optional_request_params: Dict,
litellm_params: Union[GenericLiteLLMParams, Dict],
logging_obj: LiteLLMLoggingObj,
timeout: Optional[Union[float, httpx.Timeout]],
extra_headers: Optional[Dict[str, Any]] = None,
client: Optional[AsyncHTTPHandler] = None,
) -> ImageResponse:
"""
Async version of image edit.
"""
# Handle litellm_params as dict or object
if isinstance(litellm_params, dict):
api_key = litellm_params.get("api_key")
api_base = litellm_params.get("api_base")
litellm_params_dict = litellm_params
else:
api_key = litellm_params.api_key
api_base = litellm_params.api_base
litellm_params_dict = dict(litellm_params)
if client is None:
async_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.BLACK_FOREST_LABS,
)
else:
async_client = client
# Validate environment and get headers
headers = self.config.validate_environment(
api_key=api_key,
headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
model=model,
)
if extra_headers:
headers.update(extra_headers)
# Get complete URL
api_base = self.config.get_complete_url(
model=model,
api_base=api_base,
litellm_params=litellm_params_dict,
)
# Transform request
image_input = image[0] if isinstance(image, list) and image else image
data, _ = self.config.transform_image_edit_request(
model=model,
prompt=prompt or "",
image=image_input,
image_edit_optional_request_params=image_edit_optional_request_params,
litellm_params=litellm_params_dict,
headers=headers,
)
# Logging
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": headers,
},
)
# Make initial request
try:
response = await async_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout,
)
except Exception as e:
raise BlackForestLabsError(
status_code=500,
message=f"Request failed: {str(e)}",
)
# Poll for result
final_response = await self._poll_for_result_async(
initial_response=response,
headers=headers,
async_client=async_client,
)
# Transform response
return self.config.transform_image_edit_response(
model=model,
raw_response=final_response,
logging_obj=logging_obj,
)
def _poll_for_result_sync(
self,
initial_response: httpx.Response,
headers: dict,
sync_client: HTTPHandler,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
) -> httpx.Response:
"""
Poll BFL API until result is ready (sync version).
Args:
initial_response: The initial response containing polling_url
headers: Headers to use for polling (must include x-key)
sync_client: HTTP client
max_wait: Maximum time to wait in seconds
interval: Polling interval in seconds
Returns:
Final response with completed result
"""
# Parse initial response to get polling URL
try:
response_data = initial_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"Error parsing initial response: {e}",
)
# Check for immediate errors
if "errors" in response_data:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL error: {response_data['errors']}",
)
polling_url = response_data.get("polling_url")
if not polling_url:
raise BlackForestLabsError(
status_code=500,
message="No polling_url in BFL response",
)
# Get just the auth header for polling
polling_headers = {"x-key": headers.get("x-key", "")}
start_time = time.time()
verbose_logger.debug(f"BFL starting sync polling at {polling_url}")
while time.time() - start_time < max_wait:
response = sync_client.get(
url=polling_url,
headers=polling_headers,
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
verbose_logger.debug(f"BFL poll status: {status}")
if status == "Ready":
return response
elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
time.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Polling timed out after {max_wait} seconds",
)
async def _poll_for_result_async(
self,
initial_response: httpx.Response,
headers: dict,
async_client: AsyncHTTPHandler,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
) -> httpx.Response:
"""
Poll BFL API until result is ready (async version).
"""
# Parse initial response to get polling URL
try:
response_data = initial_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"Error parsing initial response: {e}",
)
# Check for immediate errors
if "errors" in response_data:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL error: {response_data['errors']}",
)
polling_url = response_data.get("polling_url")
if not polling_url:
raise BlackForestLabsError(
status_code=500,
message="No polling_url in BFL response",
)
# Get just the auth header for polling
polling_headers = {"x-key": headers.get("x-key", "")}
start_time = time.time()
verbose_logger.debug(f"BFL starting async polling at {polling_url}")
while time.time() - start_time < max_wait:
response = await async_client.get(
url=polling_url,
headers=polling_headers,
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
verbose_logger.debug(f"BFL poll status: {status}")
if status == "Ready":
return response
elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
await asyncio.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Polling timed out after {max_wait} seconds",
)
# Singleton instance for use in images/main.py
bfl_image_edit = BlackForestLabsImageEdit()

89
litellm/llms/black_forest_labs/image_edit/transformation.py
View file

@ -9,13 +9,12 @@ API Reference: https://docs.bfl.ai/
import base64
import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
from httpx._types import RequestFiles
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
@ -23,8 +22,6 @@ from litellm.types.utils import FileTypes, ImageObject, ImageResponse
from ..common_utils import (
DEFAULT_API_BASE,
DEFAULT_MAX_POLLING_TIME,
DEFAULT_POLLING_INTERVAL,
IMAGE_EDIT_MODELS,
BlackForestLabsError,
)
@ -46,6 +43,9 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
- flux-kontext-max: Premium quality editing
- flux-pro-1.0-fill: Inpainting with mask
- flux-pro-1.0-expand: Outpainting (expand image borders)
Note: HTTP requests and polling are handled by the handler (handler.py).
This class only handles data transformation.
"""
def get_supported_openai_params(self, model: str) -> List[str]:
@ -234,52 +234,6 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
# BFL uses JSON, not multipart - return empty files
return request_body, []
def _poll_for_result(
self,
polling_url: str,
api_key: str,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
) -> Dict:
"""
Poll the BFL API until the result is ready.
Returns the result data when status is "Ready".
Raises BlackForestLabsError on failure.
"""
start_time = time.time()
httpx_client = _get_httpx_client()
while time.time() - start_time < max_wait:
response = httpx_client.get(
polling_url,
headers={"x-key": api_key},
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
if status == "Ready":
return data
elif status in ["Error", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
time.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Timeout waiting for result after {max_wait} seconds",
)
def transform_image_edit_response(
self,
model: str,
@ -289,7 +243,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
"""
Transform Black Forest Labs response to OpenAI-compatible ImageResponse.
BFL returns a task ID initially, then we poll until the result is ready.
This is called with the FINAL polled response (after handler does polling).
The response contains: {"status": "Ready", "result": {"sample": "https://..."}}
"""
try:
response_data = raw_response.json()
@ -299,29 +254,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
message=f"Error parsing BFL response: {e}",
)
# Check for immediate errors
if "errors" in response_data:
raise BlackForestLabsError(
status_code=raw_response.status_code,
message=f"BFL error: {response_data['errors']}",
)
# Get polling URL
polling_url = response_data.get("polling_url")
if not polling_url:
raise BlackForestLabsError(
status_code=500,
message="No polling_url in BFL response",
)
# Extract API key from original request headers
api_key = raw_response.request.headers.get("x-key", "")
# Poll for result
result_data = self._poll_for_result(polling_url, api_key)
# Get image URL from result
image_url = result_data.get("result", {}).get("sample")
image_url = response_data.get("result", {}).get("sample")
if not image_url:
raise BlackForestLabsError(
status_code=500,
@ -333,3 +267,12 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
created=int(time.time()),
data=[ImageObject(url=image_url)],
)
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BlackForestLabsError:
"""Return the appropriate error class for Black Forest Labs."""
return BlackForestLabsError(
status_code=status_code,
message=error_message,
)

3
litellm/llms/black_forest_labs/image_generation/__init__.py
View file

@ -1,3 +1,4 @@
from .handler import BlackForestLabsImageGeneration, bfl_image_generation
from .transformation import (
BlackForestLabsImageGenerationConfig,
get_black_forest_labs_image_generation_config,
@ -6,4 +7,6 @@ from .transformation import (
__all__ = [
"BlackForestLabsImageGenerationConfig",
"get_black_forest_labs_image_generation_config",
"BlackForestLabsImageGeneration",
"bfl_image_generation",
]

424
litellm/llms/black_forest_labs/image_generation/handler.py Normal file
View file

@ -0,0 +1,424 @@
"""
Black Forest Labs Image Generation Handler
Handles image generation requests for Black Forest Labs models.
BFL uses an async polling pattern - the initial request returns a task ID,
then we poll until the result is ready.
"""
import asyncio
import time
from typing import Any, Dict, Optional, Union
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
_get_httpx_client,
get_async_httpx_client,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageResponse
from ..common_utils import (
DEFAULT_MAX_POLLING_TIME,
DEFAULT_POLLING_INTERVAL,
BlackForestLabsError,
)
from .transformation import BlackForestLabsImageGenerationConfig
class BlackForestLabsImageGeneration:
"""
Black Forest Labs Image Generation handler.
Handles the HTTP requests and polling logic, delegating data transformation
to the BlackForestLabsImageGenerationConfig class.
"""
def __init__(self):
self.config = BlackForestLabsImageGenerationConfig()
def image_generation(
self,
model: str,
prompt: str,
model_response: ImageResponse,
optional_params: Dict,
litellm_params: Union[GenericLiteLLMParams, Dict],
logging_obj: LiteLLMLoggingObj,
timeout: Optional[Union[float, httpx.Timeout]],
extra_headers: Optional[Dict[str, Any]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
aimg_generation: bool = False,
) -> Union[ImageResponse, Any]:
"""
Main entry point for image generation requests.
Args:
model: The model to use (e.g., "black_forest_labs/flux-pro-1.1")
prompt: The text prompt for image generation
model_response: ImageResponse object to populate
optional_params: Optional parameters for the request
litellm_params: LiteLLM parameters including api_key, api_base
logging_obj: Logging object
timeout: Request timeout
extra_headers: Additional headers
client: HTTP client to use
aimg_generation: If True, return async coroutine
Returns:
ImageResponse or coroutine if aimg_generation=True
"""
# Handle litellm_params as dict or object
if isinstance(litellm_params, dict):
api_key = litellm_params.get("api_key")
api_base = litellm_params.get("api_base")
litellm_params_dict = litellm_params
else:
api_key = litellm_params.api_key
api_base = litellm_params.api_base
litellm_params_dict = dict(litellm_params)
if aimg_generation:
return self.async_image_generation(
model=model,
prompt=prompt,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
logging_obj=logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=client if isinstance(client, AsyncHTTPHandler) else None,
)
# Sync version
if client is None or not isinstance(client, HTTPHandler):
sync_client = _get_httpx_client()
else:
sync_client = client
# Validate environment and get headers
headers = self.config.validate_environment(
api_key=api_key,
headers={},
model=model,
messages=[],
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
if extra_headers:
headers.update(extra_headers)
# Get complete URL
complete_url = self.config.get_complete_url(
api_base=api_base,
api_key=api_key,
model=model,
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
# Transform request
data = self.config.transform_image_generation_request(
model=model,
prompt=prompt,
optional_params=optional_params,
litellm_params=litellm_params_dict,
headers=headers,
)
# Logging
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": complete_url,
"headers": headers,
},
)
# Make initial request
try:
response = sync_client.post(
url=complete_url,
headers=headers,
json=data,
timeout=timeout,
)
except Exception as e:
raise BlackForestLabsError(
status_code=500,
message=f"Request failed: {str(e)}",
)
# Poll for result
final_response = self._poll_for_result_sync(
initial_response=response,
headers=headers,
sync_client=sync_client,
)
# Transform response
return self.config.transform_image_generation_response(
model=model,
raw_response=final_response,
model_response=model_response,
logging_obj=logging_obj,
)
async def async_image_generation(
self,
model: str,
prompt: str,
model_response: ImageResponse,
optional_params: Dict,
litellm_params: Union[GenericLiteLLMParams, Dict],
logging_obj: LiteLLMLoggingObj,
timeout: Optional[Union[float, httpx.Timeout]],
extra_headers: Optional[Dict[str, Any]] = None,
client: Optional[AsyncHTTPHandler] = None,
) -> ImageResponse:
"""
Async version of image generation.
"""
# Handle litellm_params as dict or object
if isinstance(litellm_params, dict):
api_key = litellm_params.get("api_key")
api_base = litellm_params.get("api_base")
litellm_params_dict = litellm_params
else:
api_key = litellm_params.api_key
api_base = litellm_params.api_base
litellm_params_dict = dict(litellm_params)
if client is None:
async_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.BLACK_FOREST_LABS,
)
else:
async_client = client
# Validate environment and get headers
headers = self.config.validate_environment(
api_key=api_key,
headers={},
model=model,
messages=[],
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
if extra_headers:
headers.update(extra_headers)
# Get complete URL
complete_url = self.config.get_complete_url(
api_base=api_base,
api_key=api_key,
model=model,
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
# Transform request
data = self.config.transform_image_generation_request(
model=model,
prompt=prompt,
optional_params=optional_params,
litellm_params=litellm_params_dict,
headers=headers,
)
# Logging
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": complete_url,
"headers": headers,
},
)
# Make initial request
try:
response = await async_client.post(
url=complete_url,
headers=headers,
json=data,
timeout=timeout,
)
except Exception as e:
raise BlackForestLabsError(
status_code=500,
message=f"Request failed: {str(e)}",
)
# Poll for result
final_response = await self._poll_for_result_async(
initial_response=response,
headers=headers,
async_client=async_client,
)
# Transform response
return self.config.transform_image_generation_response(
model=model,
raw_response=final_response,
model_response=model_response,
logging_obj=logging_obj,
)
def _poll_for_result_sync(
self,
initial_response: httpx.Response,
headers: dict,
sync_client: HTTPHandler,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
) -> httpx.Response:
"""
Poll BFL API until result is ready (sync version).
"""
# Parse initial response to get polling URL
try:
response_data = initial_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"Error parsing initial response: {e}",
)
# Check for immediate errors
if "errors" in response_data:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL error: {response_data['errors']}",
)
polling_url = response_data.get("polling_url")
if not polling_url:
raise BlackForestLabsError(
status_code=500,
message="No polling_url in BFL response",
)
# Get just the auth header for polling
polling_headers = {"x-key": headers.get("x-key", "")}
start_time = time.time()
verbose_logger.debug(f"BFL starting sync polling at {polling_url}")
while time.time() - start_time < max_wait:
response = sync_client.get(
url=polling_url,
headers=polling_headers,
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
verbose_logger.debug(f"BFL poll status: {status}")
if status == "Ready":
return response
elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
time.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Polling timed out after {max_wait} seconds",
)
async def _poll_for_result_async(
self,
initial_response: httpx.Response,
headers: dict,
async_client: AsyncHTTPHandler,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
) -> httpx.Response:
"""
Poll BFL API until result is ready (async version).
"""
# Parse initial response to get polling URL
try:
response_data = initial_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"Error parsing initial response: {e}",
)
# Check for immediate errors
if "errors" in response_data:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL error: {response_data['errors']}",
)
polling_url = response_data.get("polling_url")
if not polling_url:
raise BlackForestLabsError(
status_code=500,
message="No polling_url in BFL response",
)
# Get just the auth header for polling
polling_headers = {"x-key": headers.get("x-key", "")}
start_time = time.time()
verbose_logger.debug(f"BFL starting async polling at {polling_url}")
while time.time() - start_time < max_wait:
response = await async_client.get(
url=polling_url,
headers=polling_headers,
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
verbose_logger.debug(f"BFL poll status: {status}")
if status == "Ready":
return response
elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
await asyncio.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Polling timed out after {max_wait} seconds",
)
# Singleton instance for use in images/main.py
bfl_image_generation = BlackForestLabsImageGeneration()

225
litellm/llms/black_forest_labs/image_generation/transformation.py
View file

@ -7,21 +7,14 @@ for image generation endpoints (flux-pro-1.1, flux-pro-1.1-ultra, flux-dev, flux
API Reference: https://docs.bfl.ai/
"""
import asyncio
import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
@ -31,8 +24,6 @@ from litellm.types.utils import ImageObject, ImageResponse
from ..common_utils import (
DEFAULT_API_BASE,
DEFAULT_MAX_POLLING_TIME,
DEFAULT_POLLING_INTERVAL,
IMAGE_GENERATION_MODELS,
BlackForestLabsError,
)
@ -54,6 +45,9 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig):
- flux-pro-1.1-ultra: Ultra high-resolution (up to 4MP)
- flux-dev: Development/open-source variant
- flux-pro: Original pro model
Note: HTTP requests and polling are handled by the handler (handler.py).
This class only handles data transformation.
"""
def get_supported_openai_params(
@ -249,111 +243,28 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig):
return request_body
def _poll_for_result(
def transform_image_generation_response(
self,
polling_url: str,
api_key: str,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
) -> Dict:
"""
Poll the BFL API until the result is ready.
Returns the result data when status is "Ready".
Raises BlackForestLabsError on failure.
"""
start_time = time.time()
httpx_client = _get_httpx_client()
while time.time() - start_time < max_wait:
response = httpx_client.get(
polling_url,
headers={"x-key": api_key},
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
if status == "Ready":
return data
elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
time.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Timeout waiting for result after {max_wait} seconds",
)
async def _poll_for_result_async(
self,
polling_url: str,
api_key: str,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
) -> Dict:
"""
Poll the BFL API until the result is ready (async version).
Returns the result data when status is "Ready".
Raises BlackForestLabsError on failure.
"""
start_time = time.time()
httpx_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.BLACK_FOREST_LABS
)
while time.time() - start_time < max_wait:
response = await httpx_client.get(
polling_url,
headers={"x-key": api_key},
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
verbose_logger.debug(f"BFL polling status: {status}")
if status == "Ready":
return data
elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
await asyncio.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Timeout waiting for result after {max_wait} seconds",
)
def _extract_images_from_result(
self,
result_data: Dict,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: LiteLLMLoggingObj,
) -> ImageResponse:
"""
Extract image URLs from BFL result and populate ImageResponse.
Transform Black Forest Labs response to OpenAI-compatible ImageResponse.
This is called with the FINAL polled response (after handler does polling).
The response contains: {"status": "Ready", "result": {"sample": "https://..."}}
"""
result = result_data.get("result", {})
try:
response_data = raw_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=raw_response.status_code,
message=f"Error parsing BFL response: {e}",
)
result = response_data.get("result", {})
if not model_response.data:
model_response.data = []
@ -378,102 +289,6 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig):
model_response.created = int(time.time())
return model_response
def _parse_initial_response(
self,
raw_response: httpx.Response,
) -> tuple:
"""
Parse initial BFL response and extract polling URL and API key.
Returns:
Tuple of (polling_url, api_key)
"""
try:
response_data = raw_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=raw_response.status_code,
message=f"Error parsing BFL response: {e}",
)
# Check for immediate errors
if "errors" in response_data:
raise BlackForestLabsError(
status_code=raw_response.status_code,
message=f"BFL error: {response_data['errors']}",
)
# Get polling URL
polling_url = response_data.get("polling_url")
if not polling_url:
raise BlackForestLabsError(
status_code=500,
message="No polling_url in BFL response",
)
# Extract API key from original request headers
request_api_key = raw_response.request.headers.get("x-key", "")
return polling_url, request_api_key
def transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ImageResponse:
"""
Transform Black Forest Labs response to OpenAI-compatible ImageResponse.
BFL returns a task ID initially, then we poll until the result is ready.
"""
verbose_logger.debug("BFL starting sync polling...")
polling_url, request_api_key = self._parse_initial_response(raw_response)
# Poll for result (sync)
result_data = self._poll_for_result(polling_url, request_api_key)
verbose_logger.debug("BFL polling complete, extracting images")
return self._extract_images_from_result(result_data, model_response)
async def async_transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ImageResponse:
"""
Async transform Black Forest Labs response to OpenAI-compatible ImageResponse.
BFL returns a task ID initially, then we poll until the result is ready.
"""
verbose_logger.debug("BFL starting async polling...")
polling_url, request_api_key = self._parse_initial_response(raw_response)
# Poll for result (async)
result_data = await self._poll_for_result_async(polling_url, request_api_key)
verbose_logger.debug("BFL async polling complete, extracting images")
return self._extract_images_from_result(result_data, model_response)
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BlackForestLabsError:

230
tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py
View file

@ -1,5 +1,8 @@
"""
Unit tests for Black Forest Labs image edit transformation functionality.
Note: Polling tests are now in test_bfl_image_edit_handler.py
since polling logic was moved to the handler.
"""
import base64
@ -233,219 +236,66 @@ class TestBlackForestLabsImageEditTransformation:
result = self.config._read_image_bytes(images)
assert result == image_data
def test_poll_for_result_success(self):
"""Test successful polling."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/image.png"},
}
mock_client = MagicMock()
mock_client.get.return_value = mock_response
with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client):
result = self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert result["status"] == "Ready"
assert result["result"]["sample"] == "https://example.com/image.png"
def test_poll_for_result_pending_then_ready(self):
"""Test polling that starts pending then becomes ready."""
pending_response = MagicMock()
pending_response.status_code = 200
pending_response.json.return_value = {"status": "Pending"}
ready_response = MagicMock()
ready_response.status_code = 200
ready_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/image.png"},
}
mock_client = MagicMock()
mock_client.get.side_effect = [pending_response, ready_response]
with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client):
result = self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert result["status"] == "Ready"
def test_poll_for_result_error_status(self):
"""Test polling with error status."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"status": "Error"}
mock_client = MagicMock()
mock_client.get.return_value = mock_response
with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert exc_info.value.status_code == 400
assert "Error" in exc_info.value.message
def test_poll_for_result_content_moderated(self):
"""Test polling with content moderated status."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"status": "Content Moderated"}
mock_client = MagicMock()
mock_client.get.return_value = mock_response
with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert exc_info.value.status_code == 400
assert "Content Moderated" in exc_info.value.message
def test_poll_for_result_timeout(self):
"""Test polling timeout."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"status": "Pending"}
mock_client = MagicMock()
mock_client.get.return_value = mock_response
with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=0.2,
interval=0.1,
)
assert exc_info.value.status_code == 408
assert "Timeout" in exc_info.value.message
def test_poll_for_result_http_error(self):
"""Test polling with HTTP error."""
mock_response = MagicMock()
mock_response.status_code = 500
mock_response.text = "Internal Server Error"
mock_client = MagicMock()
mock_client.get.return_value = mock_response
with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert exc_info.value.status_code == 500
def test_transform_image_edit_response_success(self):
"""Test successful response transformation."""
# Create mock initial response with polling URL
mock_request = MagicMock()
mock_request.headers = {"x-key": "test-key"}
mock_response = MagicMock()
"""Test response transformation with final polled response."""
# The response is now the FINAL polled response from handler
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {
"id": "task-123",
"polling_url": "https://api.bfl.ai/v1/get_result?id=task-123",
}
mock_response.request = mock_request
mock_response.status_code = 200
# Mock the polling result
poll_response = MagicMock()
poll_response.status_code = 200
poll_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/edited-image.png"},
"result": {"sample": "https://example.com/edited_image.png"},
}
mock_client = MagicMock()
mock_client.get.return_value = poll_response
with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client):
result = self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert isinstance(result, ImageResponse)
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/edited-image.png"
assert result.created is not None
def test_transform_image_edit_response_no_polling_url(self):
"""Test response transformation when polling URL is missing."""
mock_response = MagicMock()
mock_response.json.return_value = {"id": "task-123"} # No polling_url
mock_response.status_code = 200
with pytest.raises(BlackForestLabsError) as exc_info:
self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
result = self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert exc_info.value.status_code == 500
assert "No polling_url" in exc_info.value.message
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/edited_image.png"
def test_transform_image_edit_response_api_error(self):
"""Test response transformation with API error."""
mock_response = MagicMock()
def test_transform_image_edit_response_no_image_url(self):
"""Test response transformation when no image URL is present."""
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {
"errors": ["Invalid image format"]
"status": "Ready",
"result": {},
}
mock_response.status_code = 400
mock_response.status_code = 200
with pytest.raises(BlackForestLabsError) as exc_info:
with pytest.raises(BlackForestLabsError, match="No image URL"):
self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert "Invalid image format" in exc_info.value.message
def test_transform_image_edit_response_json_parse_error(self):
"""Test response transformation with JSON parse error."""
mock_response = MagicMock()
mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0)
mock_response.status_code = 500
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.side_effect = json.JSONDecodeError("error", "doc", 0)
mock_response.status_code = 200
with pytest.raises(BlackForestLabsError) as exc_info:
with pytest.raises(BlackForestLabsError, match="Error parsing"):
self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert "Error parsing BFL response" in exc_info.value.message
def test_get_error_class(self):
"""Test that get_error_class returns BlackForestLabsError."""
error = self.config.get_error_class(
error_message="Test error",
status_code=400,
headers={},
)
assert isinstance(error, BlackForestLabsError)
assert error.status_code == 400
assert "Test error" in str(error.message)
def test_use_multipart_form_data_returns_false(self):
"""Test that use_multipart_form_data returns False for BFL."""
assert self.config.use_multipart_form_data() is False

522
tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py
View file

@ -1,5 +1,8 @@
"""
Unit tests for Black Forest Labs image generation transformation functionality.
Note: Polling tests are now in test_bfl_image_generation_handler.py
since polling logic was moved to the handler.
"""
import json
@ -49,104 +52,81 @@ class TestBlackForestLabsImageGenerationTransformation:
optional_params = {}
result = self.config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=self.model,
drop_params=False,
non_default_params, optional_params, self.model, drop_params=False
)
# Should be empty since no params provided
# Empty input should return empty output
assert result == {}
def test_map_openai_params_size_mapping(self):
"""Test that OpenAI size param is mapped to BFL width/height."""
"""Test that OpenAI size is mapped to BFL width/height."""
non_default_params = {"size": "1024x1024"}
optional_params = {}
result = self.config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=self.model,
drop_params=False,
non_default_params, optional_params, self.model, drop_params=False
)
assert result.get("width") == 1024
assert result.get("height") == 1024
assert result["width"] == 1024
assert result["height"] == 1024
def test_map_openai_params_size_custom(self):
"""Test custom size parsing."""
non_default_params = {"size": "1920x1080"}
non_default_params = {"size": "800x600"}
optional_params = {}
result = self.config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=self.model,
drop_params=False,
non_default_params, optional_params, self.model, drop_params=False
)
assert result.get("width") == 1920
assert result.get("height") == 1080
assert result["width"] == 800
assert result["height"] == 600
def test_map_openai_params_n_for_ultra(self):
"""Test that n param is mapped to num_images for ultra model."""
"""Test that n is mapped to num_images for ultra model."""
non_default_params = {"n": 4}
optional_params = {}
result = self.config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model="flux-pro-1.1-ultra",
drop_params=False,
non_default_params, optional_params, "flux-pro-1.1-ultra", drop_params=False
)
assert result.get("num_images") == 4
assert result["num_images"] == 4
def test_map_openai_params_quality_hd_for_ultra(self):
"""Test that quality=hd is mapped to raw=True for ultra model."""
"""Test that 'hd' quality maps to raw=True for ultra model."""
non_default_params = {"quality": "hd"}
optional_params = {}
result = self.config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model="flux-pro-1.1-ultra",
drop_params=False,
non_default_params, optional_params, "flux-pro-1.1-ultra", drop_params=False
)
assert result.get("raw") is True
assert result["raw"] is True
def test_map_openai_params_unsupported_raises(self):
"""Test that unsupported param raises error when drop_params=False."""
"""Test that unsupported params raise ValueError when drop_params=False."""
non_default_params = {"unsupported_param": "value"}
optional_params = {}
with pytest.raises(ValueError) as exc_info:
with pytest.raises(ValueError, match="not supported"):
self.config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=self.model,
drop_params=False,
non_default_params, optional_params, self.model, drop_params=False
)
assert "unsupported_param" in str(exc_info.value)
def test_map_openai_params_unsupported_dropped(self):
"""Test that unsupported param is dropped when drop_params=True."""
"""Test that unsupported params are dropped when drop_params=True."""
non_default_params = {"unsupported_param": "value"}
optional_params = {}
result = self.config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=self.model,
drop_params=True,
non_default_params, optional_params, self.model, drop_params=True
)
assert "unsupported_param" not in result
def test_validate_environment_with_api_key(self):
"""Test environment validation with provided API key."""
"""Test that validate_environment sets headers correctly."""
headers = {}
result = self.config.validate_environment(
@ -155,21 +135,21 @@ class TestBlackForestLabsImageGenerationTransformation:
messages=[],
optional_params={},
litellm_params={},
api_key="test-api-key",
api_key="test_api_key",
)
assert result["x-key"] == "test-api-key"
assert result["x-key"] == "test_api_key"
assert result["Content-Type"] == "application/json"
assert result["Accept"] == "application/json"
def test_validate_environment_missing_api_key(self):
"""Test that missing API key raises error."""
"""Test that validate_environment raises error when API key is missing."""
headers = {}
with patch("litellm.llms.black_forest_labs.image_generation.transformation.get_secret_str") as mock_get_secret:
mock_get_secret.return_value = None
with pytest.raises(BlackForestLabsError) as exc_info:
with patch(
"litellm.llms.black_forest_labs.image_generation.transformation.get_secret_str",
return_value=None,
):
with pytest.raises(BlackForestLabsError, match="BFL_API_KEY"):
self.config.validate_environment(
headers=headers,
model=self.model,
@ -179,390 +159,171 @@ class TestBlackForestLabsImageGenerationTransformation:
api_key=None,
)
assert exc_info.value.status_code == 401
assert "BFL_API_KEY is not set" in exc_info.value.message
def test_get_model_endpoint_flux_pro_1_1(self):
"""Test endpoint resolution for flux-pro-1.1."""
"""Test endpoint for flux-pro-1.1 model."""
endpoint = self.config._get_model_endpoint("flux-pro-1.1")
assert endpoint == "/v1/flux-pro-1.1"
def test_get_model_endpoint_flux_pro_1_1_ultra(self):
"""Test endpoint resolution for flux-pro-1.1-ultra."""
"""Test endpoint for flux-pro-1.1-ultra model."""
endpoint = self.config._get_model_endpoint("flux-pro-1.1-ultra")
assert endpoint == "/v1/flux-pro-1.1-ultra"
def test_get_model_endpoint_flux_dev(self):
"""Test endpoint resolution for flux-dev."""
"""Test endpoint for flux-dev model."""
endpoint = self.config._get_model_endpoint("flux-dev")
assert endpoint == "/v1/flux-dev"
def test_get_model_endpoint_flux_pro(self):
"""Test endpoint resolution for flux-pro."""
"""Test endpoint for flux-pro model."""
endpoint = self.config._get_model_endpoint("flux-pro")
assert endpoint == "/v1/flux-pro"
def test_get_model_endpoint_with_provider_prefix(self):
"""Test endpoint resolution with provider prefix."""
endpoint = self.config._get_model_endpoint("black_forest_labs/flux-pro-1.1")
assert endpoint == "/v1/flux-pro-1.1"
def test_get_model_endpoint_unknown_defaults(self):
"""Test that unknown model defaults to flux-pro-1.1."""
"""Test that unknown models default to flux-pro-1.1."""
endpoint = self.config._get_model_endpoint("unknown-model")
assert endpoint == "/v1/flux-pro-1.1"
def test_get_model_endpoint_with_provider_prefix(self):
"""Test that provider prefix is stripped from model name."""
endpoint = self.config._get_model_endpoint("black_forest_labs/flux-pro-1.1")
assert endpoint == "/v1/flux-pro-1.1"
def test_get_complete_url(self):
"""Test complete URL generation."""
"""Test URL construction with default base."""
url = self.config.get_complete_url(
api_base=None,
api_key="test-key",
api_key=None,
model="flux-pro-1.1",
optional_params={},
litellm_params={},
)
assert url == "https://api.bfl.ai/v1/flux-pro-1.1"
assert "https://api.bfl.ai/v1/flux-pro-1.1" == url
def test_get_complete_url_custom_base(self):
"""Test complete URL generation with custom base."""
"""Test URL construction with custom base."""
url = self.config.get_complete_url(
api_base="https://custom.api.com/",
api_key="test-key",
api_base="https://custom.api.com",
api_key=None,
model="flux-pro-1.1",
optional_params={},
litellm_params={},
)
assert url == "https://custom.api.com/v1/flux-pro-1.1"
assert "https://custom.api.com/v1/flux-pro-1.1" == url
def test_transform_image_generation_request(self):
"""Test request transformation to BFL format."""
optional_params = {
"width": 1024,
"height": 1024,
"seed": 42,
}
result = self.config.transform_image_generation_request(
"""Test request body transformation."""
request = self.config.transform_image_generation_request(
model=self.model,
prompt=self.prompt,
optional_params=optional_params,
optional_params={},
litellm_params={},
headers={},
)
assert result["prompt"] == self.prompt
assert result["width"] == 1024
assert result["height"] == 1024
assert result["seed"] == 42
assert result["output_format"] == "png" # Default
assert request["prompt"] == self.prompt
assert request["output_format"] == "png"
def test_transform_image_generation_request_custom_format(self):
"""Test request transformation with custom output format."""
optional_params = {
"output_format": "jpeg",
}
result = self.config.transform_image_generation_request(
"""Test request body with custom output format."""
request = self.config.transform_image_generation_request(
model=self.model,
prompt=self.prompt,
optional_params=optional_params,
optional_params={"output_format": "jpeg"},
litellm_params={},
headers={},
)
assert result["output_format"] == "jpeg"
assert request["output_format"] == "jpeg"
def test_transform_image_generation_request_ultra_params(self):
"""Test request transformation with ultra-specific params."""
optional_params = {
"raw": True,
"num_images": 2,
}
result = self.config.transform_image_generation_request(
"""Test request body with ultra-specific params."""
request = self.config.transform_image_generation_request(
model="flux-pro-1.1-ultra",
prompt=self.prompt,
optional_params=optional_params,
optional_params={
"raw": True,
"num_images": 2,
"aspect_ratio": "16:9",
},
litellm_params={},
headers={},
)
assert result["raw"] is True
assert result["num_images"] == 2
assert request["raw"] is True
assert request["num_images"] == 2
assert request["aspect_ratio"] == "16:9"
def test_poll_for_result_success(self):
"""Test successful polling."""
mock_response = MagicMock()
mock_response.status_code = 200
def test_transform_image_generation_response_success(self):
"""Test response transformation with final polled response."""
# The response is now the FINAL polled response from handler
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/image.png"},
}
mock_client = MagicMock()
mock_client.get.return_value = mock_response
with patch("litellm.llms.black_forest_labs.image_generation.transformation._get_httpx_client", return_value=mock_client):
result = self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert result["status"] == "Ready"
assert result["result"]["sample"] == "https://example.com/image.png"
def test_poll_for_result_pending_then_ready(self):
"""Test polling that starts pending then becomes ready."""
pending_response = MagicMock()
pending_response.status_code = 200
pending_response.json.return_value = {"status": "Pending"}
ready_response = MagicMock()
ready_response.status_code = 200
ready_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/image.png"},
}
mock_client = MagicMock()
mock_client.get.side_effect = [pending_response, ready_response]
with patch("litellm.llms.black_forest_labs.image_generation.transformation._get_httpx_client", return_value=mock_client):
result = self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert result["status"] == "Ready"
def test_poll_for_result_error_status(self):
"""Test polling with error status."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"status": "Error"}
mock_client = MagicMock()
mock_client.get.return_value = mock_response
with patch("litellm.llms.black_forest_labs.image_generation.transformation._get_httpx_client", return_value=mock_client):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert exc_info.value.status_code == 400
assert "Error" in exc_info.value.message
def test_poll_for_result_content_moderated(self):
"""Test polling with content moderated status."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"status": "Content Moderated"}
mock_client = MagicMock()
mock_client.get.return_value = mock_response
with patch("litellm.llms.black_forest_labs.image_generation.transformation._get_httpx_client", return_value=mock_client):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert exc_info.value.status_code == 400
assert "Content Moderated" in exc_info.value.message
def test_poll_for_result_timeout(self):
"""Test polling timeout."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"status": "Pending"}
mock_client = MagicMock()
mock_client.get.return_value = mock_response
with patch("litellm.llms.black_forest_labs.image_generation.transformation._get_httpx_client", return_value=mock_client):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=0.2,
interval=0.1,
)
assert exc_info.value.status_code == 408
assert "Timeout" in exc_info.value.message
def test_poll_for_result_http_error(self):
"""Test polling with HTTP error."""
mock_response = MagicMock()
mock_response.status_code = 500
mock_response.text = "Internal Server Error"
mock_client = MagicMock()
mock_client.get.return_value = mock_response
with patch("litellm.llms.black_forest_labs.image_generation.transformation._get_httpx_client", return_value=mock_client):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert exc_info.value.status_code == 500
def test_extract_images_from_result_single(self):
"""Test extracting single image from result."""
result_data = {
"result": {"sample": "https://example.com/image.png"}
}
model_response = ImageResponse(created=0, data=[])
result = self.config._extract_images_from_result(result_data, model_response)
result = self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=self.logging_obj,
)
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/image.png"
def test_extract_images_from_result_multiple(self):
"""Test extracting multiple images from result."""
result_data = {
def test_transform_image_generation_response_multiple_images(self):
"""Test response transformation with multiple images."""
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {
"status": "Ready",
"result": [
"https://example.com/image1.png",
"https://example.com/image2.png",
]
],
}
mock_response.status_code = 200
model_response = ImageResponse(created=0, data=[])
result = self.config._extract_images_from_result(result_data, model_response)
result = self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=self.logging_obj,
)
assert len(result.data) == 2
assert result.data[0].url == "https://example.com/image1.png"
assert result.data[1].url == "https://example.com/image2.png"
def test_extract_images_from_result_no_image(self):
"""Test error when no image in result."""
result_data = {"result": {}}
model_response = ImageResponse(created=0, data=[])
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._extract_images_from_result(result_data, model_response)
assert exc_info.value.status_code == 500
assert "No image URL" in exc_info.value.message
def test_parse_initial_response_success(self):
"""Test parsing initial response."""
mock_request = MagicMock()
mock_request.headers = {"x-key": "test-key"}
mock_response = MagicMock()
def test_transform_image_generation_response_no_image(self):
"""Test response transformation when no image URL is present."""
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {
"id": "task-123",
"polling_url": "https://api.bfl.ai/v1/get_result?id=task-123",
}
mock_response.request = mock_request
mock_response.status_code = 200
polling_url, api_key = self.config._parse_initial_response(mock_response)
assert polling_url == "https://api.bfl.ai/v1/get_result?id=task-123"
assert api_key == "test-key"
def test_parse_initial_response_no_polling_url(self):
"""Test error when polling URL is missing."""
mock_response = MagicMock()
mock_response.json.return_value = {"id": "task-123"}
mock_response.status_code = 200
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._parse_initial_response(mock_response)
assert exc_info.value.status_code == 500
assert "No polling_url" in exc_info.value.message
def test_parse_initial_response_api_error(self):
"""Test parsing response with API error."""
mock_response = MagicMock()
mock_response.json.return_value = {
"errors": ["Invalid prompt"]
}
mock_response.status_code = 400
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._parse_initial_response(mock_response)
assert "Invalid prompt" in exc_info.value.message
def test_parse_initial_response_json_error(self):
"""Test parsing response with JSON parse error."""
mock_response = MagicMock()
mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0)
mock_response.status_code = 500
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._parse_initial_response(mock_response)
assert "Error parsing BFL response" in exc_info.value.message
def test_transform_image_generation_response_success(self):
"""Test successful response transformation."""
# Create mock initial response with polling URL
mock_request = MagicMock()
mock_request.headers = {"x-key": "test-key"}
mock_response = MagicMock()
mock_response.json.return_value = {
"id": "task-123",
"polling_url": "https://api.bfl.ai/v1/get_result?id=task-123",
}
mock_response.request = mock_request
mock_response.status_code = 200
# Mock the polling result
poll_response = MagicMock()
poll_response.status_code = 200
poll_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/generated-image.png"},
"result": {},
}
mock_client = MagicMock()
mock_client.get.return_value = poll_response
mock_response.status_code = 200
model_response = ImageResponse(created=0, data=[])
with patch("litellm.llms.black_forest_labs.image_generation.transformation._get_httpx_client", return_value=mock_client):
result = self.config.transform_image_generation_response(
with pytest.raises(BlackForestLabsError, match="No image URL"):
self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=self.logging_obj,
request_data={},
optional_params={},
litellm_params={},
encoding=None,
)
assert isinstance(result, ImageResponse)
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/generated-image.png"
assert result.created is not None
def test_get_error_class(self):
"""Test error class generation."""
"""Test that get_error_class returns BlackForestLabsError."""
error = self.config.get_error_class(
error_message="Test error",
status_code=400,
@ -571,87 +332,10 @@ class TestBlackForestLabsImageGenerationTransformation:
assert isinstance(error, BlackForestLabsError)
assert error.status_code == 400
assert error.message == "Test error"
assert "Test error" in str(error.message)
def test_get_black_forest_labs_image_generation_config(self):
"""Test factory function returns correct config."""
"""Test the factory function."""
config = get_black_forest_labs_image_generation_config("flux-pro-1.1")
assert isinstance(config, BlackForestLabsImageGenerationConfig)
@pytest.mark.asyncio
class TestBlackForestLabsImageGenerationTransformationAsync:
"""Async tests for Black Forest Labs image generation."""
def setup_method(self):
"""Set up test fixtures before each test method."""
self.config = BlackForestLabsImageGenerationConfig()
self.model = "flux-pro-1.1"
self.logging_obj = MagicMock()
async def test_poll_for_result_async_success(self):
"""Test successful async polling."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/image.png"},
}
mock_client = AsyncMock()
mock_client.get.return_value = mock_response
with patch("litellm.llms.black_forest_labs.image_generation.transformation.get_async_httpx_client", return_value=mock_client):
result = await self.config._poll_for_result_async(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert result["status"] == "Ready"
assert result["result"]["sample"] == "https://example.com/image.png"
async def test_async_transform_image_generation_response_success(self):
"""Test successful async response transformation."""
# Create mock initial response with polling URL
mock_request = MagicMock()
mock_request.headers = {"x-key": "test-key"}
mock_response = MagicMock()
mock_response.json.return_value = {
"id": "task-123",
"polling_url": "https://api.bfl.ai/v1/get_result?id=task-123",
}
mock_response.request = mock_request
mock_response.status_code = 200
# Mock the polling result
poll_response = MagicMock()
poll_response.status_code = 200
poll_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/generated-image.png"},
}
mock_client = AsyncMock()
mock_client.get.return_value = poll_response
model_response = ImageResponse(created=0, data=[])
with patch("litellm.llms.black_forest_labs.image_generation.transformation.get_async_httpx_client", return_value=mock_client):
result = await self.config.async_transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=self.logging_obj,
request_data={},
optional_params={},
litellm_params={},
encoding=None,
)
assert isinstance(result, ImageResponse)
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/generated-image.png"