diff --git a/litellm/images/main.py b/litellm/images/main.py index eb6aa0c209c..7e59a3f9e25 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -404,7 +404,8 @@ def image_generation( # noqa: PLR0915 litellm.LlmProviders.STABILITY, litellm.LlmProviders.RUNWAYML, litellm.LlmProviders.VERTEX_AI, - litellm.LlmProviders.OPENROUTER + litellm.LlmProviders.OPENROUTER, + litellm.LlmProviders.BLACK_FOREST_LABS, ): if image_generation_config is None: raise ValueError( diff --git a/litellm/llms/black_forest_labs/__init__.py b/litellm/llms/black_forest_labs/__init__.py index 4ddb0464f7d..7a78638c8c7 100644 --- a/litellm/llms/black_forest_labs/__init__.py +++ b/litellm/llms/black_forest_labs/__init__.py @@ -7,10 +7,12 @@ from .common_utils import ( BlackForestLabsError, ) from .image_edit import BlackForestLabsImageEditConfig +from .image_generation import BlackForestLabsImageGenerationConfig __all__ = [ "BlackForestLabsError", "BlackForestLabsImageEditConfig", + "BlackForestLabsImageGenerationConfig", "DEFAULT_API_BASE", "DEFAULT_MAX_POLLING_TIME", "DEFAULT_POLLING_INTERVAL", diff --git a/litellm/llms/black_forest_labs/image_generation/__init__.py b/litellm/llms/black_forest_labs/image_generation/__init__.py new file mode 100644 index 00000000000..905e59d6ae0 --- /dev/null +++ b/litellm/llms/black_forest_labs/image_generation/__init__.py @@ -0,0 +1,9 @@ +from .transformation import ( + BlackForestLabsImageGenerationConfig, + get_black_forest_labs_image_generation_config, +) + +__all__ = [ + "BlackForestLabsImageGenerationConfig", + "get_black_forest_labs_image_generation_config", +] diff --git a/litellm/llms/black_forest_labs/image_generation/transformation.py b/litellm/llms/black_forest_labs/image_generation/transformation.py new file mode 100644 index 00000000000..d9d674cdc45 --- /dev/null +++ b/litellm/llms/black_forest_labs/image_generation/transformation.py @@ -0,0 +1,496 @@ +""" +Black Forest Labs Image Generation Configuration + +Handles transformation between OpenAI-compatible format and Black Forest Labs API format +for image generation endpoints (flux-pro-1.1, flux-pro-1.1-ultra, flux-dev, flux-pro). + +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, + OpenAIImageGenerationOptionalParams, +) +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, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): + """ + Configuration for Black Forest Labs image generation (text-to-image). + + Supports: + - flux-pro-1.1: Fast & reliable standard generation + - flux-pro-1.1-ultra: Ultra high-resolution (up to 4MP) + - flux-dev: Development/open-source variant + - flux-pro: Original pro model + """ + + def get_supported_openai_params( + self, model: str + ) -> List[OpenAIImageGenerationOptionalParams]: + """ + Return list of OpenAI params supported by Black Forest Labs. + + Note: BFL uses different parameter names, these are mapped in map_openai_params. + """ + return [ + "n", # Number of images (BFL returns 1 per request, but ultra supports up to 4) + "size", # Maps to width/height or aspect_ratio + "response_format", # b64_json or url + "quality", # Maps to raw mode for ultra + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + """ + Map OpenAI parameters to Black Forest Labs parameters. + + BFL-specific params are passed through directly. + """ + supported_params = self.get_supported_openai_params(model) + + for k, v in non_default_params.items(): + if k in optional_params: + continue + + if k in supported_params: + # Map OpenAI 'size' to BFL width/height + if k == "size" and v: + self._map_size_param(v, optional_params) + elif k == "n": + # BFL uses num_images for ultra model + if "ultra" in model.lower(): + optional_params["num_images"] = v + elif k == "quality" and v == "hd": + # Map 'hd' quality to raw mode for more natural look + if "ultra" in model.lower(): + optional_params["raw"] = True + else: + optional_params[k] = v + elif not drop_params: + raise ValueError( + f"Parameter {k} is not supported for model {model}. " + f"Supported parameters are {supported_params}. " + f"Set drop_params=True to drop unsupported parameters." + ) + + return optional_params + + def _map_size_param(self, size: str, optional_params: dict) -> None: + """Map OpenAI size parameter to BFL width/height.""" + # Common size mappings + size_mapping = { + "1024x1024": (1024, 1024), + "1792x1024": (1792, 1024), + "1024x1792": (1024, 1792), + "512x512": (512, 512), + "256x256": (256, 256), + } + + if size in size_mapping: + width, height = size_mapping[size] + optional_params["width"] = width + optional_params["height"] = height + elif "x" in size: + # Parse custom size + try: + width, height = map(int, size.lower().split("x")) + optional_params["width"] = width + optional_params["height"] = height + except ValueError: + pass # Ignore invalid size format + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate environment and set up headers for Black Forest Labs. + + BFL uses x-key header for authentication. + """ + final_api_key: Optional[str] = ( + api_key + or get_secret_str("BFL_API_KEY") + or get_secret_str("BLACK_FOREST_LABS_API_KEY") + ) + + if not final_api_key: + raise BlackForestLabsError( + status_code=401, + message="BFL_API_KEY is not set. Please set it via environment variable or pass api_key parameter.", + ) + + headers["x-key"] = final_api_key + headers["Content-Type"] = "application/json" + headers["Accept"] = "application/json" + + return headers + + def _get_model_endpoint(self, model: str) -> str: + """ + Get the API endpoint for a given model. + """ + # Remove provider prefix if present (e.g., "black_forest_labs/flux-pro-1.1") + model_name = model.lower() + if "/" in model_name: + model_name = model_name.split("/")[-1] + + # Check if model is in our mapping + if model_name in IMAGE_GENERATION_MODELS: + return IMAGE_GENERATION_MODELS[model_name] + + # Default to flux-pro-1.1 + return IMAGE_GENERATION_MODELS["flux-pro-1.1"] + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + Get the complete URL for the Black Forest Labs API request. + """ + base_url: str = ( + api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE + ) + base_url = base_url.rstrip("/") + + endpoint = self._get_model_endpoint(model) + return f"{base_url}{endpoint}" + + def transform_image_generation_request( + self, + model: str, + prompt: str, + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + """ + Transform OpenAI-style request to Black Forest Labs request format. + + https://docs.bfl.ai/flux_models/flux_1_1_pro + """ + # Build request body with prompt + request_body: Dict[str, Any] = { + "prompt": prompt, + } + + # BFL-specific params that can be passed through + bfl_params = [ + "width", + "height", + "aspect_ratio", + "seed", + "output_format", + "safety_tolerance", + "prompt_upsampling", + # Ultra-specific + "raw", + "num_images", + "image_url", + "image_prompt_strength", + ] + + for param in bfl_params: + if param in optional_params and optional_params[param] is not None: + request_body[param] = optional_params[param] + + # Set default output format if not specified + if "output_format" not in request_body: + request_body["output_format"] = "png" + + 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", "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_response: ImageResponse, + ) -> ImageResponse: + """ + Extract image URLs from BFL result and populate ImageResponse. + """ + result = result_data.get("result", {}) + + if not model_response.data: + model_response.data = [] + + # Handle single image (sample) or multiple images + if isinstance(result, dict) and "sample" in result: + model_response.data.append(ImageObject(url=result["sample"])) + elif isinstance(result, list): + # Multiple images returned + for img in result: + if isinstance(img, str): + model_response.data.append(ImageObject(url=img)) + elif isinstance(img, dict) and "url" in img: + model_response.data.append(ImageObject(url=img["url"])) + + if not model_response.data: + raise BlackForestLabsError( + status_code=500, + message="No image URL in BFL result", + ) + + 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: + """Return the appropriate error class for Black Forest Labs.""" + return BlackForestLabsError( + status_code=status_code, + message=error_message, + ) + + +def get_black_forest_labs_image_generation_config( + model: str, +) -> BlackForestLabsImageGenerationConfig: + """ + Get the appropriate image generation config for a Black Forest Labs model. + + Currently returns a single config class, but can be extended + for model-specific configurations if needed. + """ + return BlackForestLabsImageGenerationConfig() diff --git a/litellm/utils.py b/litellm/utils.py index 6a3b18f37e0..0b3cdeb15d1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8663,6 +8663,12 @@ class ProviderConfigManager: ) return get_runwayml_image_generation_config(model) + elif LlmProviders.BLACK_FOREST_LABS == provider: + from litellm.llms.black_forest_labs.image_generation import ( + get_black_forest_labs_image_generation_config, + ) + + return get_black_forest_labs_image_generation_config(model) elif LlmProviders.VERTEX_AI == provider: from litellm.llms.vertex_ai.image_generation import ( get_vertex_ai_image_generation_config, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 26a0f926e6b..29525276481 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7822,6 +7822,42 @@ "/v1/images/edits" ] }, + "black_forest_labs/flux-pro-1.1": { + "litellm_provider": "black_forest_labs", + "mode": "image_generation", + "output_cost_per_image": 0.04, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "black_forest_labs/flux-pro-1.1-ultra": { + "litellm_provider": "black_forest_labs", + "mode": "image_generation", + "output_cost_per_image": 0.06, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "black_forest_labs/flux-dev": { + "litellm_provider": "black_forest_labs", + "mode": "image_generation", + "output_cost_per_image": 0.025, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "black_forest_labs/flux-pro": { + "litellm_provider": "black_forest_labs", + "mode": "image_generation", + "output_cost_per_image": 0.05, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "cerebras/llama-3.3-70b": { "input_cost_per_token": 8.5e-07, "litellm_provider": "cerebras", diff --git a/tests/test_litellm/llms/black_forest_labs/image_generation/__init__.py b/tests/test_litellm/llms/black_forest_labs/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d 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 new file mode 100644 index 00000000000..107c0551a22 --- /dev/null +++ b/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py @@ -0,0 +1,657 @@ +""" +Unit tests for Black Forest Labs image generation transformation functionality. +""" + +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.black_forest_labs.image_generation.transformation import ( + BlackForestLabsImageGenerationConfig, + get_black_forest_labs_image_generation_config, +) +from litellm.llms.black_forest_labs.common_utils import BlackForestLabsError +from litellm.types.utils import ImageObject, ImageResponse + + +class TestBlackForestLabsImageGenerationTransformation: + """ + Unit tests for Black Forest Labs image generation transformation functionality. + """ + + 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() + self.prompt = "A beautiful sunset over the ocean" + + def test_get_supported_openai_params(self): + """Test that supported OpenAI params are returned correctly.""" + params = self.config.get_supported_openai_params(self.model) + + assert "n" in params + assert "size" in params + assert "response_format" in params + assert "quality" in params + + def test_map_openai_params_basic(self): + """Test mapping of OpenAI params to BFL params.""" + non_default_params = {} + optional_params = {} + + result = self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=self.model, + drop_params=False, + ) + + # Should be empty since no params provided + assert result == {} + + def test_map_openai_params_size_mapping(self): + """Test that OpenAI size param 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, + ) + + assert result.get("width") == 1024 + assert result.get("height") == 1024 + + def test_map_openai_params_size_custom(self): + """Test custom size parsing.""" + non_default_params = {"size": "1920x1080"} + optional_params = {} + + result = self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=self.model, + drop_params=False, + ) + + assert result.get("width") == 1920 + assert result.get("height") == 1080 + + def test_map_openai_params_n_for_ultra(self): + """Test that n param 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, + ) + + assert result.get("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.""" + 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, + ) + + assert result.get("raw") is True + + def test_map_openai_params_unsupported_raises(self): + """Test that unsupported param raises error when drop_params=False.""" + non_default_params = {"unsupported_param": "value"} + optional_params = {} + + with pytest.raises(ValueError) as exc_info: + self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=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.""" + 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, + ) + + assert "unsupported_param" not in result + + def test_validate_environment_with_api_key(self): + """Test environment validation with provided API key.""" + headers = {} + + result = self.config.validate_environment( + headers=headers, + model=self.model, + messages=[], + optional_params={}, + litellm_params={}, + api_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.""" + 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: + self.config.validate_environment( + headers=headers, + model=self.model, + messages=[], + optional_params={}, + litellm_params={}, + 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.""" + 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.""" + 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.""" + 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.""" + 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.""" + endpoint = self.config._get_model_endpoint("unknown-model") + assert endpoint == "/v1/flux-pro-1.1" + + def test_get_complete_url(self): + """Test complete URL generation.""" + url = self.config.get_complete_url( + api_base=None, + api_key="test-key", + model="flux-pro-1.1", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://api.bfl.ai/v1/flux-pro-1.1" + + def test_get_complete_url_custom_base(self): + """Test complete URL generation with custom base.""" + url = self.config.get_complete_url( + api_base="https://custom.api.com/", + api_key="test-key", + model="flux-pro-1.1", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://custom.api.com/v1/flux-pro-1.1" + + 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( + model=self.model, + prompt=self.prompt, + 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 + + 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( + model=self.model, + prompt=self.prompt, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result["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( + model="flux-pro-1.1-ultra", + prompt=self.prompt, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result["raw"] is True + assert result["num_images"] == 2 + + 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_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) + + 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 = { + "result": [ + "https://example.com/image1.png", + "https://example.com/image2.png", + ] + } + model_response = ImageResponse(created=0, data=[]) + + result = self.config._extract_images_from_result(result_data, model_response) + + 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() + 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"}, + } + + mock_client = MagicMock() + mock_client.get.return_value = poll_response + + 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( + 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.""" + error = self.config.get_error_class( + error_message="Test error", + status_code=400, + headers={}, + ) + + assert isinstance(error, BlackForestLabsError) + assert error.status_code == 400 + assert error.message == "Test error" + + def test_get_black_forest_labs_image_generation_config(self): + """Test factory function returns correct config.""" + 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"