From cd731811d996d51bd6debead779c4a3c1b00aec8 Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 15 Dec 2025 17:50:12 -0300 Subject: [PATCH] feat(black_forest_labs): add native image edit support for Black Forest Labs Add native integration for Black Forest Labs image editing models (flux-kontext-pro, flux-kontext-max, flux-pro-1.0-fill, flux-pro-1.0-expand). Changes: - Add BlackForestLabsImageEditConfig for BFL API transformation - Add BLACK_FOREST_LABS to LlmProviders enum - Add use_multipart_form_data() to BaseImageEditConfig for JSON vs form-data - Modify image_edit_handler to support JSON request bodies - Add comprehensive unit tests Closes #11401 --- litellm/llms/black_forest_labs/__init__.py | 19 + .../llms/black_forest_labs/common_utils.py | 39 ++ .../black_forest_labs/image_edit/__init__.py | 3 + .../image_edit/transformation.py | 334 ++++++++++++++ litellm/types/utils.py | 1 + litellm/utils.py | 6 + .../llms/black_forest_labs/__init__.py | 0 .../black_forest_labs/image_edit/__init__.py | 0 .../test_bfl_image_edit_transformation.py | 430 ++++++++++++++++++ 9 files changed, 832 insertions(+) create mode 100644 litellm/llms/black_forest_labs/__init__.py create mode 100644 litellm/llms/black_forest_labs/common_utils.py create mode 100644 litellm/llms/black_forest_labs/image_edit/__init__.py create mode 100644 litellm/llms/black_forest_labs/image_edit/transformation.py create mode 100644 tests/test_litellm/llms/black_forest_labs/__init__.py create mode 100644 tests/test_litellm/llms/black_forest_labs/image_edit/__init__.py create mode 100644 tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py diff --git a/litellm/llms/black_forest_labs/__init__.py b/litellm/llms/black_forest_labs/__init__.py new file mode 100644 index 00000000000..4ddb0464f7d --- /dev/null +++ b/litellm/llms/black_forest_labs/__init__.py @@ -0,0 +1,19 @@ +from .common_utils import ( + DEFAULT_API_BASE, + DEFAULT_MAX_POLLING_TIME, + DEFAULT_POLLING_INTERVAL, + IMAGE_EDIT_MODELS, + IMAGE_GENERATION_MODELS, + BlackForestLabsError, +) +from .image_edit import BlackForestLabsImageEditConfig + +__all__ = [ + "BlackForestLabsError", + "BlackForestLabsImageEditConfig", + "DEFAULT_API_BASE", + "DEFAULT_MAX_POLLING_TIME", + "DEFAULT_POLLING_INTERVAL", + "IMAGE_EDIT_MODELS", + "IMAGE_GENERATION_MODELS", +] diff --git a/litellm/llms/black_forest_labs/common_utils.py b/litellm/llms/black_forest_labs/common_utils.py new file mode 100644 index 00000000000..4469a9df404 --- /dev/null +++ b/litellm/llms/black_forest_labs/common_utils.py @@ -0,0 +1,39 @@ +""" +Black Forest Labs Common Utilities + +Common utilities, constants, and error handling for Black Forest Labs API. +""" + +from typing import Dict + +from litellm.llms.base_llm.chat.transformation import BaseLLMException + + +class BlackForestLabsError(BaseLLMException): + """Exception class for Black Forest Labs API errors.""" + + pass + + +# API Constants +DEFAULT_API_BASE = "https://api.bfl.ai" + +# Polling configuration +DEFAULT_POLLING_INTERVAL = 1.5 # seconds +DEFAULT_MAX_POLLING_TIME = 300 # 5 minutes + +# Model to endpoint mapping for image edit +IMAGE_EDIT_MODELS: Dict[str, str] = { + "flux-kontext-pro": "/v1/flux-kontext-pro", + "flux-kontext-max": "/v1/flux-kontext-max", + "flux-pro-1.0-fill": "/v1/flux-pro-1.0-fill", + "flux-pro-1.0-expand": "/v1/flux-pro-1.0-expand", +} + +# Model to endpoint mapping for image generation +IMAGE_GENERATION_MODELS: Dict[str, str] = { + "flux-pro-1.1": "/v1/flux-pro-1.1", + "flux-pro-1.1-ultra": "/v1/flux-pro-1.1-ultra", + "flux-dev": "/v1/flux-dev", + "flux-pro": "/v1/flux-pro", +} diff --git a/litellm/llms/black_forest_labs/image_edit/__init__.py b/litellm/llms/black_forest_labs/image_edit/__init__.py new file mode 100644 index 00000000000..6f72edea9f1 --- /dev/null +++ b/litellm/llms/black_forest_labs/image_edit/__init__.py @@ -0,0 +1,3 @@ +from .transformation import BlackForestLabsImageEditConfig + +__all__ = ["BlackForestLabsImageEditConfig"] diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py new file mode 100644 index 00000000000..175e401ab23 --- /dev/null +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -0,0 +1,334 @@ +""" +Black Forest Labs Image Edit Configuration + +Handles transformation between OpenAI-compatible format and Black Forest Labs API format +for image editing endpoints (flux-kontext-pro, flux-kontext-max, etc.). + +API Reference: https://docs.bfl.ai/ +""" + +import base64 +import time +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.secret_managers.main import get_secret_str +from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.router import GenericLiteLLMParams +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, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class BlackForestLabsImageEditConfig(BaseImageEditConfig): + """ + Configuration for Black Forest Labs image editing. + + Supports: + - flux-kontext-pro: General image editing with prompts + - flux-kontext-max: Premium quality editing + - flux-pro-1.0-fill: Inpainting with mask + - flux-pro-1.0-expand: Outpainting (expand image borders) + """ + + def get_supported_openai_params(self, model: str) -> List[str]: + """ + 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) + "size", # Maps to aspect_ratio + "response_format", # b64_json or url + ] + + def map_openai_params( + self, + image_edit_optional_params: ImageEditOptionalRequestParams, + model: str, + drop_params: bool, + ) -> Dict: + """ + Map OpenAI parameters to Black Forest Labs parameters. + + BFL-specific params are passed through directly. + """ + optional_params: Dict[str, Any] = {} + + # Pass through BFL-specific params + bfl_params = [ + "seed", + "output_format", + "safety_tolerance", + "prompt_upsampling", + # Kontext-specific + "aspect_ratio", + # Fill/Inpaint-specific + "steps", + "guidance", + "grow_mask", + # Expand-specific + "top", + "bottom", + "left", + "right", + ] + + # Convert TypedDict to regular dict for access + params_dict = dict(image_edit_optional_params) + + for param in bfl_params: + if param in params_dict: + value = params_dict[param] + if value is not None: + optional_params[param] = value + + # Set default output format + if "output_format" not in optional_params: + optional_params["output_format"] = "png" + + return optional_params + + def validate_environment( + self, + headers: dict, + model: str, + api_key: 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 use_multipart_form_data(self) -> bool: + """ + BFL uses JSON requests, not multipart/form-data. + """ + return False + + 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-kontext-pro") + 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_EDIT_MODELS: + return IMAGE_EDIT_MODELS[model_name] + + # Default to kontext-pro + return IMAGE_EDIT_MODELS["flux-kontext-pro"] + + def get_complete_url( + self, + model: str, + api_base: Optional[str], + litellm_params: dict, + ) -> 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 _read_image_bytes(self, image: Any) -> bytes: + """Read image bytes from various input types.""" + if isinstance(image, bytes): + return image + elif isinstance(image, list): + # If it's a list, take the first image + return self._read_image_bytes(image[0]) + elif hasattr(image, "read"): + # File-like object + pos = getattr(image, "tell", lambda: 0)() + if hasattr(image, "seek"): + image.seek(0) + data = image.read() + if hasattr(image, "seek"): + image.seek(pos) + return data + else: + return image + + def transform_image_edit_request( + self, + model: str, + prompt: str, + image: FileTypes, + image_edit_optional_request_params: Dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> Tuple[Dict, RequestFiles]: + """ + Transform OpenAI-style request to Black Forest Labs request format. + + BFL uses JSON body with base64-encoded images, not multipart/form-data. + """ + # Read and encode image + image_bytes = self._read_image_bytes(image) + b64_image = base64.b64encode(image_bytes).decode("utf-8") + + # Build request body + request_body: Dict[str, Any] = { + "prompt": prompt, + "input_image": b64_image, + } + + # Add optional params + for key, value in image_edit_optional_request_params.items(): + if key not in ["extra_headers", "extra_body"] and value is not None: + request_body[key] = value + + # Handle mask if provided (for inpainting) + if "mask" in image_edit_optional_request_params: + mask = image_edit_optional_request_params["mask"] + mask_bytes = self._read_image_bytes(mask) + request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8") + + # 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() + + while time.time() - start_time < max_wait: + response = httpx.get( + polling_url, + headers={"x-key": api_key}, + timeout=30.0, + ) + + 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, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> ImageResponse: + """ + Transform Black Forest Labs response to OpenAI-compatible ImageResponse. + + BFL returns a task ID initially, then we poll until the result is ready. + """ + 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 + 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") + if not image_url: + raise BlackForestLabsError( + status_code=500, + message="No image URL in BFL result", + ) + + # Build ImageResponse + return ImageResponse( + created=int(time.time()), + data=[ImageObject(url=image_url)], + ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 50e4687b5a8..2d4fd0c8cc9 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3099,6 +3099,7 @@ class LlmProviders(str, Enum): GEMINI = "gemini" AI21 = "ai21" BASETEN = "baseten" + BLACK_FOREST_LABS = "black_forest_labs" AZURE = "azure" AZURE_TEXT = "azure_text" AZURE_AI = "azure_ai" diff --git a/litellm/utils.py b/litellm/utils.py index d192609eead..6a3b18f37e0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8748,6 +8748,12 @@ class ProviderConfigManager: ) return RecraftImageEditConfig() + elif LlmProviders.BLACK_FOREST_LABS == provider: + from litellm.llms.black_forest_labs.image_edit.transformation import ( + BlackForestLabsImageEditConfig, + ) + + return BlackForestLabsImageEditConfig() elif LlmProviders.AZURE_AI == provider: from litellm.llms.azure_ai.image_edit import get_azure_ai_image_edit_config diff --git a/tests/test_litellm/llms/black_forest_labs/__init__.py b/tests/test_litellm/llms/black_forest_labs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/__init__.py b/tests/test_litellm/llms/black_forest_labs/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d 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 new file mode 100644 index 00000000000..2c2d207dc72 --- /dev/null +++ b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py @@ -0,0 +1,430 @@ +""" +Unit tests for Black Forest Labs image edit transformation functionality. +""" + +import base64 +import json +import os +import sys +import time +from io import BytesIO +from typing import Dict, List +from unittest.mock import 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_edit.transformation import ( + BlackForestLabsImageEditConfig, +) +from litellm.llms.black_forest_labs.common_utils import BlackForestLabsError +from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ImageObject, ImageResponse + + +class TestBlackForestLabsImageEditTransformation: + """ + Unit tests for Black Forest Labs image edit transformation functionality. + """ + + def setup_method(self): + """Set up test fixtures before each test method.""" + self.config = BlackForestLabsImageEditConfig() + self.model = "flux-kontext-pro" + self.logging_obj = MagicMock() + self.prompt = "Add a red hat to the person in the image" + + 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 + + def test_map_openai_params_basic(self): + """Test mapping of OpenAI params to BFL params.""" + optional_params = ImageEditOptionalRequestParams() + + result = self.config.map_openai_params( + image_edit_optional_params=optional_params, + model=self.model, + drop_params=False, + ) + + # Should have default output_format + assert result.get("output_format") == "png" + + def test_map_openai_params_with_bfl_specific(self): + """Test that BFL-specific params are passed through.""" + # BFL-specific params are passed as dict keys + optional_params: ImageEditOptionalRequestParams = { + "seed": 42, + "safety_tolerance": 2, + "aspect_ratio": "16:9", + } + + result = self.config.map_openai_params( + image_edit_optional_params=optional_params, + model=self.model, + drop_params=False, + ) + + assert result.get("seed") == 42 + assert result.get("safety_tolerance") == 2 + assert result.get("aspect_ratio") == "16:9" + assert result.get("output_format") == "png" + + 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, + 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_edit.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, + 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_kontext_pro(self): + """Test endpoint resolution for flux-kontext-pro.""" + endpoint = self.config._get_model_endpoint("flux-kontext-pro") + assert endpoint == "/v1/flux-kontext-pro" + + def test_get_model_endpoint_kontext_max(self): + """Test endpoint resolution for flux-kontext-max.""" + endpoint = self.config._get_model_endpoint("flux-kontext-max") + assert endpoint == "/v1/flux-kontext-max" + + 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-kontext-pro") + assert endpoint == "/v1/flux-kontext-pro" + + def test_get_model_endpoint_fill(self): + """Test endpoint resolution for flux-pro-1.0-fill.""" + endpoint = self.config._get_model_endpoint("flux-pro-1.0-fill") + assert endpoint == "/v1/flux-pro-1.0-fill" + + def test_get_complete_url(self): + """Test complete URL generation.""" + url = self.config.get_complete_url( + model="flux-kontext-pro", + api_base=None, + litellm_params={}, + ) + + assert url == "https://api.bfl.ai/v1/flux-kontext-pro" + + def test_get_complete_url_custom_base(self): + """Test complete URL generation with custom base.""" + url = self.config.get_complete_url( + model="flux-kontext-pro", + api_base="https://custom.api.com/", + litellm_params={}, + ) + + assert url == "https://custom.api.com/v1/flux-kontext-pro" + + def test_transform_image_edit_request(self): + """Test request transformation to BFL format.""" + image_data = b"fake_image_data" + image = BytesIO(image_data) + + image_edit_optional_params = { + "seed": 123, + "output_format": "jpeg", + } + + litellm_params = GenericLiteLLMParams() + headers = {} + + data, files = self.config.transform_image_edit_request( + model=self.model, + prompt=self.prompt, + image=image, + image_edit_optional_request_params=image_edit_optional_params, + litellm_params=litellm_params, + headers=headers, + ) + + # Check that data contains the expected parameters + assert data["prompt"] == self.prompt + assert "input_image" in data + # Verify base64 encoding + decoded = base64.b64decode(data["input_image"]) + assert decoded == image_data + assert data["seed"] == 123 + assert data["output_format"] == "jpeg" + + # BFL uses JSON, not multipart - files should be empty + assert files == [] + + def test_transform_image_edit_request_with_mask(self): + """Test request transformation with mask for inpainting.""" + image_data = b"fake_image_data" + mask_data = b"fake_mask_data" + image = BytesIO(image_data) + + image_edit_optional_params = { + "mask": BytesIO(mask_data), + "output_format": "png", + } + + litellm_params = GenericLiteLLMParams() + headers = {} + + data, files = self.config.transform_image_edit_request( + model="flux-pro-1.0-fill", + prompt=self.prompt, + image=image, + image_edit_optional_request_params=image_edit_optional_params, + litellm_params=litellm_params, + headers=headers, + ) + + # Check mask is base64 encoded + assert "mask" in data + decoded_mask = base64.b64decode(data["mask"]) + assert decoded_mask == mask_data + + def test_read_image_bytes_from_bytes(self): + """Test reading image bytes from bytes input.""" + image_data = b"test_image_bytes" + result = self.config._read_image_bytes(image_data) + assert result == image_data + + def test_read_image_bytes_from_file_like(self): + """Test reading image bytes from file-like object.""" + image_data = b"test_image_bytes" + image = BytesIO(image_data) + result = self.config._read_image_bytes(image) + assert result == image_data + + def test_read_image_bytes_from_list(self): + """Test reading image bytes from list (takes first).""" + image_data = b"test_image_bytes" + images = [BytesIO(image_data), BytesIO(b"other")] + 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"}, + } + + with patch("httpx.get", return_value=mock_response): + 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"}, + } + + with patch("httpx.get", side_effect=[pending_response, ready_response]): + 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"} + + with patch("httpx.get", return_value=mock_response): + 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"} + + with patch("httpx.get", return_value=mock_response): + 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"} + + with patch("httpx.get", return_value=mock_response): + 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" + + with patch("httpx.get", return_value=mock_response): + 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() + 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"}, + } + + with patch("httpx.get", return_value=poll_response): + 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, + ) + + assert exc_info.value.status_code == 500 + assert "No polling_url" in exc_info.value.message + + def test_transform_image_edit_response_api_error(self): + """Test response transformation with API error.""" + mock_response = MagicMock() + mock_response.json.return_value = { + "errors": ["Invalid image format"] + } + mock_response.status_code = 400 + + 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, + ) + + 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 + + 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, + ) + + assert "Error parsing BFL response" in exc_info.value.message