scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-31 06:56:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 416

fix(xai): type imagine transforms and accept the handler client on video downloads

Browse source 
Video content hooks now take the client the HTTP handler passes in, image
responses are parsed through one typed helper that raises a provider error
when no images come back, and video provider resolution no longer rebinds
the custom_llm_provider argument
This commit is contained in:
hx 2026-10-15 20:39:11 +08:00
parent 2c313b6fd0
commit 9c3a560f32
9 changed files with 356 additions and 338 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
litellm/llms/xai/common_utils.py
View file

@ -7,9 +7,10 @@ from pydantic import ConfigDict, TypeAdapter
import litellm
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ProviderSpecificModelInfo
from litellm.types.utils import ImageObject, ProviderSpecificModelInfo
USD_TICKS_PER_DOLLAR: Final = 10_000_000_000
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
@ -17,6 +18,39 @@ _JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input
_MODEL_ID: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True))
def _optional_str(value: object) -> str | None:
return value if isinstance(value, str) else None
def _image_object(entry: object) -> ImageObject | None:
if not isinstance(entry, Mapping):
return None
item: Final = _JSON_OBJECT.validate_python(entry)
return ImageObject(
url=_optional_str(item.get("url")),
b64_json=_optional_str(item.get("b64_json")) or _optional_str(item.get("b64")),
)
def xai_image_response(raw_response: httpx.Response, operation: str) -> tuple[ImageObject, ...]:
try:
payload: Final = _JSON_OBJECT.validate_python(raw_response.json())
except ValueError as exc:
raise BaseLLMException(
status_code=raw_response.status_code, message=raw_response.text, headers=raw_response.headers
) from exc
data: Final = payload.get("data")
entries: Final = _JSON_VALUES.validate_python(data) if isinstance(data, list) else ()
images: Final = tuple(image for image in map(_image_object, entries) if image is not None)
if not images:
raise BaseLLMException(
status_code=raw_response.status_code,
message=f"xAI {operation} returned no image data: {payload}",
headers=raw_response.headers,
)
return images
def xai_reported_cost_in_usd(cost_in_usd_ticks: object) -> float | None:
"""xAI bills in ticks of a dollar: https://docs.x.ai/developers/cost-tracking"""
if not isinstance(cost_in_usd_ticks, int) or isinstance(cost_in_usd_ticks, bool):
@ -105,9 +139,4 @@ class XAIModelInfo(BaseLLMModelInfo):
models: Final = _JSON_VALUES.validate_python(_JSON_OBJECT.validate_python(response.json())["data"])
litellm_model_names: Final = []
for model in models:
stripped_model_name = _MODEL_ID.validate_python(_JSON_OBJECT.validate_python(model)["id"])
litellm_model_name = "xai/" + stripped_model_name
litellm_model_names.append(litellm_model_name)
return litellm_model_names
return ["xai/" + _MODEL_ID.validate_python(_JSON_OBJECT.validate_python(model)["id"]) for model in models]

170
litellm/llms/xai/image_edit/transformation.py
View file

@ -1,23 +1,31 @@
import base64
from collections.abc import Mapping
from io import BufferedReader, BytesIO
from typing import IO, TYPE_CHECKING, Final, Protocol, cast, runtime_checkable # noqa: TID251 # see _read_all_bytes
from typing import IO, TYPE_CHECKING, Final, Protocol, TypeAlias, runtime_checkable
import httpx
from httpx._types import RequestFiles
from typing_extensions import TypeIs # noqa: TID251 # narrows untyped reference images without a runtime conversion
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.llms.xai.common_utils import XAIModelInfo, xai_image_response
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
from litellm.types.utils import FileTypes, ImageResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_Params: TypeAlias = dict[str, object]
_Headers: TypeAlias = dict[str, object]
_ImageRef: TypeAlias = dict[str, str]
_SupportedParams: TypeAlias = list[str]
_ImageInputs: TypeAlias = list[FileTypes | None]
_SIZE_TO_ASPECT_RATIO: Final = {
"1024x1024": "1:1",
"1792x1024": "16:9",
@ -38,6 +46,14 @@ class _Readable(Protocol):
def read(self) -> bytes | str: ...
def _is_json_object(value: object) -> TypeIs[Mapping[str, object]]: # guard-ok: JSON-style references have str keys
return isinstance(value, Mapping)
def _is_seekable_binary(value: object) -> TypeIs[IO[bytes]]: # guard-ok: both classes are binary IO[bytes] streams
return isinstance(value, (BytesIO, BufferedReader))
def _read_seekable(image: IO[bytes]) -> bytes:
current_pos: Final = image.tell()
image.seek(0)
@ -46,10 +62,12 @@ def _read_seekable(image: IO[bytes]) -> bytes:
return data
def _count(value: object) -> int | None:
return int(value) if isinstance(value, (int, float, str)) else None
class XAIImageEditConfig(BaseImageEditConfig):
def get_supported_openai_params(
self, model: str
) -> list: # mutable-ok: provider JSON body and base-class dict signature
def get_supported_openai_params(self, model: str) -> _SupportedParams:
return list(_OPENAI_PARAMS)
def map_openai_params(
@ -57,11 +75,11 @@ class XAIImageEditConfig(BaseImageEditConfig):
image_edit_optional_params: ImageEditOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict: # mutable-ok: provider JSON body and base-class dict signature
) -> _Params:
supported: Final = frozenset(_OPENAI_PARAMS)
allowed: Final = supported | _XAI_NATIVE_PARAMS
raw: Final = image_edit_optional_params
incoming: Final = dict(raw)
incoming: Final[Mapping[str, object]] = dict(raw)
unknown: Final = tuple(key for key in incoming if key not in allowed)
if unknown and not drop_params:
raise ValueError(
@ -76,12 +94,10 @@ class XAIImageEditConfig(BaseImageEditConfig):
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")
fields: Final = (
("aspect_ratio", aspect_ratio),
("n", int(n) if n is not None else None),
("resolution", resolution),
("n", _count(mapped.get("n"))),
("resolution", mapped.get("resolution")),
)
return {key: value for key, value in fields if value is not None}
@ -92,7 +108,7 @@ class XAIImageEditConfig(BaseImageEditConfig):
self,
model: str,
api_base: str | None,
litellm_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
litellm_params: _Params,
) -> str:
from litellm.llms.xai.oauth import XAIOAuthAuthenticator, should_use_xai_oauth
@ -108,70 +124,61 @@ class XAIImageEditConfig(BaseImageEditConfig):
return f"{base}/images/edits"
return f"{base}/v1/images/edits"
def validate_environment(
self,
headers: dict, # mutable-ok: provider JSON body and base-class dict signature
model: str,
api_key: str | None = None,
litellm_params: dict[str, object] | None = None, # mutable-ok: provider JSON body and base-class dict signature
api_base: str | None = None,
) -> dict: # mutable-ok: provider JSON body and base-class dict signature
def _bearer_token(self, model: str, api_key: str | None, litellm_params: _Params) -> str:
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 dynamic_api_key:
return dynamic_api_key
if not should_use_xai_oauth(litellm_params):
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.",
)
try:
return XAIOAuthAuthenticator().get_access_token()
except XAIOAuthError as exc:
raise AuthenticationError(model=model, llm_provider="xai", message=str(exc)) from exc
if "content-type" not in headers and "Content-Type" not in headers:
headers["Content-Type"] = "application/json"
return headers
def validate_environment(
self,
headers: _Headers,
model: str,
api_key: str | None = None,
litellm_params: _Params | None = None,
api_base: str | None = None,
) -> _Headers:
has_content_type: Final = "content-type" in headers or "Content-Type" in headers
return {
**headers,
"Authorization": f"Bearer {self._bearer_token(model, api_key, litellm_params or {})}",
**({} if has_content_type else {"Content-Type": "application/json"}),
}
def transform_image_edit_request(
self,
model: str,
prompt: str | None,
image: FileTypes | None,
image_edit_optional_request_params: dict[str, object], # mutable-ok: base-class dict signature
image: FileTypes | _ImageInputs | None,
image_edit_optional_request_params: _Params,
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: provider JSON body and base-class dict signature
) -> tuple[dict, RequestFiles]: # mutable-ok: provider JSON body and base-class dict signature
if image is None:
raise ValueError("xAI image edit requires at least one reference image.")
headers: _Headers,
) -> tuple[_Params, RequestFiles]:
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")
prompt_body: Final = {"prompt": prompt} if prompt is not None else None
many: Final = {"images": list(image_payloads)}
one: Final = image_payloads[0]
image_body: Final = {"image": one} if len(image_payloads) == 1 else many
prompt_field: Final = prompt_body or {}
image_field: Final = image_body
request: Final[dict[str, object]] = { # mutable-ok: provider JSON body and base-class dict signature
n: Final = _count(image_edit_optional_request_params.get("n"))
prompt_field: Final = {"prompt": prompt} if prompt is not None else {}
image_field: Final = (
{"image": image_payloads[0]} if len(image_payloads) == 1 else {"images": list(image_payloads)}
)
request: Final[_Params] = {
"model": XAIModelInfo.get_base_model(model) or model,
**prompt_field,
**image_field,
@ -180,7 +187,7 @@ class XAIImageEditConfig(BaseImageEditConfig):
for key in ("aspect_ratio", "resolution")
if image_edit_optional_request_params.get(key) is not None
},
**({"n": int(n)} if isinstance(n, (int, float, str)) else {}),
**({"n": n} if n is not None else {}),
}
return request, []
@ -190,45 +197,17 @@ class XAIImageEditConfig(BaseImageEditConfig):
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> ImageResponse:
try:
response_data: Final = raw_response.json()
except Exception: # noqa: BLE001 # any body-parse failure becomes a provider error carrying the raw payload
raise self.get_error_class(
error_message=raw_response.text,
status_code=raw_response.status_code,
headers=raw_response.headers,
)
return ImageResponse(data=list(xai_image_response(raw_response, "image edit")))
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], # mutable-ok: the proxy passes multi-image edits as a list of uploads
) -> tuple[FileTypes, ...]:
def _as_image_list(self, image: FileTypes | _ImageInputs | None) -> tuple[FileTypes, ...]:
if isinstance(image, list):
return tuple(item for item in image if item is not None)
return (image,)
return () if image is None else (image,)
def _to_image_url(
self, image: FileTypes
) -> dict[str, str]: # mutable-ok: provider JSON body and base-class dict signature
def _to_image_url(self, image: FileTypes) -> _ImageRef:
if isinstance(image, str):
return {"url": image}
if isinstance(image, dict):
if _is_json_object(image):
url: Final = image.get("url")
if isinstance(url, str) and url:
return {"url": url}
@ -245,9 +224,8 @@ class XAIImageEditConfig(BaseImageEditConfig):
return image
if isinstance(image, bytearray):
return bytes(image)
if isinstance(image, (BytesIO, BufferedReader)):
seekable: Final = cast(IO[bytes], image) # cast-ok: isinstance drops BufferedReader's type argument
return _read_seekable(seekable)
if _is_seekable_binary(image):
return _read_seekable(image)
if isinstance(image, _Readable):
raw: Final = image.read()
if isinstance(raw, str):

141
litellm/llms/xai/image_generation/transformation.py
View file

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, TypeAlias
import httpx
@ -7,18 +7,21 @@ 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.llms.xai.common_utils import XAIModelInfo, xai_image_response
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 litellm.types.utils import ImageResponse
if TYPE_CHECKING:
import tiktoken
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
_Params: TypeAlias = dict[str, object]
_Headers: TypeAlias = dict[str, object]
_SupportedParams: TypeAlias = list[OpenAIImageGenerationOptionalParams]
_SIZE_TO_ASPECT_RATIO: Final = {
"1024x1024": "1:1",
@ -35,18 +38,16 @@ _XAI_NATIVE_PARAMS: Final = frozenset({"aspect_ratio", "n"})
class XAIImageGenerationConfig(BaseImageGenerationConfig):
def get_supported_openai_params(
self, model: str
) -> list[OpenAIImageGenerationOptionalParams]: # mutable-ok: provider JSON body and base-class dict signature
def get_supported_openai_params(self, model: str) -> _SupportedParams:
return ["n", "response_format", "size", "user"]
def map_openai_params(
self,
non_default_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
optional_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
non_default_params: _Params,
optional_params: _Params,
model: str,
drop_params: bool,
) -> dict: # mutable-ok: provider JSON body and base-class dict signature
) -> _Params:
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)
@ -73,8 +74,8 @@ class XAIImageGenerationConfig(BaseImageGenerationConfig):
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
litellm_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
optional_params: _Params,
litellm_params: _Params,
stream: bool | None = None,
) -> str:
from litellm.llms.xai.oauth import XAIOAuthAuthenticator, should_use_xai_oauth
@ -89,16 +90,7 @@ class XAIImageGenerationConfig(BaseImageGenerationConfig):
return f"{base}/images/generations"
return f"{base}/v1/images/generations"
def validate_environment(
self,
headers: dict, # mutable-ok: provider JSON body and base-class dict signature
model: str,
messages: list[AllMessageValues], # mutable-ok: provider JSON body and base-class dict signature
optional_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
litellm_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
api_key: str | None = None,
api_base: str | None = None,
) -> dict: # mutable-ok: provider JSON body and base-class dict signature
def _bearer_token(self, model: str, api_key: str | None, litellm_params: _Params) -> str:
from litellm.llms.xai.oauth import (
XAIOAuthAuthenticator,
XAIOAuthError,
@ -106,39 +98,47 @@ class XAIImageGenerationConfig(BaseImageGenerationConfig):
)
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 dynamic_api_key:
return dynamic_api_key
if not should_use_xai_oauth(litellm_params):
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."
),
)
try:
return XAIOAuthAuthenticator().get_access_token()
except XAIOAuthError as exc:
raise AuthenticationError(model=model, llm_provider="xai", message=str(exc)) from exc
if "content-type" not in headers and "Content-Type" not in headers:
headers["Content-Type"] = "application/json"
return headers
def validate_environment(
self,
headers: _Headers,
model: str,
messages: list[AllMessageValues], # mutable-ok: base-class signature
optional_params: _Params,
litellm_params: _Params,
api_key: str | None = None,
api_base: str | None = None,
) -> _Headers:
has_content_type: Final = "content-type" in headers or "Content-Type" in headers
return {
**headers,
"Authorization": f"Bearer {self._bearer_token(model, api_key, litellm_params)}",
**({} if has_content_type else {"Content-Type": "application/json"}),
}
def transform_image_generation_request(
self,
model: str,
prompt: str,
optional_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
litellm_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
headers: dict, # mutable-ok: provider JSON body and base-class dict signature
) -> dict: # mutable-ok: provider JSON body and base-class dict signature
optional_params: _Params,
litellm_params: _Params,
headers: _Headers,
) -> _Params:
n: Final = optional_params.get("n")
requested_ratio: Final = optional_params.get("aspect_ratio")
aspect: Final = {"aspect_ratio": requested_ratio} if requested_ratio is not None else None # fmt: skip
@ -156,42 +156,13 @@ class XAIImageGenerationConfig(BaseImageGenerationConfig):
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: "LiteLLMLoggingObj",
request_data: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
optional_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
litellm_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
encoding: "tiktoken.Encoding | None",
request_data: _Params,
optional_params: _Params,
litellm_params: _Params,
encoding: "Tokenizer | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ImageResponse:
try:
response_data: Final = raw_response.json()
except Exception: # noqa: BLE001 # any body-parse failure becomes a provider error carrying the raw payload
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)
images: Final = xai_image_response(raw_response, "image generation")
model_response.data = list(images) # rebind-ok: the base contract fills the caller's ImageResponse
return model_response

236
litellm/llms/xai/videos/transformation.py
View file

@ -1,10 +1,11 @@
import math
import time
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, TypeAlias
import httpx
from httpx._types import RequestFiles
from pydantic import TypeAdapter
from typing_extensions import TypeIs # noqa: TID251 # narrows untyped wire payloads without a runtime conversion
import litellm
@ -15,7 +16,7 @@ from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
get_async_httpx_client,
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory lacks typed params
)
from litellm.llms.xai.common_utils import XAIModelInfo
from litellm.secret_managers.main import get_secret_str
@ -29,11 +30,31 @@ from litellm.types.videos.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_Params: TypeAlias = dict[str, object]
_Headers: TypeAlias = dict[str, object]
_SupportedParams: TypeAlias = list[str]
_EMPTY: Final[Mapping[str, object]] = {}
_JSON_PAYLOAD: Final = TypeAdapter(object)
_DEFAULT_DURATION: Final = 6
def _is_json_object(value: object) -> TypeIs[Mapping[str, object]]: # guard-ok: provider JSON objects have str keys
return isinstance(value, Mapping)
def _json_object(raw_response: httpx.Response) -> Mapping[str, object]:
payload: Final = _JSON_PAYLOAD.validate_python(raw_response.json())
return payload if _is_json_object(payload) else _EMPTY
def _optional_str(value: object) -> str | None:
return value if isinstance(value, str) else None
def _optional_int(value: object) -> int | None:
return value if isinstance(value, int) and not isinstance(value, bool) else None
_DROPPED: Final = frozenset(("seconds", "size", "input_reference", "user", "extra_headers", "model"))
_SIZE_TO_ASPECT_RATIO: Final = {
"1024x1024": "1:1",
@ -46,13 +67,6 @@ _SIZE_TO_ASPECT_RATIO: Final = {
}
def _duration_from_seconds(seconds: object) -> int:
try:
return int(seconds) if seconds is not None else 6
except (TypeError, ValueError):
return 6
def _positive_seconds(value: object) -> float | None:
if isinstance(value, bool) or not isinstance(value, (int, float, str)):
return None
@ -63,7 +77,12 @@ def _positive_seconds(value: object) -> float | None:
return seconds if math.isfinite(seconds) and seconds > 0 else None
def _requested_video_usage(request_data: Mapping[str, object]) -> dict[str, object]: # mutable-ok: VideoObject.usage
def _duration_from_seconds(seconds: object) -> int:
parsed: Final = _positive_seconds(seconds)
return int(parsed) if parsed is not None else _DEFAULT_DURATION
def _requested_video_usage(request_data: Mapping[str, object]) -> _Params:
resolution: Final = request_data.get("resolution") or request_data.get("resolution_name")
return {
key: value
@ -88,9 +107,7 @@ _STATUS_MAP: Final = {
class XAIVideoConfig(BaseVideoConfig):
def get_supported_openai_params(
self, model: str
) -> list: # mutable-ok: provider JSON body and base-class dict signature
def get_supported_openai_params(self, model: str) -> _SupportedParams:
return [
"model",
"prompt",
@ -106,14 +123,12 @@ class XAIVideoConfig(BaseVideoConfig):
video_create_optional_params: VideoCreateOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict: # mutable-ok: provider JSON body and base-class dict signature
) -> _Params:
raw: Final = video_create_optional_params
incoming: Final = dict(raw)
incoming: Final[Mapping[str, object]] = dict(raw)
size: Final = incoming.get("size")
seconds: Final = incoming.get("seconds")
use_duration: Final = "seconds" in incoming and "duration" not in incoming
duration_value: Final = _duration_from_seconds(seconds)
duration: Final = {"duration": duration_value} if use_duration else None
duration: Final = {"duration": _duration_from_seconds(incoming.get("seconds"))} if use_duration else None
mapped_ratio: Final = incoming.get("aspect_ratio") or _SIZE_TO_ASPECT_RATIO.get(str(size), "16:9")
use_ratio: Final = bool(size) and "aspect_ratio" not in incoming
ratio: Final = {"aspect_ratio": mapped_ratio} if use_ratio else None
@ -131,24 +146,21 @@ class XAIVideoConfig(BaseVideoConfig):
self,
api_base: str | None,
api_key: str | None,
litellm_params: GenericLiteLLMParams | dict[str, object] | None, # mutable-ok: base-class dict
litellm_params: Mapping[str, object],
) -> 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):
if should_use_xai_oauth(litellm_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 _optional_str(litellm_params.get("api_base"))
or get_secret_str("XAI_API_BASE")
or get_secret_str("XAI_OAUTH_API_BASE")
or XAI_API_BASE
)
return str(resolved).rstrip("/")
return resolved.rstrip("/")
def _v1_root(self, api_base: str) -> str:
base: Final = api_base.rstrip("/")
@ -156,58 +168,55 @@ class XAIVideoConfig(BaseVideoConfig):
return base
return f"{base}/v1"
def validate_environment(
self,
headers: dict, # mutable-ok: provider JSON body and base-class dict signature
model: str,
api_key: str | None = None,
litellm_params: GenericLiteLLMParams | None = None,
) -> dict: # mutable-ok: provider JSON body and base-class dict signature
def _bearer_token(self, model: str, api_key: str | None, litellm_params: GenericLiteLLMParams | None) -> str:
from litellm.llms.xai.oauth import (
XAIOAuthAuthenticator,
XAIOAuthError,
should_use_xai_oauth,
)
dumped: Final = litellm_params.model_dump() if litellm_params is not None else None
params: Final = dumped or {}
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}"
dynamic_api_key: Final = XAIModelInfo.get_api_key(
api_key or (litellm_params.api_key if litellm_params else None)
)
if dynamic_api_key:
return dynamic_api_key
if not should_use_xai_oauth(litellm_params.model_dump() if litellm_params is not None else {}):
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."
),
)
try:
return XAIOAuthAuthenticator().get_access_token()
except XAIOAuthError as exc:
raise AuthenticationError(model=model or "xai-video", llm_provider="xai", message=str(exc)) from exc
if "content-type" not in headers and "Content-Type" not in headers:
headers["Content-Type"] = "application/json"
return headers
def validate_environment(
self,
headers: _Headers,
model: str,
api_key: str | None = None,
litellm_params: GenericLiteLLMParams | None = None,
) -> _Headers:
has_content_type: Final = "content-type" in headers or "Content-Type" in headers
return {
**headers,
"Authorization": f"Bearer {self._bearer_token(model, api_key, litellm_params)}",
**({} if has_content_type else {"Content-Type": "application/json"}),
}
def get_complete_url(
self,
model: str,
api_base: str | None,
litellm_params: dict[str, object], # mutable-ok: provider JSON body and base-class dict signature
litellm_params: _Params,
) -> str:
api_key: Final = litellm_params.get("api_key") if litellm_params else None
resolved: Final = self._resolve_api_base(
api_base=api_base,
api_key=api_key if isinstance(api_key, str) else None,
api_key=_optional_str(litellm_params.get("api_key")),
litellm_params=litellm_params,
)
if not model:
@ -219,10 +228,10 @@ class XAIVideoConfig(BaseVideoConfig):
model: str,
prompt: str,
api_base: str,
video_create_optional_request_params: dict, # mutable-ok: provider JSON body and base-class dict signature
video_create_optional_request_params: _Params,
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: provider JSON body and base-class dict signature
) -> tuple[dict, RequestFiles, str]: # mutable-ok: provider JSON body and base-class dict signature
headers: _Headers,
) -> tuple[_Params, RequestFiles, str]:
copied: Final = {
key: video_create_optional_request_params[key]
for key in (
@ -237,7 +246,7 @@ class XAIVideoConfig(BaseVideoConfig):
if video_create_optional_request_params.get(key) is not None
}
prompt_body: Final = {"prompt": prompt} if prompt else None
duration_body: Final = {"duration": 6} if "duration" not in copied else None
duration_body: Final = {"duration": _DEFAULT_DURATION} if "duration" not in copied else None
return (
{
"model": XAIModelInfo.get_base_model(model) or model,
@ -255,29 +264,28 @@ class XAIVideoConfig(BaseVideoConfig):
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
custom_llm_provider: str | None = None,
request_data: dict | None = None, # mutable-ok: provider JSON body and base-class dict signature
request_data: _Params | None = None,
) -> VideoObject:
response_data: Final = raw_response.json()
request_id: Final = response_data.get("request_id") or response_data.get("id")
response_data: Final = _json_object(raw_response)
request_id: Final = _optional_str(response_data.get("request_id")) or _optional_str(response_data.get("id"))
if not request_id:
raise ValueError(f"xAI video generation response missing request_id: {response_data}")
provider_usage: Final[object] = response_data.get("usage")
requested_usage: Final = _requested_video_usage(request_data or {})
provider_usage: Final = response_data.get("usage")
requested_usage: Final = _requested_video_usage(request_data or _EMPTY)
duration_seconds: Final = requested_usage.get("duration_seconds")
video_obj: Final = VideoObject(
id=str(request_id),
return VideoObject(
id=encode_video_id_with_provider(request_id, custom_llm_provider, model)
if custom_llm_provider
else request_id,
object="video",
status="processing",
created_at=int(time.time()),
model=XAIModelInfo.get_base_model(model) or model,
progress=0,
seconds=f"{duration_seconds:g}" if isinstance(duration_seconds, float) else None,
usage={**requested_usage, **(provider_usage if _is_json_object(provider_usage) else _EMPTY)},
)
if custom_llm_provider:
video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model)
video_obj.usage = {**requested_usage, **(provider_usage if _is_json_object(provider_usage) else {})}
return video_obj
def _video_resource_url(self, api_base: str, video_id: str) -> str:
encoded_video_id: Final = encode_url_path_segment(
@ -291,8 +299,8 @@ class XAIVideoConfig(BaseVideoConfig):
video_id: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: provider JSON body and base-class dict signature
) -> tuple[str, dict]: # mutable-ok: provider JSON body and base-class dict signature
headers: _Headers,
) -> tuple[str, _Params]:
return self._video_resource_url(api_base, video_id), {}
def transform_video_status_retrieve_response(
@ -302,34 +310,34 @@ class XAIVideoConfig(BaseVideoConfig):
custom_llm_provider: str | None = None,
client: HTTPHandler | None = None,
) -> VideoObject:
response_data: Final = raw_response.json()
response_data: Final = _json_object(raw_response)
status_raw: Final = str(response_data.get("status") or "processing").lower()
status: Final = _STATUS_MAP.get(status_raw, status_raw)
video_body: Final[object] = response_data.get("video")
video_meta: Final[Mapping[str, object]] = video_body if _is_json_object(video_body) else {}
url_value: Final = video_meta.get("url")
video_url: Final = url_value if isinstance(url_value, str) else None
video_body: Final = response_data.get("video")
video_meta: Final = video_body if _is_json_object(video_body) else _EMPTY
video_url: Final = _optional_str(video_meta.get("url"))
duration: Final = video_meta.get("duration")
seconds: Final = str(duration) if duration is not None else None
model: Final = _optional_str(response_data.get("model"))
provider_usage: Final = response_data.get("usage")
request_id: Final = (
response_data.get("request_id")
or response_data.get("id")
_optional_str(response_data.get("request_id"))
or _optional_str(response_data.get("id"))
or (video_url.split("/")[-1].replace(".mp4", "") if video_url else "unknown")
)
video_obj: Final = VideoObject(
id=str(request_id),
id=encode_video_id_with_provider(request_id, custom_llm_provider, model)
if custom_llm_provider and request_id != "unknown"
else request_id,
object="video",
status=status,
created_at=response_data.get("created_at") or int(time.time()),
created_at=_optional_int(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 {},
model=model,
progress=_optional_int(response_data.get("progress")),
seconds=str(duration) if duration is not None else None,
usage=dict(provider_usage) if _is_json_object(provider_usage) 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(
@ -337,21 +345,21 @@ class XAIVideoConfig(BaseVideoConfig):
video_id: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: provider JSON body and base-class dict signature
headers: _Headers,
variant: str | None = None,
) -> tuple[str, dict]: # mutable-ok: provider JSON body and base-class dict signature
) -> tuple[str, _Params]:
return self._video_resource_url(api_base, video_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):
payload: Final = _JSON_PAYLOAD.validate_python(raw_response.json())
if not _is_json_object(payload):
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:
video_body: Final = payload.get("video")
url: Final = _optional_str(video_body.get("url")) if _is_json_object(video_body) else None
if url:
return url
raise ValueError(f"xAI video not ready for download (status={payload.get('status')}): {payload}")
@ -359,11 +367,12 @@ class XAIVideoConfig(BaseVideoConfig):
self,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
client: HTTPHandler | None = None,
) -> bytes:
url: Final = self._video_cdn_url(raw_response)
if url is None:
return raw_response.content
video_response: Final = safe_get(litellm.module_level_client, url)
video_response: Final = safe_get(client or litellm.module_level_client, url)
video_response.raise_for_status()
return video_response.content
@ -371,11 +380,12 @@ class XAIVideoConfig(BaseVideoConfig):
self,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
client: AsyncHTTPHandler | None = None,
) -> 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(
async_httpx_client: Final[AsyncHTTPHandler] = client or get_async_httpx_client(
llm_provider=litellm.LlmProviders.XAI,
)
video_response: Final = await async_safe_get(async_httpx_client, url)
@ -388,9 +398,9 @@ class XAIVideoConfig(BaseVideoConfig):
prompt: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: provider JSON body and base-class dict signature
extra_body: dict[str, object] | None = None, # mutable-ok: provider JSON body and base-class dict signature
) -> tuple[str, dict]: # mutable-ok: provider JSON body and base-class dict signature
headers: _Headers,
extra_body: Mapping[str, object] | None = None,
) -> tuple[str, _Params]:
raise NotImplementedError("Video remix is not supported by xAI Imagine API")
def transform_video_remix_response(
@ -405,12 +415,12 @@ class XAIVideoConfig(BaseVideoConfig):
self,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: provider JSON body and base-class dict signature
headers: _Headers,
after: str | None = None,
limit: int | None = None,
order: str | None = None,
extra_query: dict[str, object] | None = None, # mutable-ok: provider JSON body and base-class dict signature
) -> tuple[str, dict]: # mutable-ok: provider JSON body and base-class dict signature
extra_query: Mapping[str, object] | None = None,
) -> tuple[str, _Params]:
raise NotImplementedError("Video listing is not supported by xAI Imagine API")
def transform_video_list_response(
@ -418,7 +428,7 @@ class XAIVideoConfig(BaseVideoConfig):
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
custom_llm_provider: str | None = None,
) -> dict[str, str]: # mutable-ok: provider JSON body and base-class dict signature
) -> dict[str, str]: # mutable-ok: base-class signature
raise NotImplementedError("Video listing is not supported by xAI Imagine API")
def transform_video_delete_request(
@ -426,8 +436,8 @@ class XAIVideoConfig(BaseVideoConfig):
video_id: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: provider JSON body and base-class dict signature
) -> tuple[str, dict]: # mutable-ok: provider JSON body and base-class dict signature
headers: _Headers,
) -> tuple[str, _Params]:
raise NotImplementedError("Video delete is not supported by xAI Imagine API")
def transform_video_delete_response(

15
litellm/proxy/video_endpoints/ownership.py
View file

@ -12,7 +12,7 @@ checks are skipped and videos remain reachable by any caller allowed on the rout
import json
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, Protocol
from fastapi import HTTPException
from typing_extensions import TypeIs # noqa: TID251 # narrows untyped wire payloads without a runtime conversion
@ -35,9 +35,20 @@ if TYPE_CHECKING:
VIDEO_OBJECT_PURPOSE: Final = "video"
# Status is polled; a short-lived cache of known owners keeps polling off the DB.
# Only positive answers are cached, so a missing row is always re-checked.
_VIDEO_OWNER_CACHE: Final = InMemoryCache(max_size_in_memory=10000, default_ttl=60)
class _OwnerCache(Protocol):
def get_cache(self, key: str) -> object: ...
def set_cache(self, key: str, value: str) -> None: ...
def _owner_cache() -> _OwnerCache:
return InMemoryCache(max_size_in_memory=10000, default_ttl=60)
_VIDEO_OWNER_CACHE: Final = _owner_cache()
def _video_model_object_id(video_id: str) -> str:

59
litellm/videos/main.py
View file

@ -3,7 +3,7 @@ import contextvars
import json
from collections.abc import Coroutine
from functools import partial
from typing import Final, Literal, cast, overload # noqa: TID251 # untyped kwargs need cast
from typing import Final, Literal, overload
from httpx._types import FileContent
@ -332,28 +332,27 @@ def video_content(
```
"""
local_vars: Final = locals()
provider: Final = _provider_for_video_id(
video_id,
custom_llm_provider,
kwargs.get("model"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # kwargs is untyped
)
try:
litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj")
litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None)
_is_async: Final = kwargs.pop("async_call", False) is True
custom_llm_provider = _provider_for_video_id(
video_id,
custom_llm_provider,
cast(object, kwargs.get("model")), # cast-ok: kwargs is untyped
)
# get llm provider logic
litellm_params: Final = GenericLiteLLMParams(**kwargs)
# get provider config
video_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config(
model=None,
provider=litellm.LlmProviders(custom_llm_provider),
provider=litellm.LlmProviders(provider),
)
if video_provider_config is None:
raise ValueError(f"video support download is not supported for {custom_llm_provider}")
raise ValueError(f"video support download is not supported for {provider}")
local_vars.update(kwargs)
# For video content download, we don't need complex optional parameter handling
@ -372,14 +371,14 @@ def video_content(
"litellm_call_id": litellm_call_id,
**video_content_request_params,
},
custom_llm_provider=custom_llm_provider,
custom_llm_provider=provider,
)
# Call the handler with _is_async flag instead of directly calling the async handler
return base_llm_http_handler.video_content_handler(
video_id=video_id,
video_content_provider_config=video_provider_config,
custom_llm_provider=custom_llm_provider,
custom_llm_provider=provider,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
timeout=timeout or DEFAULT_REQUEST_TIMEOUT,
@ -392,7 +391,7 @@ def video_content(
except Exception as e:
raise litellm.exception_type(
model="",
custom_llm_provider=custom_llm_provider,
custom_llm_provider=provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
@ -429,21 +428,20 @@ async def avideo_content(
- `bytes`: The raw video content as bytes
"""
local_vars: Final = locals()
provider: Final = _provider_for_video_id(
video_id,
custom_llm_provider,
kwargs.get("model"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # kwargs is untyped
)
try:
loop: Final = asyncio.get_event_loop()
kwargs["async_call"] = True
custom_llm_provider = _provider_for_video_id(
video_id,
custom_llm_provider,
cast(object, kwargs.get("model")), # cast-ok: kwargs is untyped
)
func: Final = partial(
video_content,
video_id=video_id,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
custom_llm_provider=provider,
variant=variant,
extra_headers=extra_headers,
extra_query=extra_query,
@ -464,7 +462,7 @@ async def avideo_content(
except Exception as e:
raise litellm.exception_type(
model="",
custom_llm_provider=custom_llm_provider,
custom_llm_provider=provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
@ -1027,6 +1025,11 @@ def video_status(
```
"""
local_vars: Final = locals()
provider: Final = _provider_for_video_id(
video_id,
custom_llm_provider,
kwargs.get("model"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # kwargs is untyped
)
try:
litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj")
litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None)
@ -1041,23 +1044,17 @@ def video_status(
response: Final = VideoObject(**mock_response)
return response
custom_llm_provider = _provider_for_video_id(
video_id,
custom_llm_provider,
cast(object, kwargs.get("model")), # cast-ok: kwargs is untyped
)
# get llm provider logic
litellm_params: Final = GenericLiteLLMParams(**kwargs)
# get provider config
video_status_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config(
model=None,
provider=litellm.LlmProviders(custom_llm_provider),
provider=litellm.LlmProviders(provider),
)
if video_status_provider_config is None:
raise ValueError(f"video status is not supported for {custom_llm_provider}")
raise ValueError(f"video status is not supported for {provider}")
local_vars.update(kwargs)
# For video status, we need the video_id
@ -1075,7 +1072,7 @@ def video_status(
"litellm_call_id": litellm_call_id,
**video_status_request_params,
},
custom_llm_provider=custom_llm_provider,
custom_llm_provider=provider,
)
# Set the correct call type for video status
@ -1085,7 +1082,7 @@ def video_status(
return base_llm_http_handler.video_status_handler(
video_id=video_id,
video_status_provider_config=video_status_provider_config,
custom_llm_provider=custom_llm_provider,
custom_llm_provider=provider,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
extra_headers=extra_headers,
@ -1098,7 +1095,7 @@ def video_status(
except Exception as e:
raise litellm.exception_type(
model="",
custom_llm_provider=custom_llm_provider,
custom_llm_provider=provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,

5
tests/unit/llms/xai/test_xai_image_edit.py
View file

@ -5,6 +5,7 @@ from unittest.mock import MagicMock, patch
import httpx
import pytest
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
@ -309,11 +310,11 @@ def test_transform_rejects_unsupported_image_type():
def test_transform_response_raises_provider_error_without_images(raw):
from litellm.llms.xai.image_edit.transformation import XAIImageEditConfig
with pytest.raises(Exception) as exc_info:
with pytest.raises(BaseLLMException) as exc_info:
XAIImageEditConfig().transform_image_edit_response(
model="grok-imagine-image",
raw_response=raw,
logging_obj=MagicMock(),
)
assert getattr(exc_info.value, "status_code", None) == raw.status_code
assert exc_info.value.status_code == raw.status_code

5
tests/unit/llms/xai/test_xai_image_generation.py
View file

@ -3,6 +3,7 @@ from unittest.mock import MagicMock, patch
import httpx
import pytest
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.xai.image_generation.transformation import XAIImageGenerationConfig
from litellm.types.utils import ImageResponse, LlmProviders
from litellm.utils import ProviderConfigManager
@ -180,7 +181,7 @@ def test_get_complete_url_appends_v1_to_a_bare_custom_base():
],
)
def test_transform_response_raises_provider_error_without_images(raw):
with pytest.raises(Exception) as exc_info:
with pytest.raises(BaseLLMException) as exc_info:
XAIImageGenerationConfig().transform_image_generation_response(
model="grok-imagine-image",
raw_response=raw,
@ -192,4 +193,4 @@ def test_transform_response_raises_provider_error_without_images(raw):
encoding=None,
)
assert getattr(exc_info.value, "status_code", None) == raw.status_code
assert exc_info.value.status_code == raw.status_code

20
tests/unit/llms/xai/test_xai_video_generation.py
View file

@ -279,6 +279,26 @@ def test_sync_video_status_through_the_http_handler():
assert client.get.call_args.kwargs["url"] == "https://api.x.ai/v1/videos/req-123"
def test_sync_video_content_downloads_the_cdn_file_with_the_handler_client():
client = MagicMock(spec=HTTPHandler)
client.get.return_value = httpx.Response(
200,
headers={"content-type": "application/json"},
json={"status": "done", "video": {"url": "https://vidgen.x.ai/x.mp4"}},
request=httpx.Request("GET", "https://api.x.ai/v1/videos/req-123"),
)
cdn_response = httpx.Response(200, content=b"mp4-bytes", request=httpx.Request("GET", "https://vidgen.x.ai/x.mp4"))
with patch("litellm.llms.xai.videos.transformation.safe_get", return_value=cdn_response) as safe_get:
body = litellm.video_content(
video_id="req-123", model="xai/grok-imagine-video", api_key="test-key", client=client
)
assert body == b"mp4-bytes"
assert client.get.call_args.kwargs["url"] == "https://api.x.ai/v1/videos/req-123"
safe_get.assert_called_once_with(client, "https://vidgen.x.ai/x.mp4")
def test_validate_environment_requires_credentials():
with pytest.raises(Exception, match="Missing xAI credentials"):
XAIVideoConfig().validate_environment(