scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-28 01:32:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 43

feat(xai): add grok imagine image generation, edit, and video download

Browse source 
Map OpenAI /images/generations, /images/edits, and /videos to xAI Imagine,
including CDN video download and existing SuperGrok OAuth or API-key auth

Co-authored-by: HX <HXhlx@users.noreply.github.com>
This commit is contained in:
hx 2026-09-08 10:55:50 +00:00
parent 1af7a403c6
commit f4d012c9a3
23 changed files with 1826 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/images/main.py
View file

@ -388,6 +388,7 @@ def image_generation(
litellm.LlmProviders.DASHSCOPE,
litellm.LlmProviders.QWENCLOUD,
litellm.LlmProviders.QWEN_AI_PLATFORM,
litellm.LlmProviders.XAI,
):
if image_generation_config is None:
raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
@ -397,7 +398,7 @@ def image_generation(
litellm_params_dict["api_base"] = _api_base
return llm_http_handler.image_generation_handler(
api_key=api_key,
api_key=api_key or dynamic_api_key,
model=model,
prompt=prompt,
image_generation_provider_config=image_generation_config,

3
litellm/llms/xai/image_edit/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .transformation import XAIImageEditConfig
__all__ = ["XAIImageEditConfig"]

237
litellm/llms/xai/image_edit/transformation.py Normal file
View file

@ -0,0 +1,237 @@
import base64
from io import BufferedReader, BytesIO
from typing import Any, Final
import httpx
from httpx._types import RequestFiles
from litellm.constants import XAI_API_BASE
from litellm.exceptions import AuthenticationError
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.xai.common_utils import XAIModelInfo
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import FileTypes, ImageObject, ImageResponse
_SIZE_TO_ASPECT_RATIO: Final = {
"1024x1024": "1:1",
"1792x1024": "16:9",
"1024x1792": "9:16",
"1536x1024": "3:2",
"1024x1536": "2:3",
"1280x720": "16:9",
"720x1280": "9:16",
"1920x1080": "16:9",
"1080x1920": "9:16",
}
_XAI_NATIVE_PARAMS: Final = frozenset({"aspect_ratio", "n", "resolution"})
def _read_seekable(image: BytesIO | BufferedReader) -> bytes:
current_pos: Final = image.tell()
image.seek(0)
data: Final = image.read()
image.seek(current_pos)
return data
class XAIImageEditConfig(BaseImageEditConfig):
def get_supported_openai_params(self, model: str) -> list:
return ["n", "response_format", "size", "user"]
def map_openai_params(
self,
image_edit_optional_params: ImageEditOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict:
supported: Final = frozenset(self.get_supported_openai_params(model))
allowed: Final = supported | _XAI_NATIVE_PARAMS
incoming: Final = dict(image_edit_optional_params)
unknown: Final = tuple(key for key in incoming if key not in allowed)
if unknown and not drop_params:
raise ValueError(
f"Parameter {unknown[0]} is not supported for model {model}. "
f"Supported parameters are {sorted(allowed)}. "
"Set drop_params=True to drop unsupported parameters."
)
mapped: Final = {key: value for key, value in incoming.items() if key in allowed}
size: Final = mapped.get("size")
aspect_ratio: Final = mapped.get("aspect_ratio") or (
_SIZE_TO_ASPECT_RATIO.get(str(size), "1:1") if size else None
)
n: Final = mapped.get("n")
resolution: Final = mapped.get("resolution")
return {
**({"aspect_ratio": aspect_ratio} if aspect_ratio is not None else {}),
**({"n": int(n)} if n is not None else {}),
**({"resolution": resolution} if resolution is not None else {}),
}
def use_multipart_form_data(self) -> bool:
return False
def get_complete_url(
self,
model: str,
api_base: str | None,
litellm_params: dict,
) -> str:
from litellm.llms.xai.oauth import XAIOAuthAuthenticator, should_use_xai_oauth
api_key: Final = litellm_params.get("api_key") if isinstance(litellm_params, dict) else None
resolved_base: Final = (
XAIOAuthAuthenticator().get_api_base()
if should_use_xai_oauth(litellm_params) and not XAIModelInfo.get_api_key(api_key)
else (
api_base
or get_secret_str("XAI_API_BASE")
or get_secret_str("XAI_OAUTH_API_BASE")
or XAI_API_BASE
)
)
base: Final = (resolved_base or XAI_API_BASE).rstrip("/")
if base.endswith("/v1"):
return f"{base}/images/edits"
return f"{base}/v1/images/edits"
def validate_environment(
self,
headers: dict,
model: str,
api_key: str | None = None,
litellm_params: dict | None = None,
api_base: str | None = None,
) -> dict:
from litellm.llms.xai.oauth import (
XAIOAuthAuthenticator,
XAIOAuthError,
should_use_xai_oauth,
)
params: Final = litellm_params or {}
dynamic_api_key: Final = XAIModelInfo.get_api_key(api_key)
if should_use_xai_oauth(params) and not dynamic_api_key:
try:
headers["Authorization"] = f"Bearer {XAIOAuthAuthenticator().get_access_token()}"
except XAIOAuthError as exc:
raise AuthenticationError(
model=model,
llm_provider="xai",
message=str(exc),
) from exc
else:
if not dynamic_api_key:
raise AuthenticationError(
model=model,
llm_provider="xai",
message=(
"Missing xAI credentials for image edit. "
"Pass api_key / XAI_API_KEY, or set use_xai_oauth=True."
),
)
headers["Authorization"] = f"Bearer {dynamic_api_key}"
if "content-type" not in headers and "Content-Type" not in headers:
headers["Content-Type"] = "application/json"
return headers
def transform_image_edit_request(
self,
model: str,
prompt: str | None,
image: FileTypes | None,
image_edit_optional_request_params: dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> tuple[dict, RequestFiles]:
if image is None:
raise ValueError("xAI image edit requires at least one reference image.")
image_payloads: Final = tuple(self._to_image_url(item) for item in self._as_image_list(image))
if not image_payloads:
raise ValueError("xAI image edit requires at least one reference image.")
n: Final = image_edit_optional_request_params.get("n")
request: Final[dict[str, Any]] = {
"model": XAIModelInfo.get_base_model(model) or model,
**({"prompt": prompt} if prompt is not None else {}),
**(
{"image": image_payloads[0]}
if len(image_payloads) == 1
else {"images": list(image_payloads)}
),
**{
key: image_edit_optional_request_params[key]
for key in ("aspect_ratio", "resolution")
if image_edit_optional_request_params.get(key) is not None
},
**({"n": int(n)} if n is not None else {}),
}
return request, []
def transform_image_edit_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: Any,
) -> ImageResponse:
try:
response_data: Final = raw_response.json()
except Exception:
raise self.get_error_class(
error_message=raw_response.text,
status_code=raw_response.status_code,
headers=raw_response.headers,
)
images: Final = tuple(
ImageObject(
url=item.get("url"),
b64_json=item.get("b64_json") or item.get("b64"),
)
for item in response_data.get("data") or ()
if isinstance(item, dict)
)
if not images:
raise self.get_error_class(
error_message=f"xAI image edit returned no image data: {response_data}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
return ImageResponse(data=list(images))
def _as_image_list(self, image: FileTypes | list[FileTypes]) -> tuple[FileTypes, ...]:
if isinstance(image, list):
return tuple(item for item in image if item is not None)
return (image,)
def _to_image_url(self, image: FileTypes) -> dict[str, str]:
if isinstance(image, str):
return {"url": image}
if isinstance(image, dict):
if image.get("url"):
return {"url": str(image["url"])}
if image.get("file_id"):
return {"file_id": str(image["file_id"])}
mime: Final = ImageEditRequestUtils.get_image_content_type(image)
encoded: Final = base64.b64encode(self._read_all_bytes(image)).decode("utf-8")
return {"url": f"data:{mime};base64,{encoded}"}
def _read_all_bytes(self, image: FileTypes) -> bytes:
if isinstance(image, bytes):
return image
if isinstance(image, bytearray):
return bytes(image)
if isinstance(image, (BytesIO, BufferedReader)):
return _read_seekable(image)
if hasattr(image, "read"):
raw: Final = image.read()
if isinstance(raw, str):
return raw.encode("utf-8")
return bytes(raw)
raise ValueError(f"Unsupported image input type for xAI image edit: {type(image)}")

11
litellm/llms/xai/image_generation/__init__.py Normal file
View file

@ -0,0 +1,11 @@
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from .transformation import XAIImageGenerationConfig
__all__ = ["XAIImageGenerationConfig", "get_xai_image_generation_config"]
def get_xai_image_generation_config(model: str) -> BaseImageGenerationConfig:
return XAIImageGenerationConfig()

207
litellm/llms/xai/image_generation/transformation.py Normal file
View file

@ -0,0 +1,207 @@
from typing import TYPE_CHECKING, Any, Final
import httpx
from litellm.constants import XAI_API_BASE
from litellm.exceptions import AuthenticationError
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.llms.xai.common_utils import XAIModelInfo
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIImageGenerationOptionalParams,
)
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
_SIZE_TO_ASPECT_RATIO: Final = {
"1024x1024": "1:1",
"1792x1024": "16:9",
"1024x1792": "9:16",
"1536x1024": "3:2",
"1024x1536": "2:3",
"1280x720": "16:9",
"720x1280": "9:16",
"1920x1080": "16:9",
"1080x1920": "9:16",
}
_XAI_NATIVE_PARAMS: Final = frozenset({"aspect_ratio", "n"})
class XAIImageGenerationConfig(BaseImageGenerationConfig):
def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]:
return ["n", "response_format", "size", "user"]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
supported_params: Final = frozenset(self.get_supported_openai_params(model))
allowed: Final = supported_params | _XAI_NATIVE_PARAMS
unknown: Final = tuple(
key
for key in non_default_params
if key not in optional_params and key not in allowed
)
if unknown and not drop_params:
raise ValueError(
f"Parameter {unknown[0]} is not supported for model {model}. "
f"Supported parameters are {sorted(allowed)}. "
"Set drop_params=True to drop unsupported parameters."
)
merged: Final = {**optional_params, **{k: v for k, v in non_default_params.items() if k in allowed}}
size: Final = merged.get("size")
aspect_ratio: Final = merged.get("aspect_ratio") or (
_SIZE_TO_ASPECT_RATIO.get(str(size), "1:1") if size else None
)
n: Final = merged.get("n")
return {
**({"aspect_ratio": aspect_ratio} if aspect_ratio is not None else {}),
**({"n": int(n)} if n is not None else {}),
}
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:
from litellm.llms.xai.oauth import XAIOAuthAuthenticator, should_use_xai_oauth
resolved_base: Final = (
XAIOAuthAuthenticator().get_api_base()
if should_use_xai_oauth(litellm_params) and not XAIModelInfo.get_api_key(api_key)
else (
api_base
or get_secret_str("XAI_API_BASE")
or get_secret_str("XAI_OAUTH_API_BASE")
or XAI_API_BASE
)
)
base: Final = (resolved_base or XAI_API_BASE).rstrip("/")
if base.endswith("/v1"):
return f"{base}/images/generations"
return f"{base}/v1/images/generations"
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:
from litellm.llms.xai.oauth import (
XAIOAuthAuthenticator,
XAIOAuthError,
should_use_xai_oauth,
)
dynamic_api_key: Final = XAIModelInfo.get_api_key(api_key)
if should_use_xai_oauth(litellm_params) and not dynamic_api_key:
try:
headers["Authorization"] = f"Bearer {XAIOAuthAuthenticator().get_access_token()}"
except XAIOAuthError as exc:
raise AuthenticationError(
model=model,
llm_provider="xai",
message=str(exc),
) from exc
else:
if not dynamic_api_key:
raise AuthenticationError(
model=model,
llm_provider="xai",
message=(
"Missing xAI credentials for image generation. "
"Pass api_key / XAI_API_KEY, or set use_xai_oauth=True."
),
)
headers["Authorization"] = f"Bearer {dynamic_api_key}"
if "content-type" not in headers and "Content-Type" not in headers:
headers["Content-Type"] = "application/json"
return headers
def transform_image_generation_request(
self,
model: str,
prompt: str,
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
n: Final = optional_params.get("n")
return {
"model": XAIModelInfo.get_base_model(model) or model,
"prompt": prompt,
**(
{"aspect_ratio": optional_params["aspect_ratio"]}
if optional_params.get("aspect_ratio") is not None
else {}
),
**({"n": int(n)} if n is not None else {}),
}
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: str | None = None,
json_mode: bool | None = None,
) -> ImageResponse:
try:
response_data: Final = raw_response.json()
except Exception:
raise self.get_error_class(
error_message=raw_response.text,
status_code=raw_response.status_code,
headers=raw_response.headers,
)
logging_obj.post_call(
input=request_data.get("prompt", ""),
api_key=api_key,
additional_args={"complete_input_dict": request_data},
original_response=response_data,
)
images: Final = tuple(
ImageObject(
url=item.get("url"),
b64_json=item.get("b64_json") or item.get("b64"),
)
for item in response_data.get("data") or ()
if isinstance(item, dict)
)
if not images:
raise self.get_error_class(
error_message=f"xAI image generation returned no image data: {response_data}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
model_response.data = list(images)
return model_response

3
litellm/llms/xai/videos/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .transformation import XAIVideoConfig
__all__ = ["XAIVideoConfig"]

409
litellm/llms/xai/videos/transformation.py Normal file
View file

@ -0,0 +1,409 @@
import time
from typing import TYPE_CHECKING, Any, Final
import httpx
from httpx._types import RequestFiles
import litellm
from litellm.constants import XAI_API_BASE
from litellm.exceptions import AuthenticationError
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
_get_httpx_client,
get_async_httpx_client,
)
from litellm.llms.xai.common_utils import XAIModelInfo
from litellm.secret_managers.main import get_secret_str
from litellm.types.router import GenericLiteLLMParams
from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject
from litellm.types.videos.utils import (
encode_video_id_with_provider,
extract_original_video_id,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
_SIZE_TO_ASPECT_RATIO: Final = {
"1024x1024": "1:1",
"1792x1024": "16:9",
"1024x1792": "9:16",
"1280x720": "16:9",
"720x1280": "9:16",
"1920x1080": "16:9",
"1080x1920": "9:16",
}
def _duration_from_seconds(seconds: object) -> int:
try:
return int(seconds) if seconds is not None else 6
except (TypeError, ValueError):
return 6
_STATUS_MAP: Final = {
"done": "completed",
"completed": "completed",
"succeeded": "completed",
"failed": "failed",
"expired": "failed",
"pending": "processing",
"processing": "processing",
"in_progress": "processing",
}
class XAIVideoConfig(BaseVideoConfig):
def get_supported_openai_params(self, model: str) -> list:
return [
"model",
"prompt",
"input_reference",
"seconds",
"size",
"user",
"extra_headers",
]
def map_openai_params(
self,
video_create_optional_params: VideoCreateOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict:
incoming: Final = dict(video_create_optional_params)
size: Final = incoming.get("size")
return {
**{
key: value
for key, value in incoming.items()
if key not in {"seconds", "size", "input_reference", "user", "extra_headers", "model"}
},
**(
{"duration": _duration_from_seconds(incoming.get("seconds"))}
if "seconds" in incoming and "duration" not in incoming
else {}
),
**(
{"aspect_ratio": incoming.get("aspect_ratio") or _SIZE_TO_ASPECT_RATIO.get(str(size), "16:9")}
if size and "aspect_ratio" not in incoming
else {}
),
**(
{"image": incoming.get("image") or incoming.get("input_reference")}
if incoming.get("input_reference") and "image" not in incoming
else {}
),
}
def _resolve_api_base(
self,
api_base: str | None,
api_key: str | None,
litellm_params: GenericLiteLLMParams | dict | None,
) -> str:
from litellm.llms.xai.oauth import XAIOAuthAuthenticator, should_use_xai_oauth
params: Final = (
litellm_params.model_dump()
if isinstance(litellm_params, GenericLiteLLMParams)
else (litellm_params or {})
)
if should_use_xai_oauth(params) and not XAIModelInfo.get_api_key(api_key):
return XAIOAuthAuthenticator().get_api_base().rstrip("/")
resolved: Final = (
api_base
or (params.get("api_base") if isinstance(params, dict) else None)
or get_secret_str("XAI_API_BASE")
or get_secret_str("XAI_OAUTH_API_BASE")
or XAI_API_BASE
)
return str(resolved).rstrip("/")
def _v1_root(self, api_base: str) -> str:
base: Final = api_base.rstrip("/")
if base.endswith("/v1"):
return base
return f"{base}/v1"
def validate_environment(
self,
headers: dict,
model: str,
api_key: str | None = None,
litellm_params: GenericLiteLLMParams | None = None,
) -> dict:
from litellm.llms.xai.oauth import (
XAIOAuthAuthenticator,
XAIOAuthError,
should_use_xai_oauth,
)
params: Final = litellm_params.model_dump() if litellm_params is not None else {}
resolved_api_key: Final = api_key or (litellm_params.api_key if litellm_params else None)
dynamic_api_key: Final = XAIModelInfo.get_api_key(resolved_api_key)
if should_use_xai_oauth(params) and not dynamic_api_key:
try:
headers["Authorization"] = f"Bearer {XAIOAuthAuthenticator().get_access_token()}"
except XAIOAuthError as exc:
raise AuthenticationError(
model=model or "xai-video",
llm_provider="xai",
message=str(exc),
) from exc
else:
if not dynamic_api_key:
raise AuthenticationError(
model=model or "xai-video",
llm_provider="xai",
message=(
"Missing xAI credentials for video generation. "
"Pass api_key / XAI_API_KEY, or set use_xai_oauth=True."
),
)
headers["Authorization"] = f"Bearer {dynamic_api_key}"
if "content-type" not in headers and "Content-Type" not in headers:
headers["Content-Type"] = "application/json"
return headers
def get_complete_url(
self,
model: str,
api_base: str | None,
litellm_params: dict,
) -> str:
resolved: Final = self._resolve_api_base(
api_base=api_base,
api_key=litellm_params.get("api_key") if litellm_params else None,
litellm_params=litellm_params,
)
if not model:
return self._v1_root(resolved)
return f"{self._v1_root(resolved)}/videos/generations"
def transform_video_create_request(
self,
model: str,
prompt: str,
api_base: str,
video_create_optional_request_params: dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> tuple[dict, RequestFiles, str]:
copied: Final = {
key: video_create_optional_request_params[key]
for key in (
"image",
"images",
"duration",
"resolution_name",
"aspect_ratio",
"size",
)
if video_create_optional_request_params.get(key) is not None
}
return (
{
"model": XAIModelInfo.get_base_model(model) or model,
**({"prompt": prompt} if prompt else {}),
**copied,
**({"duration": 6} if "duration" not in copied else {}),
},
[],
api_base,
)
def transform_video_create_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
custom_llm_provider: str | None = None,
request_data: dict | None = None,
) -> VideoObject:
response_data: Final = raw_response.json()
request_id: Final = response_data.get("request_id") or response_data.get("id")
if not request_id:
raise ValueError(f"xAI video generation response missing request_id: {response_data}")
usage: Final = response_data.get("usage") or {}
video_obj: Final = VideoObject(
id=str(request_id),
object="video",
status="processing",
created_at=int(time.time()),
model=XAIModelInfo.get_base_model(model) or model,
progress=0,
)
if custom_llm_provider:
video_obj.id = encode_video_id_with_provider(
video_obj.id, custom_llm_provider, model
)
video_obj.usage = usage if isinstance(usage, dict) else {}
video_obj._hidden_params["video_url"] = None
return video_obj
def transform_video_status_retrieve_request(
self,
video_id: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> tuple[str, dict]:
original_id: Final = extract_original_video_id(video_id)
return f"{self._v1_root(api_base)}/videos/{original_id}", {}
def transform_video_status_retrieve_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
custom_llm_provider: str | None = None,
) -> VideoObject:
response_data: Final = raw_response.json()
status_raw: Final = str(response_data.get("status") or "processing").lower()
status: Final = _STATUS_MAP.get(status_raw, status_raw)
video_meta: Final = response_data.get("video") or {}
video_url: Final = video_meta.get("url") if isinstance(video_meta, dict) else None
seconds: Final = (
str(video_meta.get("duration"))
if isinstance(video_meta, dict) and video_meta.get("duration") is not None
else None
)
request_id: Final = (
response_data.get("request_id")
or response_data.get("id")
or (video_url.split("/")[-1].replace(".mp4", "") if video_url else "unknown")
)
video_obj: Final = VideoObject(
id=str(request_id),
object="video",
status=status,
created_at=response_data.get("created_at") or int(time.time()),
completed_at=int(time.time()) if status == "completed" else None,
model=response_data.get("model"),
progress=response_data.get("progress"),
seconds=seconds,
usage=response_data.get("usage") if isinstance(response_data.get("usage"), dict) else {},
)
video_obj._hidden_params["video_url"] = video_url
if custom_llm_provider and video_obj.id and video_obj.id != "unknown":
video_obj.id = encode_video_id_with_provider(
video_obj.id, custom_llm_provider, response_data.get("model")
)
return video_obj
def transform_video_content_request(
self,
video_id: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
variant: str | None = None,
) -> tuple[str, dict]:
original_id: Final = extract_original_video_id(video_id)
return f"{self._v1_root(api_base)}/videos/{original_id}", {}
def _video_cdn_url(self, raw_response: httpx.Response) -> str | None:
content_type: Final = (raw_response.headers.get("content-type") or "").lower()
if "application/json" not in content_type and raw_response.content[:1] != b"{":
return None
payload: Final = raw_response.json()
if not isinstance(payload, dict):
return None
video_meta: Final = payload.get("video") or {}
url: Final = video_meta.get("url") if isinstance(video_meta, dict) else None
if isinstance(url, str) and url:
return url
raise ValueError(
f"xAI video not ready for download (status={payload.get('status')}): {payload}"
)
def transform_video_content_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
) -> bytes:
url: Final = self._video_cdn_url(raw_response)
if url is None:
return raw_response.content
httpx_client: Final[HTTPHandler] = _get_httpx_client()
video_response: Final = httpx_client.get(url)
video_response.raise_for_status()
return video_response.content
async def async_transform_video_content_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
) -> bytes:
url: Final = self._video_cdn_url(raw_response)
if url is None:
return raw_response.content
async_httpx_client: Final[AsyncHTTPHandler] = get_async_httpx_client(
llm_provider=litellm.LlmProviders.XAI,
)
video_response: Final = await async_httpx_client.get(url)
video_response.raise_for_status()
return video_response.content
def transform_video_remix_request(
self,
video_id: str,
prompt: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
extra_body: dict[str, Any] | None = None,
) -> tuple[str, dict]:
raise NotImplementedError("Video remix is not supported by xAI Imagine API")
def transform_video_remix_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
custom_llm_provider: str | None = None,
) -> VideoObject:
raise NotImplementedError("Video remix is not supported by xAI Imagine API")
def transform_video_list_request(
self,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
after: str | None = None,
limit: int | None = None,
order: str | None = None,
extra_query: dict[str, Any] | None = None,
) -> tuple[str, dict]:
raise NotImplementedError("Video listing is not supported by xAI Imagine API")
def transform_video_list_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
custom_llm_provider: str | None = None,
) -> dict[str, str]:
raise NotImplementedError("Video listing is not supported by xAI Imagine API")
def transform_video_delete_request(
self,
video_id: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> tuple[str, dict]:
raise NotImplementedError("Video delete is not supported by xAI Imagine API")
def transform_video_delete_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
) -> VideoObject:
raise NotImplementedError("Video delete is not supported by xAI Imagine API")

75
litellm/model_prices_and_context_window_backup.json
View file

@ -59300,7 +59300,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59316,7 +59317,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59332,7 +59334,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59348,7 +59351,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59364,7 +59368,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59380,7 +59385,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59397,7 +59403,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59413,7 +59420,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -63792,5 +63800,56 @@
"supports_tool_choice": false,
"supports_response_schema": true,
"supports_vision": false
},
"xai/grok-imagine-video": {
"litellm_provider": "xai",
"mode": "video_generation",
"output_cost_per_video_per_second": 0.05,
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/videos",
"/v1/videos/generations"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"video"
]
},
"xai/grok-imagine-video-1.5": {
"litellm_provider": "xai",
"mode": "video_generation",
"output_cost_per_video_per_second": 0.08,
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/videos",
"/v1/videos/generations"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"video"
]
},
"xai/grok-imagine-video-1.5-preview": {
"litellm_provider": "xai",
"mode": "video_generation",
"output_cost_per_video_per_second": 0.08,
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/videos",
"/v1/videos/generations"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"video"
]
}
}

6
litellm/provider_endpoints_support_backup.json
View file

@ -2342,7 +2342,8 @@
"messages": true,
"responses": true,
"embeddings": false,
"image_generations": false,
"image_generations": true,
"image_edits": true,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
@ -2350,7 +2351,8 @@
"rerank": false,
"a2a": true,
"interactions": true,
"realtime": true
"realtime": true,
"video_generations": true
}
},
"xinference": {

22
litellm/proxy/image_endpoints/endpoints.py
View file

@ -299,12 +299,22 @@ async def image_edit_api(
if mask_files:
data["mask"] = mask_files
for _field in ("image", "mask"):
if _field in data and isinstance(data[_field], str):
raise HTTPException(
status_code=422,
detail=f"'{_field}' must be provided as a multipart file upload, not a string.",
)
invalid_image_fields: Final = tuple(
field
for field in ("image", "mask")
if field in data
and isinstance(data[field], str)
and not (
data[field].startswith("http://")
or data[field].startswith("https://")
or data[field].startswith("data:image/")
)
)
if invalid_image_fields:
raise HTTPException(
status_code=422,
detail=f"'{invalid_image_fields[0]}' must be a multipart file, http(s) URL, or data:image URI.",
)
# Ensure prompt exists in data (default to None for models that don't require it)
if "prompt" not in data:

27
litellm/proxy/video_endpoints/endpoints.py
View file

@ -20,6 +20,7 @@ from litellm.proxy.video_endpoints.utils import (
encode_character_id_in_response,
extract_model_from_target_model_names,
get_custom_provider_from_data,
resolve_video_request_model,
video_reference_to_id,
)
from litellm.types.videos.utils import (
@ -262,12 +263,13 @@ async def video_status(
if custom_llm_provider:
data["custom_llm_provider"] = custom_llm_provider
# Resolve model_name from model_id if available
# This allows the router to automatically inject litellm_params from the model config
if model_id_from_decoded and llm_router:
resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded)
if resolved_model:
data["model"] = resolved_model
resolved_model: Final = resolve_video_request_model(
model_id_from_decoded=model_id_from_decoded,
query_model=request.query_params.get("model"),
llm_router=llm_router,
)
if resolved_model:
data["model"] = resolved_model
# Process request using ProxyBaseLLMRequestProcessing
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
@ -360,12 +362,13 @@ async def video_content(
if custom_llm_provider:
data["custom_llm_provider"] = custom_llm_provider
# Resolve model_name from model_id if available
# This allows the router to automatically inject litellm_params from the model config
if model_id_from_decoded and llm_router:
resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded)
if resolved_model:
data["model"] = resolved_model
resolved_content_model: Final = resolve_video_request_model(
model_id_from_decoded=model_id_from_decoded,
query_model=request.query_params.get("model"),
llm_router=llm_router,
)
if resolved_content_model:
data["model"] = resolved_content_model
# Process request using ProxyBaseLLMRequestProcessing
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
try:

23
litellm/proxy/video_endpoints/utils.py
View file

@ -1,10 +1,31 @@
from typing import Any, Final
from typing import Any, Final, Protocol
import orjson
from litellm.types.videos.utils import encode_character_id_with_provider
class VideoModelIdResolver(Protocol):
def resolve_model_name_from_model_id(self, model_id: str | None) -> str | None: ...
def resolve_video_request_model(
*,
model_id_from_decoded: str | None,
query_model: str | None,
llm_router: VideoModelIdResolver | None,
) -> str | None:
if model_id_from_decoded:
if llm_router is not None:
resolved: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded)
if isinstance(resolved, str) and resolved:
return resolved
return model_id_from_decoded
if isinstance(query_model, str) and query_model:
return query_model
return None
def extract_model_from_target_model_names(target_model_names: Any) -> str | None:
if isinstance(target_model_names, str):
target_model_names = [m.strip() for m in target_model_names.split(",") if m.strip()]

14
litellm/utils.py
View file

@ -9136,6 +9136,12 @@ class ProviderConfigManager:
)
return get_modelscope_image_generation_config(model)
elif LlmProviders.XAI == provider:
from litellm.llms.xai.image_generation import (
get_xai_image_generation_config,
)
return get_xai_image_generation_config(model)
return None
@staticmethod
@ -9163,6 +9169,10 @@ class ProviderConfigManager:
from litellm.llms.runwayml.videos.transformation import RunwayMLVideoConfig
return RunwayMLVideoConfig()
elif LlmProviders.XAI == provider:
from litellm.llms.xai.videos.transformation import XAIVideoConfig
return XAIVideoConfig()
elif LlmProviders.HOSTED_VLLM == provider:
from litellm.llms.hosted_vllm.videos import get_hosted_vllm_video_config
@ -9287,6 +9297,10 @@ class ProviderConfigManager:
)
return get_openrouter_image_edit_config(model)
elif LlmProviders.XAI == provider:
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
return XAIImageEditConfig()
return None
@staticmethod

66
litellm/videos/main.py
View file

@ -30,6 +30,51 @@ from litellm.videos.utils import VideoGenerationRequestUtils
llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler()
def _litellm_provider_from_cost_entry(info: object) -> str | None:
if isinstance(info, dict):
catalog_provider: Final = info.get("litellm_provider")
if isinstance(catalog_provider, str) and catalog_provider:
return catalog_provider
return None
def _provider_from_prefixed_model(model: str) -> str | None:
if "/" not in model:
return None
try:
_, provider, _, _ = get_llm_provider(model=model)
except Exception:
return None
return provider
def _custom_llm_provider_from_model(model: str) -> str | None:
return (
_provider_from_prefixed_model(model)
or _litellm_provider_from_cost_entry(litellm.model_cost.get(model))
or _litellm_provider_from_cost_entry(litellm.model_cost.get(f"xai/{model}"))
or ("xai" if model.startswith("grok-imagine-video") else None)
)
def _provider_for_video_id(
video_id: str,
custom_llm_provider: str | None,
model: object | None = None,
) -> str:
if custom_llm_provider is not None:
return custom_llm_provider
decoded: Final = decode_video_id_with_provider(video_id)
from_id: Final = decoded.get("custom_llm_provider")
if from_id:
return from_id
if isinstance(model, str) and model:
from_model: Final = _custom_llm_provider_from_model(model)
if from_model:
return from_model
return "openai"
##### Video Generation #######################
@client
async def avideo_generation(
@ -317,10 +362,9 @@ def video_content(
litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None)
_is_async: Final = kwargs.pop("async_call", False) is True
# Try to decode provider from video_id if not explicitly provided
if custom_llm_provider is None:
decoded: Final = decode_video_id_with_provider(video_id)
custom_llm_provider = decoded.get("custom_llm_provider") or "openai"
custom_llm_provider = _provider_for_video_id(
video_id, custom_llm_provider, kwargs.get("model")
)
# get llm provider logic
litellm_params: Final = GenericLiteLLMParams(**kwargs)
@ -412,10 +456,9 @@ async def avideo_content(
loop: Final = asyncio.get_event_loop()
kwargs["async_call"] = True
# Try to decode provider from video_id if not explicitly provided
if custom_llm_provider is None:
decoded: Final = decode_video_id_with_provider(video_id)
custom_llm_provider = decoded.get("custom_llm_provider") or "openai"
custom_llm_provider = _provider_for_video_id(
video_id, custom_llm_provider, kwargs.get("model")
)
func: Final = partial(
video_content,
@ -1019,10 +1062,9 @@ def video_status(
response: Final = VideoObject(**mock_response)
return response
# Try to decode provider from video_id if not explicitly provided
if custom_llm_provider is None:
decoded: Final = decode_video_id_with_provider(video_id)
custom_llm_provider = decoded.get("custom_llm_provider") or "openai"
custom_llm_provider = _provider_for_video_id(
video_id, custom_llm_provider, kwargs.get("model")
)
# get llm provider logic
litellm_params: Final = GenericLiteLLMParams(**kwargs)

75
model_prices_and_context_window.json
View file

@ -59300,7 +59300,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59316,7 +59317,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59332,7 +59334,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59348,7 +59351,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59364,7 +59368,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59380,7 +59385,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59397,7 +59403,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -59413,7 +59420,8 @@
"mode": "image_generation",
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/images/generations"
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
@ -63792,5 +63800,56 @@
"supports_tool_choice": false,
"supports_response_schema": true,
"supports_vision": false
},
"xai/grok-imagine-video": {
"litellm_provider": "xai",
"mode": "video_generation",
"output_cost_per_video_per_second": 0.05,
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/videos",
"/v1/videos/generations"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"video"
]
},
"xai/grok-imagine-video-1.5": {
"litellm_provider": "xai",
"mode": "video_generation",
"output_cost_per_video_per_second": 0.08,
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/videos",
"/v1/videos/generations"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"video"
]
},
"xai/grok-imagine-video-1.5-preview": {
"litellm_provider": "xai",
"mode": "video_generation",
"output_cost_per_video_per_second": 0.08,
"source": "https://docs.x.ai/docs/models",
"supported_endpoints": [
"/v1/videos",
"/v1/videos/generations"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"video"
]
}
}

6
provider_endpoints_support.json
View file

@ -2653,7 +2653,8 @@
"messages": true,
"responses": true,
"embeddings": false,
"image_generations": false,
"image_generations": true,
"image_edits": true,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
@ -2661,7 +2662,8 @@
"rerank": false,
"a2a": true,
"interactions": true,
"realtime": true
"realtime": true,
"video_generations": true
}
},
"xinference": {

167
tests/test_litellm/llms/xai/test_xai_image_edit.py Normal file
View file

@ -0,0 +1,167 @@
from unittest.mock import MagicMock, patch
import httpx
import pytest
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
def test_provider_config_manager_returns_xai_image_edit_config():
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
config = ProviderConfigManager.get_provider_image_edit_config(
model="grok-imagine-image",
provider=LlmProviders.XAI,
)
assert isinstance(config, XAIImageEditConfig)
def test_get_complete_url_default():
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
url = XAIImageEditConfig().get_complete_url(
model="grok-imagine-image",
api_base=None,
litellm_params={},
)
assert url.endswith("/v1/images/edits")
def test_uses_json_not_multipart():
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
assert XAIImageEditConfig().use_multipart_form_data() is False
def test_map_size_to_aspect_ratio():
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
mapped = XAIImageEditConfig().map_openai_params(
image_edit_optional_params={"size": "1024x1792", "n": 1},
model="grok-imagine-image",
drop_params=True,
)
assert mapped["aspect_ratio"] == "9:16"
assert mapped["n"] == 1
assert "size" not in mapped
def test_validate_environment_requires_credentials():
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
with pytest.raises(Exception, match="Missing xAI credentials"):
XAIImageEditConfig().validate_environment(
headers={},
model="grok-imagine-image",
api_key=None,
litellm_params={},
)
def test_validate_environment_oauth_injects_bearer():
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
with patch(
"litellm.llms.xai.oauth.XAIOAuthAuthenticator.get_access_token",
return_value="oauth-token",
):
headers = XAIImageEditConfig().validate_environment(
headers={},
model="grok-imagine-image",
api_key=None,
litellm_params={"use_xai_oauth": True},
)
assert headers["Authorization"] == "Bearer oauth-token"
def test_map_string_n_to_int():
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
mapped = XAIImageEditConfig().map_openai_params(
image_edit_optional_params={"n": "1"},
model="grok-imagine-image",
drop_params=True,
)
assert mapped["n"] == 1
assert isinstance(mapped["n"], int)
def test_transform_string_n_to_int():
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
data, _ = XAIImageEditConfig().transform_image_edit_request(
model="xai/grok-imagine-image",
prompt="make the cube red",
image="https://imgen.x.ai/source.jpeg",
image_edit_optional_request_params={"n": "1"},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert data["n"] == 1
assert isinstance(data["n"], int)
def test_transform_bytes_to_data_uri_and_response():
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
config = XAIImageEditConfig()
data, files = config.transform_image_edit_request(
model="xai/grok-imagine-image",
prompt="make the cube red",
image=b"\xff\xd8\xfffakejpeg",
image_edit_optional_request_params={"aspect_ratio": "1:1"},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert files == []
assert data["model"] == "grok-imagine-image"
assert data["prompt"] == "make the cube red"
assert data["aspect_ratio"] == "1:1"
assert data["image"]["url"].startswith("data:image/jpeg;base64,")
raw = httpx.Response(
200,
json={"data": [{"url": "https://imgen.x.ai/edited.jpeg", "mime_type": "image/jpeg"}]},
)
response = config.transform_image_edit_response(
model="grok-imagine-image",
raw_response=raw,
logging_obj=MagicMock(),
)
assert response.data is not None
assert response.data[0].url == "https://imgen.x.ai/edited.jpeg"
def test_transform_http_url_passthrough():
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
data, files = XAIImageEditConfig().transform_image_edit_request(
model="grok-imagine-image",
prompt="make it night",
image="https://imgen.x.ai/source.jpeg",
image_edit_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert files == []
assert data["image"] == {"url": "https://imgen.x.ai/source.jpeg"}
def test_transform_multiple_images_uses_images_array():
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
data, _ = XAIImageEditConfig().transform_image_edit_request(
model="grok-imagine-image",
prompt="combine styles",
image=["https://imgen.x.ai/a.jpeg", "https://imgen.x.ai/b.jpeg"],
image_edit_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert "image" not in data
assert data["images"] == [
{"url": "https://imgen.x.ai/a.jpeg"},
{"url": "https://imgen.x.ai/b.jpeg"},
]

129
tests/test_litellm/llms/xai/test_xai_image_generation.py Normal file
View file

@ -0,0 +1,129 @@
from unittest.mock import MagicMock, patch
import httpx
import pytest
from litellm.llms.xai.image_generation.transformation import XAIImageGenerationConfig
from litellm.types.utils import ImageResponse, LlmProviders
from litellm.utils import ProviderConfigManager
def test_provider_config_manager_returns_xai_image_config():
config = ProviderConfigManager.get_provider_image_generation_config(
model="grok-imagine-image",
provider=LlmProviders.XAI,
)
assert isinstance(config, XAIImageGenerationConfig)
def test_map_size_to_aspect_ratio():
mapped = XAIImageGenerationConfig().map_openai_params(
non_default_params={"size": "1024x1792"},
optional_params={},
model="grok-imagine-image",
drop_params=True,
)
assert mapped["aspect_ratio"] == "9:16"
assert "size" not in mapped
def test_get_complete_url_default():
url = XAIImageGenerationConfig().get_complete_url(
api_base=None,
api_key=None,
model="grok-imagine-image",
optional_params={},
litellm_params={},
)
assert url.endswith("/v1/images/generations")
def test_validate_environment_requires_credentials():
with pytest.raises(Exception, match="Missing xAI credentials"):
XAIImageGenerationConfig().validate_environment(
headers={},
model="grok-imagine-image",
messages=[],
optional_params={},
litellm_params={},
api_key=None,
)
def test_validate_environment_oauth_injects_bearer():
with patch(
"litellm.llms.xai.oauth.XAIOAuthAuthenticator.get_access_token",
return_value="oauth-token",
):
headers = XAIImageGenerationConfig().validate_environment(
headers={},
model="grok-imagine-image",
messages=[],
optional_params={},
litellm_params={"use_xai_oauth": True},
api_key=None,
)
assert headers["Authorization"] == "Bearer oauth-token"
def test_map_string_n_to_int():
mapped = XAIImageGenerationConfig().map_openai_params(
non_default_params={"n": "1"},
optional_params={},
model="grok-imagine-image",
drop_params=True,
)
assert mapped["n"] == 1
assert isinstance(mapped["n"], int)
def test_transform_string_n_to_int():
request = XAIImageGenerationConfig().transform_image_generation_request(
model="xai/grok-imagine-image",
prompt="a red apple",
optional_params={"n": "1"},
litellm_params={},
headers={},
)
assert request["n"] == 1
assert isinstance(request["n"], int)
def test_transform_request_and_response():
config = XAIImageGenerationConfig()
request = config.transform_image_generation_request(
model="xai/grok-imagine-image",
prompt="a red apple",
optional_params={"aspect_ratio": "1:1"},
litellm_params={},
headers={},
)
assert request == {
"model": "grok-imagine-image",
"prompt": "a red apple",
"aspect_ratio": "1:1",
}
raw = httpx.Response(
200,
json={
"data": [
{
"url": "https://imgen.x.ai/example.jpeg",
"mime_type": "image/jpeg",
}
]
},
)
response = config.transform_image_generation_response(
model="grok-imagine-image",
raw_response=raw,
model_response=ImageResponse(),
logging_obj=MagicMock(),
request_data=request,
optional_params={},
litellm_params={},
encoding=None,
)
assert response.data is not None
assert response.data[0].url == "https://imgen.x.ai/example.jpeg"

206
tests/test_litellm/llms/xai/test_xai_video_generation.py Normal file
View file

@ -0,0 +1,206 @@
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from litellm.llms.xai.videos.transformation import XAIVideoConfig
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
def test_provider_config_manager_returns_xai_video_config():
config = ProviderConfigManager.get_provider_video_config(
model="grok-imagine-video",
provider=LlmProviders.XAI,
)
assert isinstance(config, XAIVideoConfig)
def test_map_seconds_and_size():
mapped = XAIVideoConfig().map_openai_params(
video_create_optional_params={"seconds": "10", "size": "1280x720"},
model="grok-imagine-video",
drop_params=True,
)
assert mapped["duration"] == 10
assert mapped["aspect_ratio"] == "16:9"
def test_map_seconds_nine_is_not_clamped_to_six():
mapped = XAIVideoConfig().map_openai_params(
video_create_optional_params={"seconds": "9"},
model="grok-imagine-video-1.5",
drop_params=True,
)
assert mapped["duration"] == 9
assert "seconds" not in mapped
def test_get_complete_url_create_and_status_root():
config = XAIVideoConfig()
create_url = config.get_complete_url(
model="grok-imagine-video",
api_base="https://api.x.ai/v1",
litellm_params={},
)
assert create_url == "https://api.x.ai/v1/videos/generations"
status_root = config.get_complete_url(
model="",
api_base="https://api.x.ai/v1",
litellm_params={},
)
assert status_root == "https://api.x.ai/v1"
def test_validate_environment_oauth_injects_bearer():
with patch(
"litellm.llms.xai.oauth.XAIOAuthAuthenticator.get_access_token",
return_value="oauth-token",
):
headers = XAIVideoConfig().validate_environment(
headers={},
model="grok-imagine-video",
api_key=None,
litellm_params=GenericLiteLLMParams(use_xai_oauth=True),
)
assert headers["Authorization"] == "Bearer oauth-token"
def test_transform_create_and_status_response():
config = XAIVideoConfig()
data, files, api_base = config.transform_video_create_request(
model="xai/grok-imagine-video",
prompt="a cat walking",
api_base="https://api.x.ai/v1/videos/generations",
video_create_optional_request_params={"duration": 6},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert data["model"] == "grok-imagine-video"
assert data["prompt"] == "a cat walking"
assert data["duration"] == 6
assert files == []
created = config.transform_video_create_response(
model="grok-imagine-video",
raw_response=httpx.Response(200, json={"request_id": "req-123"}),
logging_obj=MagicMock(),
custom_llm_provider="xai",
)
assert created.status == "processing"
assert created.id
status_url, params = config.transform_video_status_retrieve_request(
video_id=created.id,
api_base="https://api.x.ai/v1",
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert status_url.endswith("/videos/req-123")
assert params == {}
status = config.transform_video_status_retrieve_response(
raw_response=httpx.Response(
200,
json={
"status": "done",
"request_id": "req-123",
"video": {"url": "https://vidgen.x.ai/x.mp4", "duration": 6},
"progress": 100,
"model": "grok-imagine-video",
},
),
logging_obj=MagicMock(),
custom_llm_provider="xai",
)
assert status.status == "completed"
assert status.seconds == "6"
assert status._hidden_params.get("video_url") == "https://vidgen.x.ai/x.mp4"
def test_validate_environment_requires_credentials():
with pytest.raises(Exception, match="Missing xAI credentials"):
XAIVideoConfig().validate_environment(
headers={},
model="grok-imagine-video",
api_key=None,
litellm_params=GenericLiteLLMParams(),
)
def test_content_request_is_get_status():
url, params = XAIVideoConfig().transform_video_content_request(
video_id="req-123",
api_base="https://api.x.ai/v1",
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert url == "https://api.x.ai/v1/videos/req-123"
assert params == {}
def test_content_response_fetches_cdn_via_shared_client():
status = httpx.Response(
200,
headers={"content-type": "application/json"},
json={"status": "done", "video": {"url": "https://vidgen.x.ai/x.mp4"}},
)
cdn = MagicMock()
video_resp = MagicMock()
video_resp.content = b"mp4-bytes"
video_resp.raise_for_status.return_value = None
cdn.get.return_value = video_resp
with patch(
"litellm.llms.xai.videos.transformation._get_httpx_client",
return_value=cdn,
):
body = XAIVideoConfig().transform_video_content_response(status, logging_obj=MagicMock())
assert body == b"mp4-bytes"
cdn.get.assert_called_once_with("https://vidgen.x.ai/x.mp4")
@pytest.mark.asyncio
async def test_async_content_response_does_not_use_sync_client():
status = httpx.Response(
200,
headers={"content-type": "application/json"},
json={"status": "done", "video": {"url": "https://vidgen.x.ai/x.mp4"}},
)
async_client = MagicMock()
video_resp = MagicMock()
video_resp.content = b"async-mp4"
video_resp.raise_for_status.return_value = None
async_client.get = AsyncMock(return_value=video_resp)
with patch(
"litellm.llms.xai.videos.transformation.get_async_httpx_client",
return_value=async_client,
), patch(
"litellm.llms.xai.videos.transformation._get_httpx_client",
) as sync_client:
body = await XAIVideoConfig().async_transform_video_content_response(
status, logging_obj=MagicMock()
)
assert body == b"async-mp4"
sync_client.assert_not_called()
async_client.get.assert_awaited_once_with("https://vidgen.x.ai/x.mp4")
def test_content_response_raises_when_status_has_no_url():
status = httpx.Response(
200,
headers={"content-type": "application/json"},
json={"status": "pending"},
)
with pytest.raises(ValueError, match="not ready"):
XAIVideoConfig().transform_video_content_response(status, logging_obj=MagicMock())
def test_content_response_returns_raw_bytes_when_not_json():
raw = httpx.Response(
200,
headers={"content-type": "video/mp4"},
content=b"already-mp4",
)
assert XAIVideoConfig().transform_video_content_response(raw, logging_obj=MagicMock()) == b"already-mp4"

48
tests/test_litellm/proxy/image_endpoints/test_endpoints.py
View file

@ -167,3 +167,51 @@ def test_image_edit_multipart_n_that_is_not_a_number_is_left_alone(monkeypatch):
assert response.status_code == 200
assert captured["n"] == "two"
def test_image_edit_http_url_is_accepted(monkeypatch):
captured: Dict[str, Any] = {}
response = _image_edit_client(monkeypatch, captured).post(
"/v1/images/edits",
data={
"model": "grok-imagine-image",
"prompt": "make it night",
"image": "https://imgen.x.ai/source.jpeg",
},
)
assert response.status_code == 200
assert captured["image"] == "https://imgen.x.ai/source.jpeg"
def test_image_edit_data_uri_is_accepted(monkeypatch):
captured: Dict[str, Any] = {}
response = _image_edit_client(monkeypatch, captured).post(
"/v1/images/edits",
data={
"model": "grok-imagine-image",
"prompt": "make it red",
"image": "data:image/jpeg;base64,abc",
},
)
assert response.status_code == 200
assert captured["image"] == "data:image/jpeg;base64,abc"
def test_image_edit_plain_string_image_is_rejected(monkeypatch):
captured: Dict[str, Any] = {}
response = _image_edit_client(monkeypatch, captured).post(
"/v1/images/edits",
data={
"model": "grok-imagine-image",
"prompt": "make it red",
"image": "not-a-url-or-file",
},
)
assert response.status_code == 422
assert "multipart file" in response.json()["detail"]

70
tests/test_litellm/proxy/video_endpoints/test_endpoints.py
View file

@ -309,6 +309,40 @@ async def test_status__header_provider_beats_decoded_id(harness):
assert data["model"] == "azure-sora"
@pytest.mark.asyncio
async def test_status__resolve_fail_keeps_decoded_model_id(harness):
encoded = encode_video_id_with_provider(
"9b444cea-aaaa-bbbb-cccc-dddddddddddd",
"xai",
"grok-imagine-video-1.5",
)
await call_status(harness, encoded)
harness.resolve_model.assert_called_once_with("grok-imagine-video-1.5")
assert harness.processor_data() == {
"video_id": encoded,
"custom_llm_provider": "xai",
"model": "grok-imagine-video-1.5",
}
@pytest.mark.asyncio
async def test_status__query_model_on_plain_id(harness):
await call_status(
harness,
"9b444cea-aaaa-bbbb-cccc-dddddddddddd",
query={"model": "grok-imagine-video-1.5"},
)
harness.resolve_model.assert_not_called()
assert harness.processor_data() == {
"video_id": "9b444cea-aaaa-bbbb-cccc-dddddddddddd",
"custom_llm_provider": "openai",
"model": "grok-imagine-video-1.5",
}
# =========================================================================== #
# GET /v1/videos/{video_id}/content - video_content #
# =========================================================================== #
@ -365,6 +399,42 @@ async def test_content__model_encoded_id(harness):
}
@pytest.mark.asyncio
async def test_content__query_model_on_plain_id(harness):
harness.base_process.return_value = b"x"
await call_content(
harness,
"9b444cea-aaaa-bbbb-cccc-dddddddddddd",
query={"model": "grok-imagine-video-1.5"},
)
harness.resolve_model.assert_not_called()
assert harness.processor_data() == {
"video_id": "9b444cea-aaaa-bbbb-cccc-dddddddddddd",
"model": "grok-imagine-video-1.5",
}
@pytest.mark.asyncio
async def test_content__resolve_fail_keeps_decoded_model_id(harness):
harness.base_process.return_value = b"x"
encoded = encode_video_id_with_provider(
"9b444cea-aaaa-bbbb-cccc-dddddddddddd",
"xai",
"grok-imagine-video-1.5",
)
await call_content(harness, encoded)
harness.resolve_model.assert_called_once_with("grok-imagine-video-1.5")
assert harness.processor_data() == {
"video_id": encoded,
"custom_llm_provider": "xai",
"model": "grok-imagine-video-1.5",
}
# =========================================================================== #
# POST /v1/videos/edits - video_edit #
# =========================================================================== #

47
tests/test_litellm/proxy/video_endpoints/test_utils.py
View file

@ -21,6 +21,7 @@ from litellm.proxy.video_endpoints.utils import (
encode_character_id_in_response,
extract_model_from_target_model_names,
get_custom_provider_from_data,
resolve_video_request_model,
video_reference_to_id,
)
from litellm.types.videos.utils import (
@ -28,6 +29,52 @@ from litellm.types.videos.utils import (
encode_character_id_with_provider,
)
# =========================================================================== #
# resolve_video_request_model
# =========================================================================== #
class _Resolver:
def __init__(self, mapping: dict[str, str | None]):
self.mapping = mapping
def resolve_model_name_from_model_id(self, model_id: str | None) -> str | None:
return self.mapping.get(model_id) if model_id else None
def test_resolve_video_request_model__router_hit():
assert (
resolve_video_request_model(
model_id_from_decoded="deployment-123",
query_model="ignored",
llm_router=_Resolver({"deployment-123": "azure-sora"}),
)
== "azure-sora"
)
def test_resolve_video_request_model__keeps_decoded_id_when_router_misses():
assert (
resolve_video_request_model(
model_id_from_decoded="grok-imagine-video-1.5",
query_model=None,
llm_router=_Resolver({}),
)
== "grok-imagine-video-1.5"
)
def test_resolve_video_request_model__query_model_on_plain_id():
assert (
resolve_video_request_model(
model_id_from_decoded=None,
query_model="grok-imagine-video-1.5",
llm_router=None,
)
== "grok-imagine-video-1.5"
)
# =========================================================================== #
# extract_model_from_target_model_names
# =========================================================================== #

24
tests/test_litellm/videos/test_main.py
View file

@ -178,6 +178,15 @@ def test_video_content__plain_id_defaults_to_openai(seams):
assert seams.kwargs_of("video_content_handler")["custom_llm_provider"] == "openai"
def test_video_content__plain_id_with_grok_model_uses_xai(seams):
videos_main.video_content(
video_id="9b444cea-aaaa-bbbb-cccc-dddddddddddd",
model="grok-imagine-video-1.5",
)
assert seams.kwargs_of("video_content_handler")["custom_llm_provider"] == "xai"
def test_video_remix__dispatch_and_provider_from_id(seams):
result = videos_main.video_remix(video_id=AZURE_VIDEO_ID, prompt="new colors")
@ -377,6 +386,21 @@ async def test_avideo_content__pre_decodes_provider_before_delegating():
assert sync.call_args.kwargs["custom_llm_provider"] == "azure"
@pytest.mark.asyncio
async def test_avideo_content__plain_id_with_grok_model_uses_xai():
sentinel = b"mp4-bytes"
with patch.object(
videos_main, "video_content", MagicMock(return_value=sentinel)
) as sync:
result = await videos_main.avideo_content(
video_id="9b444cea-aaaa-bbbb-cccc-dddddddddddd",
model="grok-imagine-video-1.5",
)
assert result is sentinel
assert sync.call_args.kwargs["custom_llm_provider"] == "xai"
# =========================================================================== #
# Credential passthrough - DB/YAML model-config credentials the router injects
# via kwargs must reach the provider call for EVERY video handler, carried in