From 9c3a560f322e0dcda5fa6095200d0cf11f70e0bb Mon Sep 17 00:00:00 2001 From: hx <1367557521@qq.com> Date: Thu, 15 Oct 2026 20:39:11 +0800 Subject: [PATCH] fix(xai): type imagine transforms and accept the handler client on video downloads 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 --- litellm/llms/xai/common_utils.py | 43 +++- litellm/llms/xai/image_edit/transformation.py | 170 ++++++------- .../xai/image_generation/transformation.py | 141 +++++------ litellm/llms/xai/videos/transformation.py | 236 +++++++++--------- litellm/proxy/video_endpoints/ownership.py | 15 +- litellm/videos/main.py | 59 +++-- tests/unit/llms/xai/test_xai_image_edit.py | 5 +- .../llms/xai/test_xai_image_generation.py | 5 +- .../llms/xai/test_xai_video_generation.py | 20 ++ 9 files changed, 356 insertions(+), 338 deletions(-) diff --git a/litellm/llms/xai/common_utils.py b/litellm/llms/xai/common_utils.py index 06e3f32a767..4ffb92f2e26 100644 --- a/litellm/llms/xai/common_utils.py +++ b/litellm/llms/xai/common_utils.py @@ -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] diff --git a/litellm/llms/xai/image_edit/transformation.py b/litellm/llms/xai/image_edit/transformation.py index 86cb235d6d9..736aed0dbd3 100644 --- a/litellm/llms/xai/image_edit/transformation.py +++ b/litellm/llms/xai/image_edit/transformation.py @@ -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): diff --git a/litellm/llms/xai/image_generation/transformation.py b/litellm/llms/xai/image_generation/transformation.py index 3ffbb403751..e3ab3aa09d1 100644 --- a/litellm/llms/xai/image_generation/transformation.py +++ b/litellm/llms/xai/image_generation/transformation.py @@ -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 diff --git a/litellm/llms/xai/videos/transformation.py b/litellm/llms/xai/videos/transformation.py index 77ad2c043b4..c88692dfa34 100644 --- a/litellm/llms/xai/videos/transformation.py +++ b/litellm/llms/xai/videos/transformation.py @@ -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( diff --git a/litellm/proxy/video_endpoints/ownership.py b/litellm/proxy/video_endpoints/ownership.py index 7b3c82ada29..fc2624be9e8 100644 --- a/litellm/proxy/video_endpoints/ownership.py +++ b/litellm/proxy/video_endpoints/ownership.py @@ -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: diff --git a/litellm/videos/main.py b/litellm/videos/main.py index 40d2569c651..bbb46fca61a 100644 --- a/litellm/videos/main.py +++ b/litellm/videos/main.py @@ -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, diff --git a/tests/unit/llms/xai/test_xai_image_edit.py b/tests/unit/llms/xai/test_xai_image_edit.py index e2da75c5d4f..b2d6b4fe2a2 100644 --- a/tests/unit/llms/xai/test_xai_image_edit.py +++ b/tests/unit/llms/xai/test_xai_image_edit.py @@ -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 diff --git a/tests/unit/llms/xai/test_xai_image_generation.py b/tests/unit/llms/xai/test_xai_image_generation.py index e061f47fb8f..06116040957 100644 --- a/tests/unit/llms/xai/test_xai_image_generation.py +++ b/tests/unit/llms/xai/test_xai_image_generation.py @@ -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 diff --git a/tests/unit/llms/xai/test_xai_video_generation.py b/tests/unit/llms/xai/test_xai_video_generation.py index d4321a416ea..43f3242a6a9 100644 --- a/tests/unit/llms/xai/test_xai_video_generation.py +++ b/tests/unit/llms/xai/test_xai_video_generation.py @@ -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(