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(