[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
This commit is contained in:
Ishaan Jaff
GitHub
parent
774af8085e
commit
2941a555a8
@ -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
|
||||
|
||||
161
docs/my-website/docs/providers/recraft.md
Normal file
161
docs/my-website/docs/providers/recraft.md
Normal file
@ -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/<model_name>` 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"
|
||||
```
|
||||
@ -445,6 +445,7 @@ const sidebars = {
|
||||
"providers/github_copilot",
|
||||
"providers/ai21",
|
||||
"providers/nlp_cloud",
|
||||
"providers/recraft",
|
||||
"providers/replicate",
|
||||
"providers/togetherai",
|
||||
"providers/v0",
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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"
|
||||
)
|
||||
|
||||
@ -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,
|
||||
|
||||
13
litellm/llms/recraft/image_generation/__init__.py
Normal file
13
litellm/llms/recraft/image_generation/__init__.py
Normal file
@ -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()
|
||||
163
litellm/llms/recraft/image_generation/transformation.py
Normal file
163
litellm/llms/recraft/image_generation/transformation.py
Normal file
@ -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
|
||||
@ -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,
|
||||
|
||||
17
litellm/types/llms/recraft.py
Normal file
17
litellm/types/llms/recraft.py
Normal file
@ -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]
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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)
|
||||
Loading…
Reference in New Issue
Block a user