mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
feat(machgen): add MachGen image generation provider
This commit is contained in:
Devin AI
parent
998a372417
commit
2835a22650
13 changed files with 895 additions and 0 deletions
|
|
@ -646,6 +646,7 @@ inception_models: Set = set()
|
|||
hyperbolic_models: Set = set()
|
||||
black_forest_labs_models: Set = set()
|
||||
recraft_models: Set = set()
|
||||
machgen_models: Set = set()
|
||||
cometapi_models: Set = set()
|
||||
oci_models: Set = set()
|
||||
vercel_ai_gateway_models: Set = set()
|
||||
|
|
@ -907,6 +908,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
|
|||
black_forest_labs_models.add(key)
|
||||
elif value.get("litellm_provider") == "recraft":
|
||||
recraft_models.add(key)
|
||||
elif value.get("litellm_provider") == "machgen":
|
||||
machgen_models.add(key)
|
||||
elif value.get("litellm_provider") == "cometapi":
|
||||
cometapi_models.add(key)
|
||||
elif value.get("litellm_provider") == "oci":
|
||||
|
|
@ -1048,6 +1051,7 @@ model_list = list(
|
|||
| inception_models
|
||||
| black_forest_labs_models
|
||||
| recraft_models
|
||||
| machgen_models
|
||||
| cometapi_models
|
||||
| oci_models
|
||||
| heroku_models
|
||||
|
|
@ -1155,6 +1159,7 @@ models_by_provider: dict = {
|
|||
"hyperbolic": hyperbolic_models,
|
||||
"black_forest_labs": black_forest_labs_models,
|
||||
"recraft": recraft_models,
|
||||
"machgen": machgen_models,
|
||||
"cometapi": cometapi_models,
|
||||
"oci": oci_models,
|
||||
"volcengine": volcengine_models,
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ from openai.types.audio.transcription_create_params import FileTypes # type: ig
|
|||
# BFL handlers
|
||||
from litellm.llms.black_forest_labs.image_edit.handler import bfl_image_edit
|
||||
from litellm.llms.black_forest_labs.image_generation.handler import bfl_image_generation
|
||||
from litellm.llms.machgen.image_generation.handler import machgen_image_generation
|
||||
from litellm.main import (
|
||||
azure_chat_completions,
|
||||
base_llm_aiohttp_handler,
|
||||
|
|
@ -413,6 +414,23 @@ def image_generation(
|
|||
timeout=timeout,
|
||||
client=client,
|
||||
)
|
||||
elif custom_llm_provider == "machgen":
|
||||
if model is None:
|
||||
raise Exception("Model needs to be set for machgen")
|
||||
return machgen_image_generation.image_generation(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
model_response=model_response,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
logging_obj=litellm_logging_obj,
|
||||
timeout=timeout,
|
||||
api_key=api_key or dynamic_api_key,
|
||||
api_base=api_base,
|
||||
extra_headers=extra_headers,
|
||||
client=client,
|
||||
aimg_generation=aimg_generation,
|
||||
)
|
||||
elif custom_llm_provider == "black_forest_labs":
|
||||
# Route to BFL-specific handler (polling required)
|
||||
if model is None:
|
||||
|
|
|
|||
8
litellm/llms/machgen/__init__.py
Normal file
8
litellm/llms/machgen/__init__.py
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
from .common_utils import DEFAULT_API_BASE, MachGenError
|
||||
from .image_generation import MachGenImageGenerationConfig
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_API_BASE",
|
||||
"MachGenError",
|
||||
"MachGenImageGenerationConfig",
|
||||
]
|
||||
34
litellm/llms/machgen/common_utils.py
Normal file
34
litellm/llms/machgen/common_utils.py
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
"""
|
||||
MachGen common utilities.
|
||||
|
||||
API reference: https://www.machgen.ai/docs/rest_api
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
||||
|
||||
class MachGenError(BaseLLMException):
|
||||
"""Exception raised for MachGen API errors."""
|
||||
|
||||
|
||||
DEFAULT_API_BASE = "https://api.machgen.ai"
|
||||
GENERATE_PATH = "/api/v0/generate"
|
||||
TASKS_PATH = "/api/v0/tasks"
|
||||
|
||||
DEFAULT_POLLING_INTERVAL = 1.5
|
||||
DEFAULT_MAX_POLLING_TIME = 300.0
|
||||
|
||||
STATUS_COMPLETED = "COMPLETED"
|
||||
STATUS_FAILED = "FAILED"
|
||||
|
||||
TEXT_TO_IMAGE_TASK_TYPE = "T2I"
|
||||
|
||||
IMAGE_OUTPUT_KEY = "image"
|
||||
|
||||
MACHGEN_IMAGE_CONFIG_PARAMS: frozenset[str] = frozenset(
|
||||
{"width", "height", "aspect_ratio", "infer_steps", "guidance_scale"}
|
||||
)
|
||||
|
||||
MACHGEN_TOP_LEVEL_PARAMS: frozenset[str] = frozenset({"seed", "enhance_prompt", "moderate"})
|
||||
8
litellm/llms/machgen/image_generation/__init__.py
Normal file
8
litellm/llms/machgen/image_generation/__init__.py
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
from .handler import MachGenImageGeneration, machgen_image_generation
|
||||
from .transformation import MachGenImageGenerationConfig
|
||||
|
||||
__all__ = [
|
||||
"MachGenImageGeneration",
|
||||
"MachGenImageGenerationConfig",
|
||||
"machgen_image_generation",
|
||||
]
|
||||
315
litellm/llms/machgen/image_generation/handler.py
Normal file
315
litellm/llms/machgen/image_generation/handler.py
Normal file
|
|
@ -0,0 +1,315 @@
|
|||
"""
|
||||
MachGen image generation handler.
|
||||
|
||||
Submits a text-to-image task, polls it to completion, and returns an OpenAI shaped
|
||||
`ImageResponse`. With `response_format="b64_json"` the asset is downloaded with the
|
||||
caller's MachGen key and inlined, since MachGen asset URLs require authentication.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import time
|
||||
from collections.abc import Coroutine
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
from urllib.parse import quote
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
_get_httpx_client,
|
||||
get_async_httpx_client,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import ImageObject, ImageResponse
|
||||
|
||||
from ..common_utils import (
|
||||
DEFAULT_MAX_POLLING_TIME,
|
||||
DEFAULT_POLLING_INTERVAL,
|
||||
STATUS_COMPLETED,
|
||||
STATUS_FAILED,
|
||||
TASKS_PATH,
|
||||
MachGenError,
|
||||
)
|
||||
from .transformation import MachGenImageGenerationConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PreparedRequest:
|
||||
url: str
|
||||
headers: dict[str, str]
|
||||
body: dict[str, object]
|
||||
api_base: str
|
||||
|
||||
|
||||
class MachGenImageGeneration:
|
||||
def __init__(
|
||||
self,
|
||||
config: MachGenImageGenerationConfig | None = None,
|
||||
polling_interval: float = DEFAULT_POLLING_INTERVAL,
|
||||
max_polling_time: float = DEFAULT_MAX_POLLING_TIME,
|
||||
) -> None:
|
||||
self.config = config or MachGenImageGenerationConfig()
|
||||
self.polling_interval = polling_interval
|
||||
self.max_polling_time = max_polling_time
|
||||
|
||||
def image_generation(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
model_response: ImageResponse,
|
||||
optional_params: dict,
|
||||
litellm_params: GenericLiteLLMParams | dict,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
extra_headers: dict | None = None,
|
||||
client: HTTPHandler | AsyncHTTPHandler | None = None,
|
||||
aimg_generation: bool = False,
|
||||
) -> ImageResponse | Coroutine[None, None, ImageResponse]:
|
||||
if aimg_generation:
|
||||
return self.async_image_generation(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
model_response=model_response,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
extra_headers=extra_headers,
|
||||
client=client if isinstance(client, AsyncHTTPHandler) else None,
|
||||
)
|
||||
|
||||
sync_client = client if isinstance(client, HTTPHandler) else _get_httpx_client()
|
||||
request = self._prepare_request(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
extra_headers=extra_headers,
|
||||
)
|
||||
|
||||
submit_response = sync_client.post(
|
||||
url=request.url,
|
||||
headers=request.headers,
|
||||
json=request.body,
|
||||
timeout=timeout,
|
||||
)
|
||||
task_url = self._task_url(request, submit_response)
|
||||
|
||||
deadline = time.time() + self.max_polling_time
|
||||
while True:
|
||||
task_response = sync_client.get(url=task_url, headers=request.headers)
|
||||
if self._is_terminal(task_response):
|
||||
break
|
||||
if time.time() >= deadline:
|
||||
raise self._timeout_error()
|
||||
time.sleep(self.polling_interval)
|
||||
|
||||
response = self._transform_response(
|
||||
model=model,
|
||||
task_response=task_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request=request,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
if self._wants_b64(optional_params):
|
||||
asset = sync_client.get(url=self.config.get_asset_url(task_response), headers=request.headers)
|
||||
return self._inline_asset(response, asset)
|
||||
return response
|
||||
|
||||
async def async_image_generation(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
model_response: ImageResponse,
|
||||
optional_params: dict,
|
||||
litellm_params: GenericLiteLLMParams | dict,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
extra_headers: dict | None = None,
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
) -> ImageResponse:
|
||||
async_client = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.MACHGEN)
|
||||
request = self._prepare_request(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
extra_headers=extra_headers,
|
||||
)
|
||||
|
||||
submit_response = await async_client.post(
|
||||
url=request.url,
|
||||
headers=request.headers,
|
||||
json=request.body,
|
||||
timeout=timeout,
|
||||
)
|
||||
task_url = self._task_url(request, submit_response)
|
||||
|
||||
deadline = time.time() + self.max_polling_time
|
||||
while True:
|
||||
task_response = await async_client.get(url=task_url, headers=request.headers)
|
||||
if self._is_terminal(task_response):
|
||||
break
|
||||
if time.time() >= deadline:
|
||||
raise self._timeout_error()
|
||||
await asyncio.sleep(self.polling_interval)
|
||||
|
||||
response = self._transform_response(
|
||||
model=model,
|
||||
task_response=task_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request=request,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
if self._wants_b64(optional_params):
|
||||
asset = await async_client.get(url=self.config.get_asset_url(task_response), headers=request.headers)
|
||||
return self._inline_asset(response, asset)
|
||||
return response
|
||||
|
||||
def _prepare_request(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
optional_params: dict,
|
||||
litellm_params: GenericLiteLLMParams | dict,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
extra_headers: dict | None,
|
||||
) -> PreparedRequest:
|
||||
litellm_params_dict = litellm_params if isinstance(litellm_params, dict) else dict(litellm_params)
|
||||
resolved_api_key = api_key or litellm_params_dict.get("api_key")
|
||||
resolved_api_base = self.config.get_api_base(api_base or litellm_params_dict.get("api_base"))
|
||||
|
||||
headers = self.config.validate_environment(
|
||||
headers=dict(extra_headers or {}),
|
||||
model=model,
|
||||
messages=[],
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
api_key=resolved_api_key,
|
||||
api_base=resolved_api_base,
|
||||
)
|
||||
body = self.config.transform_image_generation_request(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
headers=headers,
|
||||
)
|
||||
url = self.config.get_complete_url(
|
||||
api_base=resolved_api_base,
|
||||
api_key=resolved_api_key,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
)
|
||||
|
||||
logging_obj.pre_call(
|
||||
input=prompt,
|
||||
api_key="",
|
||||
additional_args={"complete_input_dict": body, "api_base": url, "headers": headers},
|
||||
)
|
||||
return PreparedRequest(url=url, headers=headers, body=body, api_base=resolved_api_base)
|
||||
|
||||
def _task_url(self, request: PreparedRequest, submit_response: httpx.Response) -> str:
|
||||
if submit_response.status_code >= 400:
|
||||
raise MachGenError(
|
||||
status_code=submit_response.status_code,
|
||||
message=f"MachGen task submission failed: {submit_response.text}",
|
||||
)
|
||||
|
||||
try:
|
||||
task_id = submit_response.json().get("task_id")
|
||||
except ValueError as e:
|
||||
raise MachGenError(
|
||||
status_code=submit_response.status_code,
|
||||
message=f"Error parsing MachGen submit response: {e}",
|
||||
) from e
|
||||
|
||||
if not task_id:
|
||||
raise MachGenError(status_code=500, message="No task_id in MachGen submit response")
|
||||
|
||||
verbose_logger.debug("MachGen polling task %s", task_id)
|
||||
return f"{request.api_base}{TASKS_PATH}/{quote(str(task_id), safe='')}"
|
||||
|
||||
def _timeout_error(self) -> MachGenError:
|
||||
return MachGenError(
|
||||
status_code=408,
|
||||
message=f"MachGen task polling timed out after {self.max_polling_time} seconds",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_terminal(task_response: httpx.Response) -> bool:
|
||||
if task_response.status_code >= 400:
|
||||
raise MachGenError(
|
||||
status_code=task_response.status_code,
|
||||
message=f"MachGen task polling failed: {task_response.text}",
|
||||
)
|
||||
return task_response.json().get("status") in (STATUS_COMPLETED, STATUS_FAILED)
|
||||
|
||||
def _transform_response(
|
||||
self,
|
||||
model: str,
|
||||
task_response: httpx.Response,
|
||||
model_response: ImageResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
request: PreparedRequest,
|
||||
optional_params: dict,
|
||||
litellm_params: GenericLiteLLMParams | dict,
|
||||
) -> ImageResponse:
|
||||
return self.config.transform_image_generation_response(
|
||||
model=model,
|
||||
raw_response=task_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_data=dict(request.body),
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params if isinstance(litellm_params, dict) else dict(litellm_params),
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _wants_b64(optional_params: dict) -> bool:
|
||||
return optional_params.get("response_format") == "b64_json"
|
||||
|
||||
@staticmethod
|
||||
def _inline_asset(response: ImageResponse, asset: httpx.Response) -> ImageResponse:
|
||||
if asset.status_code >= 400:
|
||||
raise MachGenError(
|
||||
status_code=asset.status_code,
|
||||
message=f"MachGen asset download failed: {asset.text}",
|
||||
)
|
||||
response.data = [ImageObject(b64_json=base64.b64encode(asset.content).decode("utf-8"))]
|
||||
return response
|
||||
|
||||
|
||||
machgen_image_generation = MachGenImageGeneration()
|
||||
196
litellm/llms/machgen/image_generation/transformation.py
Normal file
196
litellm/llms/machgen/image_generation/transformation.py
Normal file
|
|
@ -0,0 +1,196 @@
|
|||
"""
|
||||
MachGen image generation (text-to-image) transformation.
|
||||
|
||||
MachGen is task based: `POST /api/v0/generate` returns a `task_id`, the caller polls
|
||||
`GET /api/v0/tasks/{task_id}` until the task is terminal, then downloads the asset.
|
||||
Polling lives in the handler; this module only translates data.
|
||||
|
||||
API reference: https://www.machgen.ai/docs/rest_api
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
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.utils import ImageObject, ImageResponse
|
||||
|
||||
from ..common_utils import (
|
||||
DEFAULT_API_BASE,
|
||||
GENERATE_PATH,
|
||||
IMAGE_OUTPUT_KEY,
|
||||
MACHGEN_IMAGE_CONFIG_PARAMS,
|
||||
MACHGEN_TOP_LEVEL_PARAMS,
|
||||
STATUS_COMPLETED,
|
||||
STATUS_FAILED,
|
||||
TEXT_TO_IMAGE_TASK_TYPE,
|
||||
MachGenError,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
DEFAULT_HEIGHT = 1024
|
||||
|
||||
|
||||
class MachGenImageGenerationConfig(BaseImageGenerationConfig):
|
||||
"""Translate OpenAI image generation requests to/from the MachGen task API."""
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]:
|
||||
return ["size", "response_format", "seed", "aspect_ratio"]
|
||||
|
||||
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)
|
||||
passthrough_params = MACHGEN_IMAGE_CONFIG_PARAMS | MACHGEN_TOP_LEVEL_PARAMS
|
||||
|
||||
for key, value in non_default_params.items():
|
||||
if key in optional_params or value is None:
|
||||
continue
|
||||
if key == "size":
|
||||
width, height = self._parse_size(value)
|
||||
optional_params["width"] = width
|
||||
optional_params["height"] = height
|
||||
elif key in supported_params or key in passthrough_params:
|
||||
optional_params[key] = value
|
||||
elif not drop_params:
|
||||
raise ValueError(
|
||||
f"Parameter {key} is not supported for model {model}. "
|
||||
f"Supported parameters are {sorted(set(supported_params) | passthrough_params)}. "
|
||||
"Set drop_params=True to drop unsupported parameters."
|
||||
)
|
||||
|
||||
return optional_params
|
||||
|
||||
@staticmethod
|
||||
def _parse_size(size: str) -> tuple[int, int]:
|
||||
try:
|
||||
width, height = (int(part) for part in size.lower().split("x"))
|
||||
except ValueError as e:
|
||||
raise ValueError(f"Invalid size format: '{size}'. Expected 'WIDTHxHEIGHT' (e.g. '1024x1024').") from e
|
||||
return width, height
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
model: str,
|
||||
messages: list[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict:
|
||||
resolved_api_key = api_key or get_secret_str("MACHGEN_API_KEY")
|
||||
if not resolved_api_key:
|
||||
raise MachGenError(
|
||||
status_code=401,
|
||||
message="MACHGEN_API_KEY is not set. Set the environment variable or pass api_key.",
|
||||
)
|
||||
|
||||
return {
|
||||
**headers,
|
||||
"Authorization": f"Bearer {resolved_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def get_api_base(self, api_base: str | None) -> str:
|
||||
return (api_base or get_secret_str("MACHGEN_API_BASE") or DEFAULT_API_BASE).rstrip("/")
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
stream: bool | None = None,
|
||||
) -> str:
|
||||
return f"{self.get_api_base(api_base)}{GENERATE_PATH}"
|
||||
|
||||
def transform_image_generation_request(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
image_config = {
|
||||
key: optional_params[key] for key in MACHGEN_IMAGE_CONFIG_PARAMS if optional_params.get(key) is not None
|
||||
}
|
||||
top_level_params = {
|
||||
key: optional_params[key] for key in MACHGEN_TOP_LEVEL_PARAMS if optional_params.get(key) is not None
|
||||
}
|
||||
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
"task_type": TEXT_TO_IMAGE_TASK_TYPE,
|
||||
"image_config": {"height": DEFAULT_HEIGHT, **image_config},
|
||||
**top_level_params,
|
||||
}
|
||||
|
||||
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: object,
|
||||
api_key: str | None = None,
|
||||
json_mode: bool | None = None,
|
||||
) -> ImageResponse:
|
||||
model_response.data = [ImageObject(url=self.get_asset_url(raw_response))]
|
||||
model_response.created = int(time.time())
|
||||
return model_response
|
||||
|
||||
def get_asset_url(self, raw_response: httpx.Response) -> str:
|
||||
try:
|
||||
task = raw_response.json()
|
||||
except ValueError as e:
|
||||
raise MachGenError(
|
||||
status_code=raw_response.status_code,
|
||||
message=f"Error parsing MachGen task response: {e}",
|
||||
) from e
|
||||
|
||||
status = task.get("status")
|
||||
if status == STATUS_FAILED:
|
||||
raise MachGenError(
|
||||
status_code=400,
|
||||
message=f"MachGen image generation failed: {task.get('error_msg') or task.get('moderation')}",
|
||||
)
|
||||
if status != STATUS_COMPLETED:
|
||||
raise MachGenError(
|
||||
status_code=500,
|
||||
message=f"Unexpected MachGen task status: {status}",
|
||||
)
|
||||
|
||||
asset_url = (task.get("task_output") or {}).get(IMAGE_OUTPUT_KEY)
|
||||
if not asset_url:
|
||||
raise MachGenError(
|
||||
status_code=500,
|
||||
message="MachGen task completed without an image asset",
|
||||
)
|
||||
return asset_url
|
||||
|
||||
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> MachGenError:
|
||||
return MachGenError(status_code=status_code, message=error_message, headers=headers)
|
||||
|
|
@ -3478,6 +3478,7 @@ class LlmProviders(str, Enum):
|
|||
HELICONE = "helicone"
|
||||
HYPERBOLIC = "hyperbolic"
|
||||
RECRAFT = "recraft"
|
||||
MACHGEN = "machgen"
|
||||
FAL_AI = "fal_ai"
|
||||
STABILITY = "stability"
|
||||
HEROKU = "heroku"
|
||||
|
|
|
|||
|
|
@ -8589,6 +8589,12 @@ class ProviderConfigManager:
|
|||
)
|
||||
|
||||
return get_recraft_image_generation_config(model)
|
||||
elif LlmProviders.MACHGEN == provider:
|
||||
from litellm.llms.machgen.image_generation import (
|
||||
MachGenImageGenerationConfig,
|
||||
)
|
||||
|
||||
return MachGenImageGenerationConfig()
|
||||
elif LlmProviders.AIML == provider:
|
||||
from litellm.llms.aiml.image_generation import (
|
||||
get_aiml_image_generation_config,
|
||||
|
|
|
|||
|
|
@ -2829,6 +2829,23 @@
|
|||
"responses": true
|
||||
}
|
||||
},
|
||||
"machgen": {
|
||||
"display_name": "MachGen (`machgen`)",
|
||||
"url": "https://docs.litellm.ai/docs/providers/machgen",
|
||||
"endpoints": {
|
||||
"chat_completions": false,
|
||||
"messages": false,
|
||||
"responses": false,
|
||||
"embeddings": false,
|
||||
"image_generations": true,
|
||||
"image_edits": false,
|
||||
"audio_transcriptions": false,
|
||||
"audio_speech": false,
|
||||
"moderations": false,
|
||||
"batches": false,
|
||||
"rerank": false
|
||||
}
|
||||
},
|
||||
"manus": {
|
||||
"display_name": "Manus (`manus`)",
|
||||
"url": "https://docs.litellm.ai/docs/providers/manus",
|
||||
|
|
|
|||
0
tests/test_litellm/llms/machgen/__init__.py
Normal file
0
tests/test_litellm/llms/machgen/__init__.py
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
"""
|
||||
Unit tests for the MachGen image generation provider (submit -> poll -> asset).
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.machgen.common_utils import MachGenError
|
||||
from litellm.llms.machgen.image_generation.handler import MachGenImageGeneration
|
||||
from litellm.llms.machgen.image_generation.transformation import (
|
||||
MachGenImageGenerationConfig,
|
||||
)
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
MODEL = "FLUX.2-dev"
|
||||
PROMPT = "an isometric reading nook"
|
||||
ASSET_URL = "https://api.machgen.ai/api/v0/assets/t-abc123"
|
||||
|
||||
|
||||
def _response(payload=None, status_code: int = 200, content: bytes = b"") -> httpx.Response:
|
||||
return httpx.Response(
|
||||
status_code=status_code,
|
||||
json=payload,
|
||||
content=None if payload is not None else content,
|
||||
request=httpx.Request("GET", "https://api.machgen.ai"),
|
||||
)
|
||||
|
||||
|
||||
def _submit_response(task_id: str = "t-abc123") -> httpx.Response:
|
||||
return _response({"task_id": task_id})
|
||||
|
||||
|
||||
def _task_response(status: str, **extra) -> httpx.Response:
|
||||
return _response({"task_id": "t-abc123", "status": status, **extra})
|
||||
|
||||
|
||||
def _completed_response() -> httpx.Response:
|
||||
return _task_response("COMPLETED", task_output={"image": ASSET_URL})
|
||||
|
||||
|
||||
def _sync_client(get_side_effect) -> MagicMock:
|
||||
client = MagicMock(spec=HTTPHandler)
|
||||
client.post.return_value = _submit_response()
|
||||
client.get.side_effect = get_side_effect
|
||||
return client
|
||||
|
||||
|
||||
def _generate(client, optional_params=None, api_key: str = "MGA_key:secret") -> ImageResponse:
|
||||
return MachGenImageGeneration(polling_interval=0).image_generation(
|
||||
model=MODEL,
|
||||
prompt=PROMPT,
|
||||
model_response=ImageResponse(),
|
||||
optional_params=optional_params or {},
|
||||
litellm_params={},
|
||||
logging_obj=MagicMock(),
|
||||
timeout=60,
|
||||
api_key=api_key,
|
||||
client=client,
|
||||
)
|
||||
|
||||
|
||||
class TestMachGenTransformation:
|
||||
def setup_method(self):
|
||||
self.config = MachGenImageGenerationConfig()
|
||||
|
||||
def test_request_body_is_a_text_to_image_task(self):
|
||||
body = self.config.transform_image_generation_request(
|
||||
model=MODEL,
|
||||
prompt=PROMPT,
|
||||
optional_params={"width": 1024, "height": 768, "seed": 7},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert body == {
|
||||
"prompt": PROMPT,
|
||||
"model": MODEL,
|
||||
"task_type": "T2I",
|
||||
"image_config": {"width": 1024, "height": 768},
|
||||
"seed": 7,
|
||||
}
|
||||
|
||||
def test_height_defaults_because_machgen_rejects_a_missing_height(self):
|
||||
body = self.config.transform_image_generation_request(
|
||||
model=MODEL,
|
||||
prompt=PROMPT,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert body["image_config"] == {"height": 1024}
|
||||
|
||||
def test_size_maps_to_width_and_height(self):
|
||||
params = self.config.map_openai_params(
|
||||
non_default_params={"size": "1536x1024"},
|
||||
optional_params={},
|
||||
model=MODEL,
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert params == {"width": 1536, "height": 1024}
|
||||
|
||||
def test_invalid_size_is_rejected(self):
|
||||
with pytest.raises(ValueError, match="Invalid size format"):
|
||||
self.config.map_openai_params(
|
||||
non_default_params={"size": "big"},
|
||||
optional_params={},
|
||||
model=MODEL,
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
def test_machgen_specific_params_pass_through(self):
|
||||
params = self.config.map_openai_params(
|
||||
non_default_params={"infer_steps": 30, "guidance_scale": [5.0], "enhance_prompt": True},
|
||||
optional_params={},
|
||||
model=MODEL,
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert params == {"infer_steps": 30, "guidance_scale": [5.0], "enhance_prompt": True}
|
||||
|
||||
body = self.config.transform_image_generation_request(
|
||||
model=MODEL,
|
||||
prompt=PROMPT,
|
||||
optional_params=params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
assert body["image_config"]["infer_steps"] == 30
|
||||
assert body["image_config"]["guidance_scale"] == [5.0]
|
||||
assert body["enhance_prompt"] is True
|
||||
|
||||
def test_unsupported_param_raises_unless_dropped(self):
|
||||
with pytest.raises(ValueError, match="style"):
|
||||
self.config.map_openai_params(
|
||||
non_default_params={"style": "vivid"},
|
||||
optional_params={},
|
||||
model=MODEL,
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert (
|
||||
self.config.map_openai_params(
|
||||
non_default_params={"style": "vivid"},
|
||||
optional_params={},
|
||||
model=MODEL,
|
||||
drop_params=True,
|
||||
)
|
||||
== {}
|
||||
)
|
||||
|
||||
def test_validate_environment_sets_bearer_auth(self):
|
||||
headers = self.config.validate_environment(
|
||||
headers={"x-trace": "1"},
|
||||
model=MODEL,
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key="MGA_key:secret",
|
||||
)
|
||||
|
||||
assert headers["Authorization"] == "Bearer MGA_key:secret"
|
||||
assert headers["x-trace"] == "1"
|
||||
|
||||
def test_validate_environment_requires_a_key(self, monkeypatch):
|
||||
monkeypatch.delenv("MACHGEN_API_KEY", raising=False)
|
||||
with pytest.raises(MachGenError, match="MACHGEN_API_KEY"):
|
||||
self.config.validate_environment(
|
||||
headers={},
|
||||
model=MODEL,
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
|
||||
def test_api_base_is_configurable(self):
|
||||
assert (
|
||||
self.config.get_complete_url(
|
||||
api_base="https://machgen.internal/",
|
||||
api_key="MGA_key:secret",
|
||||
model=MODEL,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
== "https://machgen.internal/api/v0/generate"
|
||||
)
|
||||
|
||||
|
||||
class TestMachGenHandler:
|
||||
def test_polls_until_the_task_completes(self):
|
||||
client = _sync_client([_task_response("PENDING"), _task_response("RUNNING"), _completed_response()])
|
||||
|
||||
response = _generate(client, optional_params={"height": 1024})
|
||||
|
||||
assert [image.url for image in response.data] == [ASSET_URL]
|
||||
assert client.post.call_args.kwargs["url"] == "https://api.machgen.ai/api/v0/generate"
|
||||
assert client.post.call_args.kwargs["headers"]["Authorization"] == "Bearer MGA_key:secret"
|
||||
assert client.get.call_count == 3
|
||||
assert all(
|
||||
call.kwargs["url"] == "https://api.machgen.ai/api/v0/tasks/t-abc123" for call in client.get.call_args_list
|
||||
)
|
||||
|
||||
def test_task_id_is_url_escaped(self):
|
||||
client = _sync_client([_completed_response()])
|
||||
client.post.return_value = _submit_response("../../admin/keys")
|
||||
|
||||
_generate(client)
|
||||
|
||||
assert client.get.call_args.kwargs["url"] == "https://api.machgen.ai/api/v0/tasks/..%2F..%2Fadmin%2Fkeys"
|
||||
|
||||
def test_failed_task_surfaces_the_provider_error(self):
|
||||
client = _sync_client([_task_response("FAILED", error_msg="prompt blocked by moderation")])
|
||||
|
||||
with pytest.raises(MachGenError, match="prompt blocked by moderation"):
|
||||
_generate(client)
|
||||
|
||||
def test_completed_task_without_an_asset_raises(self):
|
||||
client = _sync_client([_task_response("COMPLETED", task_output={})])
|
||||
|
||||
with pytest.raises(MachGenError, match="without an image asset"):
|
||||
_generate(client)
|
||||
|
||||
def test_missing_task_id_raises(self):
|
||||
client = _sync_client([_completed_response()])
|
||||
client.post.return_value = _response({})
|
||||
|
||||
with pytest.raises(MachGenError, match="No task_id"):
|
||||
_generate(client)
|
||||
|
||||
def test_b64_json_downloads_the_authenticated_asset(self):
|
||||
asset = _response(content=b"image-bytes")
|
||||
client = _sync_client([_completed_response(), asset])
|
||||
|
||||
response = _generate(client, optional_params={"response_format": "b64_json"})
|
||||
|
||||
assert response.data[0].b64_json == base64.b64encode(b"image-bytes").decode("utf-8")
|
||||
assert response.data[0].url is None
|
||||
download_call = client.get.call_args_list[-1]
|
||||
assert download_call.kwargs["url"] == ASSET_URL
|
||||
assert download_call.kwargs["headers"]["Authorization"] == "Bearer MGA_key:secret"
|
||||
|
||||
def test_litellm_image_generation_routes_machgen_models(self):
|
||||
import litellm
|
||||
|
||||
client = _sync_client([_completed_response()])
|
||||
|
||||
response = litellm.image_generation(
|
||||
model=f"machgen/{MODEL}",
|
||||
prompt=PROMPT,
|
||||
size="1024x1024",
|
||||
api_key="MGA_key:secret",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert [image.url for image in response.data] == [ASSET_URL]
|
||||
assert client.post.call_args.kwargs["json"]["model"] == MODEL
|
||||
assert client.post.call_args.kwargs["json"]["image_config"] == {"width": 1024, "height": 1024}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_generation_polls_and_returns_the_asset(self):
|
||||
client = MagicMock(spec=AsyncHTTPHandler)
|
||||
client.post = AsyncMock(return_value=_submit_response())
|
||||
client.get = AsyncMock(side_effect=[_task_response("RUNNING"), _completed_response()])
|
||||
|
||||
response = await MachGenImageGeneration(polling_interval=0).async_image_generation(
|
||||
model=MODEL,
|
||||
prompt=PROMPT,
|
||||
model_response=ImageResponse(),
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
logging_obj=MagicMock(),
|
||||
timeout=60,
|
||||
api_key="MGA_key:secret",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert [image.url for image in response.data] == [ASSET_URL]
|
||||
assert client.get.await_count == 2
|
||||
Loading…
Add table
Reference in a new issue