From 2835a2265092cb96f82f058fbe7b23f3d6c58eb9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 20:18:11 +0000 Subject: [PATCH] feat(machgen): add MachGen image generation provider --- litellm/__init__.py | 5 + litellm/images/main.py | 18 + litellm/llms/machgen/__init__.py | 8 + litellm/llms/machgen/common_utils.py | 34 ++ .../llms/machgen/image_generation/__init__.py | 8 + .../llms/machgen/image_generation/handler.py | 315 ++++++++++++++++++ .../image_generation/transformation.py | 196 +++++++++++ litellm/types/utils.py | 1 + litellm/utils.py | 6 + provider_endpoints_support.json | 17 + tests/test_litellm/llms/machgen/__init__.py | 0 .../llms/machgen/image_generation/__init__.py | 0 .../test_machgen_image_generation.py | 287 ++++++++++++++++ 13 files changed, 895 insertions(+) create mode 100644 litellm/llms/machgen/__init__.py create mode 100644 litellm/llms/machgen/common_utils.py create mode 100644 litellm/llms/machgen/image_generation/__init__.py create mode 100644 litellm/llms/machgen/image_generation/handler.py create mode 100644 litellm/llms/machgen/image_generation/transformation.py create mode 100644 tests/test_litellm/llms/machgen/__init__.py create mode 100644 tests/test_litellm/llms/machgen/image_generation/__init__.py create mode 100644 tests/test_litellm/llms/machgen/image_generation/test_machgen_image_generation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 3f8c742c5a2..5f77cf03615 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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, diff --git a/litellm/images/main.py b/litellm/images/main.py index 17ea9aa177b..b675c24655e 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -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: diff --git a/litellm/llms/machgen/__init__.py b/litellm/llms/machgen/__init__.py new file mode 100644 index 00000000000..9c18fd188cc --- /dev/null +++ b/litellm/llms/machgen/__init__.py @@ -0,0 +1,8 @@ +from .common_utils import DEFAULT_API_BASE, MachGenError +from .image_generation import MachGenImageGenerationConfig + +__all__ = [ + "DEFAULT_API_BASE", + "MachGenError", + "MachGenImageGenerationConfig", +] diff --git a/litellm/llms/machgen/common_utils.py b/litellm/llms/machgen/common_utils.py new file mode 100644 index 00000000000..0e8a458b9a1 --- /dev/null +++ b/litellm/llms/machgen/common_utils.py @@ -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"}) diff --git a/litellm/llms/machgen/image_generation/__init__.py b/litellm/llms/machgen/image_generation/__init__.py new file mode 100644 index 00000000000..644aa287afa --- /dev/null +++ b/litellm/llms/machgen/image_generation/__init__.py @@ -0,0 +1,8 @@ +from .handler import MachGenImageGeneration, machgen_image_generation +from .transformation import MachGenImageGenerationConfig + +__all__ = [ + "MachGenImageGeneration", + "MachGenImageGenerationConfig", + "machgen_image_generation", +] diff --git a/litellm/llms/machgen/image_generation/handler.py b/litellm/llms/machgen/image_generation/handler.py new file mode 100644 index 00000000000..a751f6f8112 --- /dev/null +++ b/litellm/llms/machgen/image_generation/handler.py @@ -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() diff --git a/litellm/llms/machgen/image_generation/transformation.py b/litellm/llms/machgen/image_generation/transformation.py new file mode 100644 index 00000000000..613a0af1d8c --- /dev/null +++ b/litellm/llms/machgen/image_generation/transformation.py @@ -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) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e4dfac48141..cd20c8e085e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3478,6 +3478,7 @@ class LlmProviders(str, Enum): HELICONE = "helicone" HYPERBOLIC = "hyperbolic" RECRAFT = "recraft" + MACHGEN = "machgen" FAL_AI = "fal_ai" STABILITY = "stability" HEROKU = "heroku" diff --git a/litellm/utils.py b/litellm/utils.py index a11c5500503..04900000615 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 65db63dc045..b41c386e3f9 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -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", diff --git a/tests/test_litellm/llms/machgen/__init__.py b/tests/test_litellm/llms/machgen/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/machgen/image_generation/__init__.py b/tests/test_litellm/llms/machgen/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/machgen/image_generation/test_machgen_image_generation.py b/tests/test_litellm/llms/machgen/image_generation/test_machgen_image_generation.py new file mode 100644 index 00000000000..76de9848db9 --- /dev/null +++ b/tests/test_litellm/llms/machgen/image_generation/test_machgen_image_generation.py @@ -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