From 727fd7684118c1653e97f48b4cc83b7bc6cf4fc3 Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 19 Jan 2026 12:07:33 -0300 Subject: [PATCH] refactor(bfl): separate HTTP logic into dedicated handlers - 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 --- litellm/images/main.py | 38 +- .../black_forest_labs/image_edit/__init__.py | 7 +- .../black_forest_labs/image_edit/handler.py | 427 ++++++++++++++ .../image_edit/transformation.py | 89 +-- .../image_generation/__init__.py | 3 + .../image_generation/handler.py | 424 ++++++++++++++ .../image_generation/transformation.py | 225 +------- .../test_bfl_image_edit_transformation.py | 230 ++------ ...est_bfl_image_generation_transformation.py | 522 ++++-------------- 9 files changed, 1076 insertions(+), 889 deletions(-) create mode 100644 litellm/llms/black_forest_labs/image_edit/handler.py create mode 100644 litellm/llms/black_forest_labs/image_generation/handler.py diff --git a/litellm/images/main.py b/litellm/images/main.py index 7e59a3f9e25..553aa26da98 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -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, diff --git a/litellm/llms/black_forest_labs/image_edit/__init__.py b/litellm/llms/black_forest_labs/image_edit/__init__.py index 6f72edea9f1..73af716e062 100644 --- a/litellm/llms/black_forest_labs/image_edit/__init__.py +++ b/litellm/llms/black_forest_labs/image_edit/__init__.py @@ -1,3 +1,8 @@ +from .handler import BlackForestLabsImageEdit, bfl_image_edit from .transformation import BlackForestLabsImageEditConfig -__all__ = ["BlackForestLabsImageEditConfig"] +__all__ = [ + "BlackForestLabsImageEditConfig", + "BlackForestLabsImageEdit", + "bfl_image_edit", +] diff --git a/litellm/llms/black_forest_labs/image_edit/handler.py b/litellm/llms/black_forest_labs/image_edit/handler.py new file mode 100644 index 00000000000..26a651cb7f5 --- /dev/null +++ b/litellm/llms/black_forest_labs/image_edit/handler.py @@ -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() diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 65d83159f75..23b3b276c3e 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -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, + ) diff --git a/litellm/llms/black_forest_labs/image_generation/__init__.py b/litellm/llms/black_forest_labs/image_generation/__init__.py index 905e59d6ae0..2ccee2069ef 100644 --- a/litellm/llms/black_forest_labs/image_generation/__init__.py +++ b/litellm/llms/black_forest_labs/image_generation/__init__.py @@ -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", ] diff --git a/litellm/llms/black_forest_labs/image_generation/handler.py b/litellm/llms/black_forest_labs/image_generation/handler.py new file mode 100644 index 00000000000..11c69a4b43a --- /dev/null +++ b/litellm/llms/black_forest_labs/image_generation/handler.py @@ -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() diff --git a/litellm/llms/black_forest_labs/image_generation/transformation.py b/litellm/llms/black_forest_labs/image_generation/transformation.py index d9d674cdc45..7e1c2d286d3 100644 --- a/litellm/llms/black_forest_labs/image_generation/transformation.py +++ b/litellm/llms/black_forest_labs/image_generation/transformation.py @@ -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: diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py index 167aef03140..2d5e31e90c1 100644 --- a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py +++ b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py @@ -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 diff --git a/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py b/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py index 107c0551a22..b36130bb622 100644 --- a/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py +++ b/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py @@ -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"