From 2941a555a8784bcfa0e24c3336635d27b28ee28d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 21 Jul 2025 15:01:32 -0700 Subject: [PATCH] [Feat] Add Recraft Image Generation API Support - New LLM Provider (#12832) * add recraft * init RecraftImageGenerationConfig * add get_complete_url + validate_environment * add image_generation_handler in llm http clas * fixes for transform * working recraft request * fixed img gen transform * fixes for llm http handler * test: TestRecraftImageGeneration * fixes for llm_http_handler * fix RecraftImageGenerationConfig * TestRecraftImageGenerationTransformation * add recraft API * docs recraft API * fix code QA * map_openai_params * fix recraft * cost tracking for recraft/recraftv3 * fix code qa check --- docs/my-website/docs/image_generation.md | 19 ++ docs/my-website/docs/providers/recraft.md | 161 +++++++++++ docs/my-website/sidebars.js | 1 + litellm/__init__.py | 5 + litellm/images/main.py | 98 ++++++- .../image_generation/transformation.py | 28 +- litellm/llms/custom_httpx/llm_http_handler.py | 214 ++++++++++++++ .../llms/recraft/image_generation/__init__.py | 13 + .../image_generation/transformation.py | 163 +++++++++++ ...odel_prices_and_context_window_backup.json | 18 ++ litellm/types/llms/recraft.py | 17 ++ litellm/types/utils.py | 1 + litellm/utils.py | 6 + model_prices_and_context_window.json | 18 ++ .../image_gen_tests/test_image_generation.py | 5 + .../test_recraft_image_gen_transformation.py | 270 ++++++++++++++++++ 16 files changed, 1018 insertions(+), 19 deletions(-) create mode 100644 docs/my-website/docs/providers/recraft.md create mode 100644 litellm/llms/recraft/image_generation/__init__.py create mode 100644 litellm/llms/recraft/image_generation/transformation.py create mode 100644 litellm/types/llms/recraft.py create mode 100644 tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py diff --git a/docs/my-website/docs/image_generation.md b/docs/my-website/docs/image_generation.md index 33e4c7b4cc..792a21fc1a 100644 --- a/docs/my-website/docs/image_generation.md +++ b/docs/my-website/docs/image_generation.md @@ -207,7 +207,26 @@ Use this for Stable Diffusion models hosted on Xinference See Xinference usage with LiteLLM [here](./providers/xinference.md#image-generation) +## Recraft Image Generation Models +Use this for AI-powered design and image generation with Recraft + +#### Usage + +```python showLineNumbers +from litellm import image_generation +import os + +os.environ['RECRAFT_API_KEY'] = "your-api-key" + +response = image_generation( + model="recraft/recraftv3", + prompt="A beautiful sunset over a calm ocean", +) +print(response) +``` + +See Recraft usage with LiteLLM [here](./providers/recraft.md#image-generation) ## OpenAI Compatible Image Generation Models Use this for calling `/image_generation` endpoints on OpenAI Compatible Servers, example https://github.com/xorbitsai/inference diff --git a/docs/my-website/docs/providers/recraft.md b/docs/my-website/docs/providers/recraft.md new file mode 100644 index 0000000000..dd4a91ddcf --- /dev/null +++ b/docs/my-website/docs/providers/recraft.md @@ -0,0 +1,161 @@ +# Recraft +https://www.recraft.ai/ + +## Overview + +| Property | Details | +|-------|-------| +| Description | Recraft is an AI-powered design tool that generates high-quality images with precise control over style and content. | +| Provider Route on LiteLLM | `recraft/` | +| Link to Provider Doc | [Recraft ↗](https://www.recraft.ai/docs) | +| Supported Operations | [`/images/generations`](#image-generation) | + +LiteLLM supports Recraft Image Generation calls. + +## API Base, Key +```python +# env variable +os.environ['RECRAFT_API_KEY'] = "your-api-key" +os.environ['RECRAFT_API_BASE'] = "https://external.api.recraft.ai" # [optional] +``` + +## Image Generation + +### Usage - LiteLLM Python SDK + +```python showLineNumbers +from litellm import image_generation +import os + +os.environ['RECRAFT_API_KEY'] = "your-api-key" + +# recraft image generation call +response = image_generation( + model="recraft/recraftv3", + prompt="A beautiful sunset over a calm ocean", +) +print(response) +``` + +### Usage - LiteLLM Proxy Server + +#### 1. Setup config.yaml + +```yaml showLineNumbers +model_list: + - model_name: recraft-v3 + litellm_params: + model: recraft/recraftv3 + api_key: os.environ/RECRAFT_API_KEY + model_info: + mode: image_generation + +general_settings: + master_key: sk-1234 +``` + +#### 2. Start the proxy + +```bash showLineNumbers +litellm --config config.yaml + +# RUNNING on http://0.0.0.0:4000 +``` + +#### 3. Test it + +```bash showLineNumbers +curl --location 'http://0.0.0.0:4000/v1/images/generations' \ +--header 'Content-Type: application/json' \ +--header 'Authorization: Bearer sk-1234' \ +--data '{ + "model": "recraft-v3", + "prompt": "A beautiful sunset over a calm ocean", +}' +``` + +### Advanced Usage - With Additional Parameters + +```python showLineNumbers +from litellm import image_generation +import os + +os.environ['RECRAFT_API_KEY'] = "your-api-key" + +response = image_generation( + model="recraft/recraftv3", + prompt="A beautiful sunset over a calm ocean", +) +print(response) +``` + +### Supported Parameters + +Recraft supports the following OpenAI-compatible parameters: + +| Parameter | Type | Description | Example | +|-----------|------|-------------|---------| +| `n` | integer | Number of images to generate (1-4) | `1` | +| `response_format` | string | Format of response (`url` or `b64_json`) | `"url"` | +| `size` | string | Image dimensions | `"1024x1024"` | +| `style` | string | Image style/artistic direction | `"realistic"` | + +### Using Non-OpenAI Parameters + +If you want to pass parameters that are not supported by OpenAI, you can pass them in your request body, LiteLLM will automatically route it to recraft. + +In this example we will pass `style_id` parameter to the recraft image generation call. + +**Usage with LiteLLM Python SDK** + +```python showLineNumbers +from litellm import image_generation +import os + +os.environ['RECRAFT_API_KEY'] = "your-api-key" + +response = image_generation( + model="recraft/recraftv3", + prompt="A beautiful sunset over a calm ocean", + style_id="your-style-id", +) +``` + +**Usage with LiteLLM Proxy Server + OpenAI Python SDK** + +```python showLineNumbers +from openai import OpenAI +import os + +os.environ['RECRAFT_API_KEY'] = "your-api-key" + +client = OpenAI(api_key=os.environ['RECRAFT_API_KEY']) + +response = client.images.generate( + model="recraft/recraftv3", + prompt="A beautiful sunset over a calm ocean", + extra_body={ + "style_id": "your-style-id", + }, +) +print(response) +``` + +### Supported Image Generation Models + +**Note: All recraft models are supported by LiteLLM** Just pass the model name with `recraft/` and litellm will route it to recraft. + +| Model Name | Function Call | +|------------|---------------| +| recraftv3 | `image_generation(model="recraft/recraftv3", prompt="...")` | +| recraftv2 | `image_generation(model="recraft/recraftv2", prompt="...")` | + +For more details on available models and features, see: https://www.recraft.ai/docs + +## API Key Setup + +Get your API key from [Recraft's website](https://www.recraft.ai/) and set it as an environment variable: + +```bash +export RECRAFT_API_KEY="your-api-key" +``` diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 65b0817682..df7d47678f 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -445,6 +445,7 @@ const sidebars = { "providers/github_copilot", "providers/ai21", "providers/nlp_cloud", + "providers/recraft", "providers/replicate", "providers/togetherai", "providers/v0", diff --git a/litellm/__init__.py b/litellm/__init__.py index 1f08a6c03a..66850ee209 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -505,6 +505,7 @@ moonshot_models: List = [] v0_models: List = [] morph_models: List = [] lambda_ai_models: List = [] +recraft_models: List = [] def is_bedrock_pricing_only_model(key: str) -> bool: """ @@ -689,6 +690,8 @@ def add_known_models(): morph_models.append(key) elif value.get("litellm_provider") == "lambda_ai": lambda_ai_models.append(key) + elif value.get("litellm_provider") == "recraft": + recraft_models.append(key) add_known_models() @@ -776,6 +779,7 @@ model_list = ( + v0_models + morph_models + lambda_ai_models + + recraft_models ) model_list_set = set(model_list) @@ -846,6 +850,7 @@ models_by_provider: dict = { "v0": v0_models, "morph": morph_models, "lambda_ai": lambda_ai_models, + "recraft": recraft_models, } # mapping for those models which have larger equivalents diff --git a/litellm/images/main.py b/litellm/images/main.py index cf9ea7a662..3a675a8168 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -1,7 +1,7 @@ import asyncio import contextvars from functools import partial -from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast +from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast, overload import httpx @@ -14,9 +14,11 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.mock_functions import mock_image_generation from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.custom_llm import CustomLLM #################### Initialize provider clients #################### +llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler() from litellm.main import ( azure_chat_completions, base_llm_aiohttp_handler, @@ -26,6 +28,8 @@ from litellm.main import ( openai_image_variations, vertex_image_generation, ) + +########################################### from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.openai import ImageGenerationRequestQuality @@ -78,17 +82,20 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse: # Await normally init_response = await loop.run_in_executor(None, func_with_context) - if isinstance(init_response, dict) or isinstance( - init_response, ImageResponse - ): ## CACHING SCENARIO - if isinstance(init_response, dict): - init_response = ImageResponse(**init_response) + + response: Optional[ImageResponse] = None + if isinstance(init_response, dict): + response = ImageResponse(**init_response) + elif isinstance(init_response, ImageResponse): ## CACHING SCENARIO response = init_response elif asyncio.iscoroutine(init_response): response = await init_response # type: ignore - else: - # Call the synchronous function using run_in_executor - response = await loop.run_in_executor(None, func_with_context) + + if response is None: + raise ValueError( + "Unable to get Image Response. Please pass a valid llm_provider." + ) + return response except Exception as e: custom_llm_provider = custom_llm_provider or "openai" @@ -101,6 +108,54 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse: ) +# Overload for when aimg_generation=True (returns Coroutine) +@overload +def image_generation( + prompt: str, + model: Optional[str] = None, + n: Optional[int] = None, + quality: Optional[Union[str, ImageGenerationRequestQuality]] = None, + response_format: Optional[str] = None, + size: Optional[str] = None, + style: Optional[str] = None, + user: Optional[str] = None, + input_fidelity: Optional[str] = None, + timeout=600, # default to 10 minutes + api_key: Optional[str] = None, + api_base: Optional[str] = None, + api_version: Optional[str] = None, + custom_llm_provider=None, + *, + aimg_generation: Literal[True], + **kwargs, +) -> Coroutine[Any, Any, ImageResponse]: + ... + + +# Overload for when aimg_generation=False or not specified (returns ImageResponse) +@overload +def image_generation( + prompt: str, + model: Optional[str] = None, + n: Optional[int] = None, + quality: Optional[Union[str, ImageGenerationRequestQuality]] = None, + response_format: Optional[str] = None, + size: Optional[str] = None, + style: Optional[str] = None, + user: Optional[str] = None, + input_fidelity: Optional[str] = None, + timeout=600, # default to 10 minutes + api_key: Optional[str] = None, + api_base: Optional[str] = None, + api_version: Optional[str] = None, + custom_llm_provider=None, + *, + aimg_generation: Literal[False] = False, + **kwargs, +) -> ImageResponse: + ... + + @client def image_generation( # noqa: PLR0915 prompt: str, @@ -118,7 +173,10 @@ def image_generation( # noqa: PLR0915 api_version: Optional[str] = None, custom_llm_provider=None, **kwargs, -) -> ImageResponse: +) -> Union[ + ImageResponse, + Coroutine[Any, Any, ImageResponse], + ]: """ Maps the https://api.openai.com/v1/images/generations endpoint. @@ -348,6 +406,26 @@ def image_generation( # noqa: PLR0915 api_base=api_base, client=client, ) + ######################################################### + # Providers using llm_http_handler + ######################################################### + elif custom_llm_provider in ( + litellm.LlmProviders.RECRAFT, + ): + if image_generation_config is None: + raise ValueError(f"image generation config is not supported for {custom_llm_provider}") + + return llm_http_handler.image_generation_handler( + model=model, + prompt=prompt, + image_generation_provider_config=image_generation_config, + image_generation_optional_request_params=optional_params, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params_dict, + logging_obj=litellm_logging_obj, + timeout=timeout, + client=client, + ) elif ( custom_llm_provider in litellm._custom_providers ): # Assume custom LLM provider diff --git a/litellm/llms/base_llm/image_generation/transformation.py b/litellm/llms/base_llm/image_generation/transformation.py index 134c95b1c8..fc8db8c65c 100644 --- a/litellm/llms/base_llm/image_generation/transformation.py +++ b/litellm/llms/base_llm/image_generation/transformation.py @@ -3,12 +3,12 @@ from typing import TYPE_CHECKING, Any, List, Optional, Union import httpx -from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import ( AllMessageValues, OpenAIImageGenerationOptionalParams, ) -from litellm.types.utils import ModelResponse +from litellm.types.utils import ImageResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -18,12 +18,23 @@ else: LiteLLMLoggingObj = Any -class BaseImageGenerationConfig(BaseConfig, ABC): +class BaseImageGenerationConfig(ABC): @abstractmethod def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: pass + + @abstractmethod + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + pass + def get_complete_url( self, @@ -64,10 +75,10 @@ class BaseImageGenerationConfig(BaseConfig, ABC): headers=headers, ) - def transform_request( + def transform_image_generation_request( self, model: str, - messages: List[AllMessageValues], + prompt: str, optional_params: dict, litellm_params: dict, headers: dict, @@ -76,20 +87,19 @@ class BaseImageGenerationConfig(BaseConfig, ABC): "ImageVariationConfig implementa 'transform_request_image_variation' for image variation models" ) - def transform_response( + def transform_image_generation_response( self, model: str, raw_response: httpx.Response, - model_response: ModelResponse, + model_response: ImageResponse, logging_obj: LiteLLMLoggingObj, request_data: dict, - messages: List[AllMessageValues], optional_params: dict, litellm_params: dict, encoding: Any, api_key: Optional[str] = None, json_mode: Optional[bool] = None, - ) -> ModelResponse: + ) -> ImageResponse: raise NotImplementedError( "ImageVariationConfig implements 'transform_response_image_variation' for image variation models" ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f77f87507e..46fd866be2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -35,6 +35,9 @@ from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, +) from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig @@ -2370,6 +2373,7 @@ class BaseLLMHTTPHandler: BaseRerankConfig, BaseResponsesAPIConfig, BaseImageEditConfig, + BaseImageGenerationConfig, BaseVectorStoreConfig, BaseGoogleGenAIGenerateContentConfig, BaseAnthropicMessagesConfig, @@ -2657,6 +2661,216 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) + def image_generation_handler( + self, + model: str, + prompt: str, + image_generation_provider_config: BaseImageGenerationConfig, + image_generation_optional_request_params: Dict, + custom_llm_provider: str, + litellm_params: Dict, + logging_obj: LiteLLMLoggingObj, + timeout: Union[float, httpx.Timeout], + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + _is_async: bool = False, + fake_stream: bool = False, + litellm_metadata: Optional[Dict[str, Any]] = None, + ) -> Union[ + ImageResponse, + Coroutine[Any, Any, ImageResponse], + ]: + """ + Handles image generation requests. + When _is_async=True, returns a coroutine instead of making the call directly. + """ + if _is_async: + # Return the async coroutine if called with _is_async=True + return self.async_image_generation_handler( + model=model, + prompt=prompt, + image_generation_provider_config=image_generation_provider_config, + image_generation_optional_request_params=image_generation_optional_request_params, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + logging_obj=logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout, + client=client if isinstance(client, AsyncHTTPHandler) else None, + fake_stream=fake_stream, + litellm_metadata=litellm_metadata, + ) + + if client is None or not isinstance(client, HTTPHandler): + sync_httpx_client = _get_httpx_client( + params={"ssl_verify": litellm_params.get("ssl_verify", None)} + ) + else: + sync_httpx_client = client + + headers = image_generation_provider_config.validate_environment( + api_key=litellm_params.get("api_key", None), + headers=image_generation_optional_request_params.get("extra_headers", {}) or {}, + model=model, + messages=[], + optional_params=image_generation_optional_request_params, + litellm_params=dict(litellm_params), + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = image_generation_provider_config.get_complete_url( + model=model, + api_base=litellm_params.get("api_base", None), + api_key=litellm_params.get("api_key", None), + optional_params=image_generation_optional_request_params, + litellm_params=dict(litellm_params), + ) + + data = image_generation_provider_config.transform_image_generation_request( + model=model, + prompt=prompt, + optional_params=image_generation_optional_request_params, + litellm_params=dict(litellm_params), + headers=headers, + ) + + ## LOGGING + logging_obj.pre_call( + input=prompt, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = sync_httpx_client.post( + url=api_base, + headers=headers, + json=data, + timeout=timeout, + ) + + except Exception as e: + raise self._handle_error( + e=e, + provider_config=image_generation_provider_config, + ) + + model_response: ImageResponse = image_generation_provider_config.transform_image_generation_response( + model=model, + raw_response=response, + model_response=litellm.ImageResponse(), + logging_obj=logging_obj, + request_data=data, + optional_params=image_generation_optional_request_params, + litellm_params=dict(litellm_params), + encoding=None, + ) + + return model_response + + async def async_image_generation_handler( + self, + model: str, + prompt: str, + image_generation_provider_config: BaseImageGenerationConfig, + image_generation_optional_request_params: Dict, + custom_llm_provider: str, + litellm_params: Dict, + logging_obj: LiteLLMLoggingObj, + timeout: Union[float, httpx.Timeout], + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + fake_stream: bool = False, + litellm_metadata: Optional[Dict[str, Any]] = None, + ) -> ImageResponse: + """ + Async version of the image generation handler. + Uses async HTTP client to make requests. + """ + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + else: + async_httpx_client = client + + + headers = image_generation_provider_config.validate_environment( + api_key=litellm_params.get("api_key", None), + headers=image_generation_optional_request_params.get("extra_headers", {}) or {}, + model=model, + messages=[], + optional_params=image_generation_optional_request_params, + litellm_params=dict(litellm_params), + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = image_generation_provider_config.get_complete_url( + model=model, + api_base=litellm_params.get("api_base", None), + api_key=litellm_params.get("api_key", None), + optional_params=image_generation_optional_request_params, + litellm_params=dict(litellm_params), + ) + + data = image_generation_provider_config.transform_image_generation_request( + model=model, + prompt=prompt, + optional_params=image_generation_optional_request_params, + litellm_params=dict(litellm_params), + headers=headers, + ) + + ## LOGGING + logging_obj.pre_call( + input=prompt, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = await async_httpx_client.post( + url=api_base, + headers=headers, + json=data, + timeout=timeout, + ) + + except Exception as e: + raise self._handle_error( + e=e, + provider_config=image_generation_provider_config, + ) + + model_response: ImageResponse = image_generation_provider_config.transform_image_generation_response( + model=model, + raw_response=response, + model_response=litellm.ImageResponse(), + logging_obj=logging_obj, + request_data=data, + optional_params=image_generation_optional_request_params, + litellm_params=dict(litellm_params), + encoding=None, + ) + + return model_response + ###### VECTOR STORE HANDLER ###### async def async_vector_store_search_handler( self, diff --git a/litellm/llms/recraft/image_generation/__init__.py b/litellm/llms/recraft/image_generation/__init__.py new file mode 100644 index 0000000000..cb8c5624db --- /dev/null +++ b/litellm/llms/recraft/image_generation/__init__.py @@ -0,0 +1,13 @@ +from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, +) + +from .transformation import RecraftImageGenerationConfig + +__all__ = [ + "RecraftImageGenerationConfig", +] + + +def get_recraft_image_generation_config(model: str) -> BaseImageGenerationConfig: + return RecraftImageGenerationConfig() diff --git a/litellm/llms/recraft/image_generation/transformation.py b/litellm/llms/recraft/image_generation/transformation.py new file mode 100644 index 0000000000..f632b49f3a --- /dev/null +++ b/litellm/llms/recraft/image_generation/transformation.py @@ -0,0 +1,163 @@ +from typing import TYPE_CHECKING, Any, List, Optional + +import httpx + +from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ( + AllMessageValues, + OpenAIImageGenerationOptionalParams, +) +from litellm.types.llms.recraft import RecraftImageGenerationRequestParams +from litellm.types.utils import ImageObject, ImageResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class RecraftImageGenerationConfig(BaseImageGenerationConfig): + DEFAULT_BASE_URL: str = "https://external.api.recraft.ai" + IMAGE_GENERATION_ENDPOINT: str = "v1/images/generations" + + def get_supported_openai_params( + self, model: str + ) -> List[OpenAIImageGenerationOptionalParams]: + """ + https://www.recraft.ai/docs#generate-image + """ + return [ + "n", + "response_format", + "size", + "style" + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + supported_params = self.get_supported_openai_params(model) + for k in non_default_params.keys(): + if k not in optional_params.keys(): + if k in supported_params: + optional_params[k] = non_default_params[k] + elif drop_params: + pass + else: + raise ValueError( + f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." + ) + + return optional_params + + 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 request + + Some providers need `model` in `api_base` + """ + complete_url: str = ( + api_base + or get_secret_str("RECRAFT_API_BASE") + or self.DEFAULT_BASE_URL + ) + + complete_url = complete_url.rstrip("/") + complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" + return complete_url + + 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: + final_api_key: Optional[str] = ( + api_key or + get_secret_str("RECRAFT_API_KEY") + ) + if not final_api_key: + raise ValueError("RECRAFT_API_KEY is not set") + + headers["Authorization"] = f"Bearer {final_api_key}" + return headers + + + + def transform_image_generation_request( + self, + model: str, + prompt: str, + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + """ + Transform the image generation request to the recraft image generation request body + + https://www.recraft.ai/docs#generate-image + """ + recratft_image_generation_request_body: RecraftImageGenerationRequestParams = RecraftImageGenerationRequestParams( + prompt=prompt, + model=model, + **optional_params, + ) + return dict(recratft_image_generation_request_body) + + 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 the image generation response to the litellm image response + + https://www.recraft.ai/docs#generate-image + """ + try: + response_data = raw_response.json() + except Exception as e: + raise self.get_error_class( + error_message=f"Error transforming image generation response: {e}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + if not model_response.data: + model_response.data = [] + + for image_data in response_data["data"]: + model_response.data.append(ImageObject( + url=image_data.get("url", None), + b64_json=image_data.get("b64_json", None), + )) + + return model_response \ No newline at end of file diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 92b48fba76..78fd56effc 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -16890,6 +16890,24 @@ "mode": "chat", "source": "https://platform.moonshot.ai/docs/pricing" }, + "recraft/recraftv3": { + "mode": "image_generation", + "input_cost_per_image": 0.04, + "litellm_provider": "recraft", + "supported_endpoints": [ + "/v1/images/generations" + ], + "source": "https://www.recraft.ai/docs#pricing" + }, + "recraft/recraftv2": { + "mode": "image_generation", + "input_cost_per_image": 0.022, + "litellm_provider": "recraft", + "supported_endpoints": [ + "/v1/images/generations" + ], + "source": "https://www.recraft.ai/docs#pricing" + }, "morph/morph-v3-fast": { "max_tokens": 16000, "max_input_tokens": 16000, diff --git a/litellm/types/llms/recraft.py b/litellm/types/llms/recraft.py new file mode 100644 index 0000000000..176810970b --- /dev/null +++ b/litellm/types/llms/recraft.py @@ -0,0 +1,17 @@ +from typing import Dict, List, Optional + +from typing_extensions import TypedDict + + +class RecraftImageGenerationRequestParams(TypedDict, total=False): + prompt: str + text_layout: Optional[List[Dict]] + n: Optional[int] + style_id: Optional[str] + style: Optional[str] + substyle: Optional[str] + model: Optional[str] + response_format: Optional[str] + size: Optional[str] + negative_prompt: Optional[str] + controls: Optional[Dict] \ No newline at end of file diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 45714a18d0..a025f387a1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2315,6 +2315,7 @@ class LlmProviders(str, Enum): LLAMA = "meta_llama" NSCALE = "nscale" PG_VECTOR = "pg_vector" + RECRAFT = "recraft" # Create a set of all provider values for quick lookup diff --git a/litellm/utils.py b/litellm/utils.py index 1556351e60..70fc77e467 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7151,6 +7151,12 @@ class ProviderConfigManager: ) return get_xinference_image_generation_config(model) + elif LlmProviders.RECRAFT == provider: + from litellm.llms.recraft.image_generation import ( + get_recraft_image_generation_config, + ) + + return get_recraft_image_generation_config(model) return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 92b48fba76..78fd56effc 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -16890,6 +16890,24 @@ "mode": "chat", "source": "https://platform.moonshot.ai/docs/pricing" }, + "recraft/recraftv3": { + "mode": "image_generation", + "input_cost_per_image": 0.04, + "litellm_provider": "recraft", + "supported_endpoints": [ + "/v1/images/generations" + ], + "source": "https://www.recraft.ai/docs#pricing" + }, + "recraft/recraftv2": { + "mode": "image_generation", + "input_cost_per_image": 0.022, + "litellm_provider": "recraft", + "supported_endpoints": [ + "/v1/images/generations" + ], + "source": "https://www.recraft.ai/docs#pricing" + }, "morph/morph-v3-fast": { "max_tokens": 16000, "max_input_tokens": 16000, diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index cc277f7481..79f7f42e55 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -165,6 +165,11 @@ class TestOpenAIGPTImage1(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: return {"model": "gpt-image-1"} +class TestRecraftImageGeneration(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + return {"model": "recraft/recraftv3"} + + class TestAzureOpenAIDalle3(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: litellm.set_verbose = True diff --git a/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py b/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py new file mode 100644 index 0000000000..4dd610ac86 --- /dev/null +++ b/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py @@ -0,0 +1,270 @@ +import json +import os +import sys +from typing import List, Optional +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.recraft.image_generation.transformation import ( + RecraftImageGenerationConfig, +) +from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams +from litellm.types.utils import ImageObject, ImageResponse + + +class TestRecraftImageGenerationTransformation: + def setup_method(self): + """Set up test fixtures before each test method.""" + self.config = RecraftImageGenerationConfig() + self.model = "recraft-v3" + self.logging_obj = MagicMock() + + + def test_map_openai_params_supported_params(self): + """Test that map_openai_params correctly maps supported parameters.""" + non_default_params = { + "n": 2, + "response_format": "url", + "size": "1024x1024", + "style": "photographic" + } + 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 == non_default_params + + def test_map_openai_params_unsupported_param_drop_true(self): + """Test that map_openai_params drops unsupported parameters when drop_params=True.""" + non_default_params = { + "n": 2, + "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 result == {"n": 2} + assert "unsupported_param" not in result + + def test_map_openai_params_unsupported_param_drop_false(self): + """Test that map_openai_params raises ValueError for unsupported parameters when drop_params=False.""" + non_default_params = { + "n": 2, + "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) + assert "is not supported for model" in str(exc_info.value) + + @patch("litellm.llms.recraft.image_generation.transformation.get_secret_str") + def test_get_complete_url_with_api_base(self, mock_get_secret): + """Test that get_complete_url returns correct URL when api_base is provided.""" + api_base = "https://custom.api.recraft.ai" + + result = self.config.get_complete_url( + api_base=api_base, + api_key="test_key", + model=self.model, + optional_params={}, + litellm_params={} + ) + + expected_url = f"{api_base}/{self.config.IMAGE_GENERATION_ENDPOINT}" + assert result == expected_url + mock_get_secret.assert_not_called() + + @patch("litellm.llms.recraft.image_generation.transformation.get_secret_str") + def test_get_complete_url_with_secret_base(self, mock_get_secret): + """Test that get_complete_url uses secret when api_base is None.""" + mock_get_secret.return_value = "https://secret.api.recraft.ai" + + result = self.config.get_complete_url( + api_base=None, + api_key="test_key", + model=self.model, + optional_params={}, + litellm_params={} + ) + + expected_url = f"https://secret.api.recraft.ai/{self.config.IMAGE_GENERATION_ENDPOINT}" + assert result == expected_url + mock_get_secret.assert_called_once_with("RECRAFT_API_BASE") + + @patch("litellm.llms.recraft.image_generation.transformation.get_secret_str") + def test_get_complete_url_with_default_base(self, mock_get_secret): + """Test that get_complete_url uses default base URL when no other options are available.""" + mock_get_secret.return_value = None + + result = self.config.get_complete_url( + api_base=None, + api_key="test_key", + model=self.model, + optional_params={}, + litellm_params={} + ) + + expected_url = f"{self.config.DEFAULT_BASE_URL}/{self.config.IMAGE_GENERATION_ENDPOINT}" + assert result == expected_url + + @patch("litellm.llms.recraft.image_generation.transformation.get_secret_str") + def test_validate_environment_with_api_key(self, mock_get_secret): + """Test that validate_environment correctly sets authorization header when api_key is provided.""" + headers = {} + api_key = "test_api_key" + + result = self.config.validate_environment( + headers=headers, + model=self.model, + messages=[], + optional_params={}, + litellm_params={}, + api_key=api_key + ) + + assert result["Authorization"] == f"Bearer {api_key}" + mock_get_secret.assert_not_called() + + @patch("litellm.llms.recraft.image_generation.transformation.get_secret_str") + def test_validate_environment_with_secret_key(self, mock_get_secret): + """Test that validate_environment uses secret API key when api_key is None.""" + mock_get_secret.return_value = "secret_api_key" + headers = {} + + result = self.config.validate_environment( + headers=headers, + model=self.model, + messages=[], + optional_params={}, + litellm_params={}, + api_key=None + ) + + assert result["Authorization"] == "Bearer secret_api_key" + mock_get_secret.assert_called_once_with("RECRAFT_API_KEY") + + @patch("litellm.llms.recraft.image_generation.transformation.get_secret_str") + def test_validate_environment_no_api_key_raises_error(self, mock_get_secret): + """Test that validate_environment raises ValueError when no API key is available.""" + mock_get_secret.return_value = None + headers = {} + + with pytest.raises(ValueError) as exc_info: + self.config.validate_environment( + headers=headers, + model=self.model, + messages=[], + optional_params={}, + litellm_params={}, + api_key=None + ) + + assert "RECRAFT_API_KEY is not set" in str(exc_info.value) + + def test_transform_image_generation_request(self): + """Test that transform_image_generation_request correctly transforms request parameters.""" + prompt = "A beautiful sunset over mountains" + optional_params = { + "n": 2, + "size": "1024x1024", + "style": "photographic" + } + litellm_params = {} + headers = {} + + result = self.config.transform_image_generation_request( + model=self.model, + prompt=prompt, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers + ) + + assert result["prompt"] == prompt + assert result["model"] == self.model + assert result["n"] == 2 + assert result["size"] == "1024x1024" + assert result["style"] == "photographic" + + def test_transform_image_generation_response_success(self): + """Test that transform_image_generation_response correctly transforms successful response.""" + # Mock response data + response_data = { + "data": [ + {"url": "https://example.com/image1.jpg", "b64_json": None}, + {"url": None, "b64_json": "base64encodeddata"} + ] + } + + # Create mock response + mock_response = MagicMock() + mock_response.json.return_value = response_data + + # Create empty model response + model_response = ImageResponse(data=[]) + + 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 len(result.data) == 2 + assert result.data[0].url == "https://example.com/image1.jpg" + assert result.data[0].b64_json is None + assert result.data[1].url is None + assert result.data[1].b64_json == "base64encodeddata" + + def test_transform_image_generation_response_json_error(self): + """Test that transform_image_generation_response raises error when response JSON is invalid.""" + # Create mock response that raises JSON decode error + mock_response = MagicMock() + mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0) + mock_response.status_code = 500 + mock_response.headers = {} + + model_response = ImageResponse(data=[]) + + with pytest.raises(Exception) as exc_info: + 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 "Error transforming image generation response" in str(exc_info.value) \ No newline at end of file