diff --git a/litellm/constants.py b/litellm/constants.py index 066b682362d..b4dcac764f9 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -163,6 +163,8 @@ DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = get_env_float("DEFAULT_COO DEFAULT_REDIS_BATCH_CACHE_EXPIRY_SECONDS: Final = 10 DEFAULT_REPLICATE_POLLING_RETRIES: Final = get_env_int("DEFAULT_REPLICATE_POLLING_RETRIES", 5) DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_int("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1) +DEFAULT_REPLICATE_SYNC_POLLING_RETRIES: Final = get_env_int("DEFAULT_REPLICATE_SYNC_POLLING_RETRIES", 120) +REPLICATE_MAX_OUTPUT_DOWNLOAD_BYTES: Final = get_env_int("REPLICATE_MAX_OUTPUT_DOWNLOAD_BYTES", 512 * 1024 * 1024) DEFAULT_IMAGE_TOKEN_COUNT: Final = get_env_int("DEFAULT_IMAGE_TOKEN_COUNT", 250) HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0 HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range( diff --git a/litellm/images/main.py b/litellm/images/main.py index b2da69bb452..e1ecfc35a02 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -465,6 +465,7 @@ def image_generation( litellm.LlmProviders.AIML, litellm.LlmProviders.GEMINI, litellm.LlmProviders.FAL_AI, + litellm.LlmProviders.REPLICATE, litellm.LlmProviders.STABILITY, litellm.LlmProviders.RUNWAYML, litellm.LlmProviders.VERTEX_AI, @@ -496,6 +497,7 @@ def image_generation( logging_obj=litellm_logging_obj, timeout=timeout, client=client, + _is_async=aimg_generation is True and custom_llm_provider == litellm.LlmProviders.REPLICATE, ) elif custom_llm_provider == "black_forest_labs": # Route to BFL-specific handler (polling required) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index f4ba3f43712..cbee7c818c5 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -318,6 +318,7 @@ class _BodyPrefixReader: headers: Mapping[str, str] | None = None, follow_redirects: bool = False, timeout: float | httpx.Timeout | None = None, + max_response_bytes: int | None = None, ) -> Response: with self.client.stream( "GET", @@ -342,6 +343,7 @@ class _ImageBodyReader: headers: Mapping[str, str] | None = None, follow_redirects: bool = False, timeout: float | httpx.Timeout | None = None, + max_response_bytes: int | None = None, ) -> Response: with self.client.stream( "GET", diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index eae8ac04e8e..490d96746bd 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -51,6 +51,7 @@ class _UrlFetcher(Protocol): headers: dict[str, str] | None = None, follow_redirects: bool = False, timeout: float | httpx.Timeout | None = None, + max_response_bytes: int | None = None, ) -> httpx.Response: ... @@ -80,10 +81,11 @@ class _FetchOptions(TypedDict, total=False): headers: ReadOnly[dict[str, str]] follow_redirects: ReadOnly[bool] timeout: ReadOnly[float | httpx.Timeout | None] + max_response_bytes: ReadOnly[int | None] class _AsyncFetchOptions(_FetchOptions, total=False): - max_response_bytes: ReadOnly[int | None] + pass # Cloud-internal IPs that ``ip.is_global`` can report as public. Everything diff --git a/litellm/llms/base_llm/image_generation/transformation.py b/litellm/llms/base_llm/image_generation/transformation.py index 7ac440c6f48..28bf3d407f3 100644 --- a/litellm/llms/base_llm/image_generation/transformation.py +++ b/litellm/llms/base_llm/image_generation/transformation.py @@ -78,7 +78,7 @@ class BaseImageGenerationConfig(ABC): optional_params: dict, litellm_params: dict, headers: dict, - ) -> dict: + ) -> dict[str, object]: # mutable-ok: inherited contract raise NotImplementedError( "ImageVariationConfig implementa 'transform_request_image_variation' for image variation models" ) @@ -100,6 +100,32 @@ class BaseImageGenerationConfig(ABC): "ImageVariationConfig implements 'transform_response_image_variation' for image variation models" ) + async def async_transform_image_generation_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ImageResponse, + logging_obj: LiteLLMLoggingObj, + request_data: dict[str, object], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + encoding: "Tokenizer | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ImageResponse: + return self.transform_image_generation_response( # pyright: ignore[reportUnknownMemberType] # legacy sync hook has untyped dict params + model=model, + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=request_data, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + api_key=api_key, + json_mode=json_mode, + ) + def use_multipart_form_data(self) -> bool: """ Returns True if this provider requires multipart/form-data instead of JSON. diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index b290a7abf3d..24ee7958e45 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1671,6 +1671,7 @@ class HTTPHandler: follow_redirects: bool | None = None, timeout: float | httpx.Timeout | None = None, stream: bool = False, + max_response_bytes: int | None = None, ): # Set follow_redirects to UseClientDefault if None _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT @@ -1681,6 +1682,16 @@ class HTTPHandler: url, params=query, headers=headers, follow_redirects=follow_redirects, timeout=timeout ) + if max_response_bytes is not None: + return self._get_with_response_limit( + url, + params=query, + headers=headers, + max_bytes=max_response_bytes, + follow_redirects=self.client.follow_redirects if follow_redirects is None else follow_redirects, + timeout=self.client.timeout if timeout is None else httpx.Timeout(timeout), + ) + response: Final = self.client.get( url, params=query, @@ -1691,6 +1702,67 @@ class HTTPHandler: return response + def _get_with_response_limit( + self, + url: str, + *, + params: httpx.QueryParams, + headers: Mapping[str, str] | None, + timeout: httpx.Timeout, + max_bytes: int, + follow_redirects: bool, + ) -> httpx.Response: + request: Final = self.client.build_request( + "GET", + url, + headers=MappingProxyType({**(headers or {}), "accept-encoding": "identity"}), + params=params, + timeout=timeout, + ) + response: Final = self.client.send(request, stream=True, follow_redirects=False) + return self._read_with_response_limit(response, max_bytes=max_bytes, follow_redirects=follow_redirects) + + def _next_redirect_hop(self, hop: httpx.Response) -> httpx.Response: + if hop.next_request is None: + return hop + hop.close() + return self.client.send(hop.next_request, auth=None, stream=True, follow_redirects=False) + + def _last_redirect_hop(self, response: httpx.Response, *, follow_redirects: bool) -> httpx.Response: + hop = response # rebind-ok: walks the redirect chain without recursion + for _ in range(10): + if hop.next_request is None or not follow_redirects: + return hop + hop = self._next_redirect_hop(hop) + if hop.next_request is None or not follow_redirects: + return hop + hop.close() + raise ValueError("Too many redirects") + + def _read_with_response_limit( + self, response: httpx.Response, *, max_bytes: int, follow_redirects: bool + ) -> httpx.Response: + final: Final = self._last_redirect_hop(response, follow_redirects=follow_redirects) + try: + if final.is_redirect or final.is_error: + return httpx.Response(final.status_code, headers=final.headers, request=final.request) + content_encoding: Final = header_value(final.headers, "content-encoding") + if ("identity" if content_encoding is None else content_encoding).lower() != "identity": + raise HTTPResponseLimitError("Response size limits require an uncompressed response") + content_length: Final = header_value(final.headers, "content-length") + if int("0" if content_length is None else content_length) > max_bytes: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + with BytesIO() as body: + for chunk in final.iter_bytes(chunk_size=65536): + if body.tell() + len(chunk) > max_bytes: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + body.write(chunk) + return httpx.Response( + final.status_code, headers=final.headers, content=body.getvalue(), request=final.request + ) + finally: + final.close() + def _send_streaming_get( self, url: str, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 68b13c6172f..c34b696bb5d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -7881,14 +7881,16 @@ class BaseLLMHTTPHandler: provider_config=image_generation_provider_config, ) - model_response: Final[ImageResponse] = image_generation_provider_config.transform_image_generation_response( + model_response: Final[ + ImageResponse + ] = await image_generation_provider_config.async_transform_image_generation_response( model=model, raw_response=response, model_response=litellm.ImageResponse(), logging_obj=logging_obj, request_data=data, - optional_params=image_generation_optional_request_params, - litellm_params=dict(litellm_params), + optional_params=image_generation_optional_request_params, # pyright: ignore[reportUnknownArgumentType] # legacy handler params are untyped dicts + litellm_params=dict(litellm_params), # pyright: ignore[reportUnknownArgumentType] # legacy handler params are untyped dicts encoding=None, ) set_provider_response_headers_in_hidden_params(model_response, response.headers) diff --git a/litellm/llms/replicate/common_utils.py b/litellm/llms/replicate/common_utils.py index e1fa3828379..4626babe14c 100644 --- a/litellm/llms/replicate/common_utils.py +++ b/litellm/llms/replicate/common_utils.py @@ -1,6 +1,38 @@ -import httpx +import asyncio +import time +from collections import deque +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final +import httpx +from pydantic import ConfigDict, TypeAdapter + +import litellm +from litellm.constants import ( + DEFAULT_REPLICATE_POLLING_DELAY_SECONDS, + DEFAULT_REPLICATE_SYNC_POLLING_RETRIES, + REPLICATE_MAX_OUTPUT_DOWNLOAD_BYTES, +) +from litellm.litellm_core_utils.url_utils import ( + assert_same_origin, + async_safe_get, + encode_url_path_segment, + ensure_api_route, + join_api_base, + safe_get, +) from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler, HTTPResponseLimitError +from litellm.secret_managers.main import get_secret_str + +REPLICATE_DEFAULT_API_BASE: Final[str] = "https://api.replicate.com/v1" +REPLICATE_TERMINAL_STATUSES: Final[frozenset[str]] = frozenset({"succeeded", "failed", "canceled"}) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_REQUEST_HEADERS_NOT_FORWARDED: Final[frozenset[str]] = frozenset( + {"host", "content-length", "content-type", "transfer-encoding", "prefer"} +) class ReplicateError(BaseLLMException): @@ -8,6 +40,231 @@ class ReplicateError(BaseLLMException): self, status_code: int, message: str, - headers: dict | httpx.Headers | None, + headers: dict[str, str] | httpx.Headers | None, ): super().__init__(status_code=status_code, message=message, headers=headers) + + +@dataclass(frozen=True, slots=True) +class ReplicatePredictionRequest: + url: str + body: Mapping[str, object] + + +def replicate_api_base(api_base: str | None) -> str: + return ensure_api_route(api_base or get_secret_str("REPLICATE_API_BASE") or REPLICATE_DEFAULT_API_BASE, "v1") + + +def replicate_auth_headers( + headers: Mapping[str, str], api_key: str | None +) -> dict[str, str]: # mutable-ok: base class contracts return a plain dict + resolved_key: Final = ( + api_key or litellm.replicate_key or get_secret_str("REPLICATE_API_KEY") or get_secret_str("REPLICATE_API_TOKEN") + ) + if not resolved_key: + raise ReplicateError( + status_code=401, + message="Replicate API key is required. Set REPLICATE_API_KEY or pass api_key.", + headers=None, + ) + return {**headers, "Authorization": f"Bearer {resolved_key}", "Content-Type": "application/json"} + + +def replicate_model_path(model: str) -> str: + segments: Final = tuple(model.split("/")) + if any(segment in ("", ".", "..") for segment in segments): + raise ReplicateError(status_code=400, message=f"Invalid Replicate model path: {model!r}", headers={}) + return "/".join(encode_url_path_segment(segment, field_name="model") for segment in segments) + + +def replicate_prediction_request( + api_base: str, model: str, prediction_input: Mapping[str, object] +) -> ReplicatePredictionRequest: + if model.startswith("deployments/"): + return ReplicatePredictionRequest( + url=join_api_base(api_base, f"{replicate_model_path(model)}/predictions"), + body=MappingProxyType({"input": prediction_input}), + ) + if ":" in model: + model_path, _, version = model.partition(":") + if not version: + raise ReplicateError(status_code=400, message=f"Invalid Replicate model version: {model!r}", headers={}) + return ReplicatePredictionRequest( + url=join_api_base( + api_base, + f"models/{replicate_model_path(model_path)}/versions/" + f"{encode_url_path_segment(version, field_name='version')}/predictions", + ), + body=MappingProxyType({"input": prediction_input}), + ) + return ReplicatePredictionRequest( + url=join_api_base(api_base, f"models/{replicate_model_path(model)}/predictions"), + body=MappingProxyType({"input": prediction_input}), + ) + + +def replicate_prediction_url(api_base: str, prediction_id: str) -> str: + return join_api_base(api_base, f"predictions/{encode_url_path_segment(prediction_id, field_name='prediction id')}") + + +def replicate_model_owner(model: str) -> str: + return model.removeprefix("deployments/").split(":", 1)[0].split("/", 1)[0] + + +def replicate_prediction(raw_response: httpx.Response) -> Mapping[str, object]: + try: + return _JSON_OBJECT.validate_python(raw_response.json()) + except ValueError as error: + raise ReplicateError( + status_code=raw_response.status_code, + message=f"Replicate returned a non-JSON response: {raw_response.text[:500]}", + headers=raw_response.headers, + ) from error + + +def replicate_prediction_string(prediction: Mapping[str, object], key: str) -> str | None: + value: Final = prediction.get(key) + return value if isinstance(value, str) and value else None + + +def replicate_prediction_input(prediction: Mapping[str, object]) -> Mapping[str, object]: + raw_input: Final = prediction.get("input") + return _JSON_OBJECT.validate_python(raw_input) if isinstance(raw_input, Mapping) else MappingProxyType({}) + + +def replicate_prediction_metrics(prediction: Mapping[str, object]) -> Mapping[str, object]: + raw_metrics: Final = prediction.get("metrics") + return _JSON_OBJECT.validate_python(raw_metrics) if isinstance(raw_metrics, Mapping) else MappingProxyType({}) + + +def replicate_positive_int(value: object) -> int | None: + return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None + + +def _output_leaves(output: object) -> Iterator[object]: + pending: Final = deque([output]) + while pending: + if isinstance((item := pending.popleft()), Mapping): + pending.extendleft(reversed(tuple(_JSON_OBJECT.validate_python(item).values()))) + elif isinstance(item, (list, tuple)): + pending.extendleft(reversed(TypeAdapter(tuple[object, ...]).validate_python(item))) + else: + yield item + + +def replicate_output_urls(output: object) -> tuple[str, ...]: + return tuple( + leaf for leaf in _output_leaves(output) if isinstance(leaf, str) and leaf.startswith(("http://", "https://")) + ) + + +def replicate_prediction_failure(prediction: Mapping[str, object]) -> str | None: + status: Final = replicate_prediction_string(prediction, "status") + error: Final = replicate_prediction_string(prediction, "error") + if status in ("failed", "canceled"): + return error or f"Replicate prediction {status}" + return None + + +def replicate_prediction_is_done(prediction: Mapping[str, object]) -> bool: + return replicate_prediction_string(prediction, "status") in REPLICATE_TERMINAL_STATUSES + + +def replicate_poll_headers( + request: httpx.Request, +) -> dict[str, str]: # mutable-ok: HTTP handler clients take a plain dict + return {key: value for key, value in request.headers.items() if key.lower() not in _REQUEST_HEADERS_NOT_FORWARDED} + + +def _poll_url(prediction: Mapping[str, object], raw_response: httpx.Response) -> str: + urls: Final = prediction.get("urls") + poll_url: Final = ( + replicate_prediction_string(_JSON_OBJECT.validate_python(urls), "get") if isinstance(urls, Mapping) else None + ) + if poll_url is not None: + assert_same_origin(poll_url, str(raw_response.request.url)) + return poll_url + prediction_id: Final = replicate_prediction_string(prediction, "id") + if prediction_id is None: + raise ReplicateError( + status_code=raw_response.status_code, + message="Replicate prediction has no id to poll", + headers=raw_response.headers, + ) + api_base: Final = str(raw_response.request.url).split("/predictions", 1)[0].rsplit("/models/", 1)[0] + return replicate_prediction_url(api_base.rsplit("/deployments/", 1)[0], prediction_id) + + +def _finished_prediction(raw_response: httpx.Response) -> Mapping[str, object] | None: + if raw_response.status_code >= 400: + raise ReplicateError( + status_code=raw_response.status_code, message=raw_response.text, headers=raw_response.headers + ) + prediction: Final = replicate_prediction(raw_response) + return prediction if replicate_prediction_is_done(prediction) else None + + +def wait_for_replicate_prediction( + raw_response: httpx.Response, + client: HTTPHandler, +) -> Mapping[str, object]: + prediction: Final = replicate_prediction(raw_response) + if replicate_prediction_is_done(prediction): + return prediction + poll_url: Final = _poll_url(prediction, raw_response) + headers: Final = replicate_poll_headers(raw_response.request) + for _ in range(DEFAULT_REPLICATE_SYNC_POLLING_RETRIES): + time.sleep(DEFAULT_REPLICATE_POLLING_DELAY_SECONDS) + if (polled := _finished_prediction(client.get(url=poll_url, headers=headers))) is not None: # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped + return polled + raise ReplicateError( + status_code=504, + message=f"Replicate prediction did not finish after {DEFAULT_REPLICATE_SYNC_POLLING_RETRIES} polls", + headers=raw_response.headers, + ) + + +async def async_wait_for_replicate_prediction( + raw_response: httpx.Response, + client: AsyncHTTPHandler, +) -> Mapping[str, object]: + prediction: Final = replicate_prediction(raw_response) + if replicate_prediction_is_done(prediction): + return prediction + poll_url: Final = _poll_url(prediction, raw_response) + headers: Final = replicate_poll_headers(raw_response.request) + for _ in range(DEFAULT_REPLICATE_SYNC_POLLING_RETRIES): + await asyncio.sleep(DEFAULT_REPLICATE_POLLING_DELAY_SECONDS) + if (polled := _finished_prediction(await client.get(url=poll_url, headers=headers))) is not None: # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped + return polled + raise ReplicateError( + status_code=504, + message=f"Replicate prediction did not finish after {DEFAULT_REPLICATE_SYNC_POLLING_RETRIES} polls", + headers=raw_response.headers, + ) + + +def _download_too_large(url: str, error: HTTPResponseLimitError) -> ReplicateError: + return ReplicateError( + status_code=502, + message=f"Replicate output at {url} exceeds REPLICATE_MAX_OUTPUT_DOWNLOAD_BYTES={REPLICATE_MAX_OUTPUT_DOWNLOAD_BYTES}: {error}", + headers={}, + ) + + +def replicate_download(client: HTTPHandler, url: str) -> httpx.Response: + try: + response: Final = safe_get(client, url, max_response_bytes=REPLICATE_MAX_OUTPUT_DOWNLOAD_BYTES) + except HTTPResponseLimitError as error: + raise _download_too_large(url, error) from error + response.raise_for_status() + return response + + +async def async_replicate_download(client: AsyncHTTPHandler, url: str) -> httpx.Response: + try: + response: Final = await async_safe_get(client, url, max_response_bytes=REPLICATE_MAX_OUTPUT_DOWNLOAD_BYTES) + except HTTPResponseLimitError as error: + raise _download_too_large(url, error) from error + response.raise_for_status() + return response diff --git a/litellm/llms/replicate/image_generation/transformation.py b/litellm/llms/replicate/image_generation/transformation.py new file mode 100644 index 00000000000..b80a5f6811d --- /dev/null +++ b/litellm/llms/replicate/image_generation/transformation.py @@ -0,0 +1,228 @@ +import base64 +import math +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +import httpx +from pydantic import ConfigDict, TypeAdapter + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.image_generation.transformation import BaseImageGenerationConfig +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory lacks typed params + get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared HTTP factory is private +) +from litellm.llms.replicate.common_utils import ( + ReplicateError, + async_replicate_download, + async_wait_for_replicate_prediction, + replicate_api_base, + replicate_auth_headers, + replicate_download, + replicate_output_urls, + replicate_positive_int, + replicate_prediction_failure, + replicate_prediction_input, + replicate_prediction_request, + wait_for_replicate_prediction, +) +from litellm.types.llms.openai import AllMessageValues, OpenAIImageGenerationOptionalParams +from litellm.types.utils import ImageObject, ImageResponse, LlmProviders + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer + +_ImageParams: TypeAlias = dict[str, object] +_ImageHeaders: TypeAlias = dict[str, str] + + +def _aspect_ratio(size: object) -> Mapping[str, str]: + if not isinstance(size, str) or size.lower().count("x") != 1: + return MappingProxyType({}) + width_text, height_text = size.lower().split("x") + if not (width_text.isdigit() and height_text.isdigit()) or int(width_text) <= 0 or int(height_text) <= 0: + return MappingProxyType({}) + divisor: Final = math.gcd(int(width_text), int(height_text)) + return MappingProxyType({"aspect_ratio": f"{int(width_text) // divisor}:{int(height_text) // divisor}"}) + + +_NATIVE_INPUT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + + +def _get_replicate_httpx_client(ssl_verify: object) -> HTTPHandler: + return get_httpx_client(params={"ssl_verify": ssl_verify}) + + +def _get_replicate_async_httpx_client(ssl_verify: object) -> AsyncHTTPHandler: + return get_async_httpx_client(llm_provider=LlmProviders.REPLICATE, params={"ssl_verify": ssl_verify}) + + +class ReplicateImageGenerationConfig(BaseImageGenerationConfig): + def __init__( + self, + sync_client_factory: Callable[[object], HTTPHandler] = _get_replicate_httpx_client, + async_client_factory: Callable[[object], AsyncHTTPHandler] = _get_replicate_async_httpx_client, + ) -> None: + self._sync_client_factory: Final = sync_client_factory + self._async_client_factory: Final = async_client_factory + + def get_supported_openai_params( + self, model: str + ) -> list[OpenAIImageGenerationOptionalParams]: # mutable-ok: base class contract returns a plain list + return ["n", "response_format", "size"] + + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: _ImageParams, + model: str, + drop_params: bool, + ) -> _ImageParams: + supported: Final = frozenset(self.get_supported_openai_params(model)) + unsupported: Final = tuple(key for key in non_default_params if key not in supported) + if unsupported and not drop_params: + raise ValueError( + f"Parameters {', '.join(unsupported)} are not supported for model {model}. " + f"Supported parameters are {sorted(supported)}. Set drop_params=True to drop unsupported parameters." + ) + count: Final = non_default_params.get("n") + response_format: Final = non_default_params.get("response_format") + return { + **optional_params, + **({"num_outputs": count} if isinstance(count, int) and not isinstance(count, bool) else {}), + **_aspect_ratio(non_default_params.get("size")), + **({"response_format": response_format} if isinstance(response_format, str) else {}), + } + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: _ImageParams, + litellm_params: _ImageParams, + stream: bool | None = None, + ) -> str: + return replicate_prediction_request(replicate_api_base(api_base), model, MappingProxyType({})).url + + def validate_environment( + self, + headers: _ImageHeaders, + model: str, + messages: list[AllMessageValues], + optional_params: _ImageParams, + litellm_params: _ImageParams, + api_key: str | None = None, + api_base: str | None = None, + ) -> _ImageHeaders: + return {**replicate_auth_headers(headers, api_key), "Prefer": "wait"} + + def transform_image_generation_request( + self, + model: str, + prompt: str, + optional_params: _ImageParams, + litellm_params: _ImageParams, + headers: _ImageHeaders, + ) -> _ImageParams: + extra_body: Final = optional_params.get("extra_body") + prediction_input: Final = { + **{key: value for key, value in optional_params.items() if key not in ("response_format", "extra_body")}, + **(_NATIVE_INPUT.validate_python(extra_body) if isinstance(extra_body, Mapping) else {}), + "prompt": prompt, + } + return dict(replicate_prediction_request("", model, prediction_input).body) + + @staticmethod + def _image_urls(prediction: Mapping[str, object], raw_response: httpx.Response) -> tuple[str, ...]: + failure: Final = replicate_prediction_failure(prediction) + if failure is not None: + raise ReplicateError(status_code=raw_response.status_code, message=failure, headers=raw_response.headers) + urls: Final = replicate_output_urls(prediction.get("output")) + if not urls: + raise ReplicateError( + status_code=raw_response.status_code, + message="Replicate prediction finished without image output", + headers=raw_response.headers, + ) + return urls + + @staticmethod + def _native_size(prediction: Mapping[str, object]) -> Mapping[str, str]: + prediction_input: Final = replicate_prediction_input(prediction) + width: Final = replicate_positive_int(prediction_input.get("width")) + height: Final = replicate_positive_int(prediction_input.get("height")) + return {"size": f"{width}x{height}"} if width is not None and height is not None else {} + + @staticmethod + def _image_objects(urls: tuple[str, ...], contents: tuple[bytes, ...] | None) -> tuple[ImageObject, ...]: + if contents is None: + return tuple(ImageObject(url=url, b64_json=None) for url in urls) + return tuple(ImageObject(url=None, b64_json=base64.b64encode(content).decode()) for content in contents) + + def transform_image_generation_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ImageResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: _ImageParams, + optional_params: _ImageParams, + litellm_params: _ImageParams, + encoding: "Tokenizer | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ImageResponse: + client: Final = self._sync_client_factory(litellm_params.get("ssl_verify")) + prediction: Final = wait_for_replicate_prediction(raw_response, client) + urls: Final = self._image_urls(prediction, raw_response) + contents: Final = ( + tuple(self._download(client, url) for url in urls) + if optional_params.get("response_format") == "b64_json" + else None + ) + return model_response.model_copy( + update={"data": list(self._image_objects(urls, contents)), **self._native_size(prediction)} + ) + + async def async_transform_image_generation_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ImageResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: _ImageParams, + optional_params: _ImageParams, + litellm_params: _ImageParams, + encoding: "Tokenizer | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ImageResponse: + client: Final = self._async_client_factory(litellm_params.get("ssl_verify")) + prediction: Final = await async_wait_for_replicate_prediction(raw_response, client) + urls: Final = self._image_urls(prediction, raw_response) + contents: Final = ( + tuple([await self._async_download(client, url) for url in urls]) + if optional_params.get("response_format") == "b64_json" + else None + ) + return model_response.model_copy( + update={"data": list(self._image_objects(urls, contents)), **self._native_size(prediction)} + ) + + @staticmethod + def _download(client: HTTPHandler, url: str) -> bytes: + return replicate_download(client, url).content + + @staticmethod + async def _async_download(client: AsyncHTTPHandler, url: str) -> bytes: + return (await async_replicate_download(client, url)).content + + def get_error_class( + self, error_message: str, status_code: int, headers: _ImageHeaders | httpx.Headers + ) -> BaseLLMException: + return ReplicateError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/replicate/text_to_speech/transformation.py b/litellm/llms/replicate/text_to_speech/transformation.py new file mode 100644 index 00000000000..d75fc29ba54 --- /dev/null +++ b/litellm/llms/replicate/text_to_speech/transformation.py @@ -0,0 +1,204 @@ +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +import httpx +from pydantic import ConfigDict, TypeAdapter + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig, TextToSpeechRequestData +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory lacks typed params + get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared HTTP factory is private +) +from litellm.llms.replicate.common_utils import ( + ReplicateError, + async_replicate_download, + async_wait_for_replicate_prediction, + replicate_api_base, + replicate_auth_headers, + replicate_download, + replicate_model_owner, + replicate_output_urls, + replicate_positive_int, + replicate_prediction_failure, + replicate_prediction_metrics, + replicate_prediction_request, + wait_for_replicate_prediction, +) +from litellm.types.utils import LlmProviders, Usage + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import HttpxBinaryResponseContent + +_SpeechParams: TypeAlias = dict[str, object] +_SpeechHeaders: TypeAlias = dict[str, str] +_NATIVE_INPUT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_VOICE_KEY_BY_OWNER: Final[Mapping[str, str]] = MappingProxyType( + {"minimax": "voice_id", "inworld": "voice_id", "qwen": "speaker"} +) + + +def _voice_key(model: str) -> str: + return _VOICE_KEY_BY_OWNER.get(replicate_model_owner(model), "voice") + + +def _get_replicate_httpx_client(ssl_verify: object) -> HTTPHandler: + return get_httpx_client(params={"ssl_verify": ssl_verify}) + + +def _get_replicate_async_httpx_client(ssl_verify: object) -> AsyncHTTPHandler: + return get_async_httpx_client(llm_provider=LlmProviders.REPLICATE, params={"ssl_verify": ssl_verify}) + + +def _ssl_verify(logging_obj: "LiteLLMLoggingObj") -> object: + litellm_params: Final = getattr(logging_obj, "litellm_params", None) + return litellm_params.get("ssl_verify") if isinstance(litellm_params, Mapping) else None + + +class ReplicateTextToSpeechConfig(BaseTextToSpeechConfig): + def __init__( + self, + sync_client_factory: Callable[[object], HTTPHandler] = _get_replicate_httpx_client, + async_client_factory: Callable[[object], AsyncHTTPHandler] = _get_replicate_async_httpx_client, + ) -> None: + super().__init__() + self._sync_client_factory: Final = sync_client_factory + self._async_client_factory: Final = async_client_factory + + def get_supported_openai_params( + self, model: str + ) -> list[str]: # mutable-ok: base class contract returns a plain list + return ["voice", "speed"] + + def map_openai_params( + self, + model: str, + optional_params: Mapping[str, object], + voice: object = None, + drop_params: bool = False, + kwargs: Mapping[str, object] | None = None, + ) -> tuple[str | None, _SpeechParams]: + speed: Final = optional_params.get("speed") + extra_body: Final = (kwargs or {}).get("extra_body") + return ( + voice if isinstance(voice, str) and voice else None, + { + **(_NATIVE_INPUT.validate_python(extra_body) if isinstance(extra_body, Mapping) else {}), + **({"speed": speed} if isinstance(speed, (int, float)) and not isinstance(speed, bool) else {}), + }, + ) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None = None, + api_base: str | None = None, + ) -> _SpeechHeaders: + return {**replicate_auth_headers(headers, api_key), "Prefer": "wait"} + + def get_complete_url(self, model: str, api_base: str | None, litellm_params: Mapping[str, object]) -> str: + return replicate_prediction_request(replicate_api_base(api_base), model, MappingProxyType({})).url + + def transform_text_to_speech_request( + self, + model: str, + input: str, + voice: str | None, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, str], + ) -> TextToSpeechRequestData: + prediction_input: Final = { + **optional_params, + **({_voice_key(model): voice} if voice else {}), + "text": input, + } + return {"dict_body": dict(replicate_prediction_request("", model, prediction_input).body)} + + @staticmethod + def _audio_url(prediction: Mapping[str, object], raw_response: httpx.Response) -> str: + failure: Final = replicate_prediction_failure(prediction) + if failure is not None: + raise ReplicateError(status_code=raw_response.status_code, message=failure, headers=raw_response.headers) + urls: Final = replicate_output_urls(prediction.get("output")) + if not urls: + raise ReplicateError( + status_code=raw_response.status_code, + message="Replicate prediction finished without audio output", + headers=raw_response.headers, + ) + return urls[0] + + @staticmethod + def _usage(prediction: Mapping[str, object]) -> Usage | None: + metrics: Final = replicate_prediction_metrics(prediction) + prompt_tokens: Final = replicate_positive_int( + metrics.get("token_input_count", metrics.get("input_token_count")) + ) + completion_tokens: Final = replicate_positive_int( + metrics.get("token_output_count", metrics.get("output_token_count")) + ) + if prompt_tokens is None and completion_tokens is None: + return None + return Usage( + prompt_tokens=prompt_tokens or 0, + completion_tokens=completion_tokens or 0, + total_tokens=(prompt_tokens or 0) + (completion_tokens or 0), + ) + + @staticmethod + def _binary_response( + audio_response: httpx.Response, raw_response: httpx.Response, usage: Usage | None + ) -> "HttpxBinaryResponseContent": + from litellm.types.llms.openai import HttpxBinaryResponseContent + + audio_response.raise_for_status() + return HttpxBinaryResponseContent( + httpx.Response( + status_code=200, + headers={ + "content-type": audio_response.headers.get("content-type", "application/octet-stream"), + "content-length": str(len(audio_response.content)), + }, + content=audio_response.content, + request=raw_response.request, + ), + usage=usage, + ) + + def transform_text_to_speech_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> "HttpxBinaryResponseContent": + client: Final = self._sync_client_factory(_ssl_verify(logging_obj)) + prediction: Final = wait_for_replicate_prediction(raw_response, client) + audio_url: Final = self._audio_url(prediction, raw_response) + return self._binary_response(replicate_download(client, audio_url), raw_response, self._usage(prediction)) + + async def async_transform_text_to_speech_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> "HttpxBinaryResponseContent": + client: Final = self._async_client_factory(_ssl_verify(logging_obj)) + prediction: Final = await async_wait_for_replicate_prediction(raw_response, client) + audio_url: Final = self._audio_url(prediction, raw_response) + return self._binary_response( + await async_replicate_download(client, audio_url), raw_response, self._usage(prediction) + ) + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: _SpeechHeaders | httpx.Headers, + ) -> BaseLLMException: + return ReplicateError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/replicate/videos/transformation.py b/litellm/llms/replicate/videos/transformation.py new file mode 100644 index 00000000000..cd01c832686 --- /dev/null +++ b/litellm/llms/replicate/videos/transformation.py @@ -0,0 +1,486 @@ +import math +import sys +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, TypeAlias + +import httpx +from httpx._types import FileContent, RequestFiles + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.videos.async_task_utils import unsupported_video_fields +from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory lacks typed params + get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared HTTP factory is private +) +from litellm.llms.replicate.common_utils import ( + ReplicateError, + async_replicate_download, + replicate_api_base, + replicate_auth_headers, + replicate_download, + replicate_model_owner, + replicate_output_urls, + replicate_prediction, + replicate_prediction_failure, + replicate_prediction_input, + replicate_prediction_request, + replicate_prediction_string, + replicate_prediction_url, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.types.videos.main import CharacterObject, VideoCreateOptionalRequestParams, VideoObject +from litellm.types.videos.utils import decode_video_id_with_provider, encode_video_id_with_provider + +_REPLICATE_PROVIDER: Final[str] = LlmProviders.REPLICATE.value +_STATUS_MAP: Final[Mapping[str, str]] = MappingProxyType( + { + "starting": "queued", + "processing": "in_progress", + "succeeded": "completed", + "failed": "failed", + "canceled": "failed", + } +) +_VideoParams: TypeAlias = dict[str, object] +_VideoHeaders: TypeAlias = dict[str, str] +_VideoStringParams: TypeAlias = dict[str, str] +_VideoFiles: TypeAlias = list[object] + + +@dataclass(frozen=True, slots=True) +class _ModelProfile: + resolution_tiers: tuple[tuple[int, str], ...] + aspect_ratio_key: str + reference_key: str + + +_GENERIC_PROFILE: Final[_ModelProfile] = _ModelProfile( + resolution_tiers=((480, "480p"), (720, "720p"), (1080, "1080p"), (sys.maxsize, "4k")), + aspect_ratio_key="aspect_ratio", + reference_key="image", +) +_MINIMAX_PROFILE: Final[_ModelProfile] = _ModelProfile( + resolution_tiers=((480, "480P"), (768, "768P"), (1440, "2K"), (sys.maxsize, "4K")), + aspect_ratio_key="ratio", + reference_key="first_frame_image", +) + + +def _profile_for_model(model: str) -> _ModelProfile: + return _MINIMAX_PROFILE if replicate_model_owner(model) == "minimax" else _GENERIC_PROFILE + + +def _duration_value(seconds: object) -> int | float: + if isinstance(seconds, bool) or not isinstance(seconds, (int, float, str)): + raise TypeError("Replicate seconds must be a numeric value") + try: + duration: Final = float(seconds) + except ValueError as error: + raise ValueError("Replicate seconds must be a numeric value") from error + return int(duration) if duration.is_integer() else duration + + +def _size_params(size: object, profile: _ModelProfile) -> Mapping[str, object]: + if not isinstance(size, str): + return MappingProxyType({}) + normalized: Final = size.strip().lower() + named_resolution: Final = next( + (resolution for _, resolution in profile.resolution_tiers if resolution.lower() == normalized), None + ) + if named_resolution is not None: + return MappingProxyType({"resolution": named_resolution}) + if normalized.count("x") != 1: + return MappingProxyType({}) + width_text, height_text = normalized.split("x") + if not (width_text.isdigit() and height_text.isdigit()) or int(width_text) <= 0 or int(height_text) <= 0: + return MappingProxyType({}) + width: Final = int(width_text) + height: Final = int(height_text) + divisor: Final = math.gcd(width, height) + short_side: Final = min(width, height) + resolution: Final = next( + (resolution for threshold, resolution in profile.resolution_tiers if short_side <= threshold), + profile.resolution_tiers[-1][1], + ) + return MappingProxyType( + {"resolution": resolution, profile.aspect_ratio_key: f"{width // divisor}:{height // divisor}"} + ) + + +def _usage_from_input(prediction_input: Mapping[str, object]) -> _VideoParams: + duration: Final = prediction_input.get("duration") + resolution: Final = prediction_input.get("resolution") + return { + key: value + for key, value in ( + ( + "duration_seconds", + float(duration) if isinstance(duration, (int, float)) and not isinstance(duration, bool) else None, + ), + ("video_resolution", resolution if isinstance(resolution, str) else None), + ) + if value is not None + } + + +def _video_object( + prediction: Mapping[str, object], + model: str | None, + custom_llm_provider: str | None, + created_at: int, +) -> VideoObject: + prediction_id: Final = replicate_prediction_string(prediction, "id") or "" + failure: Final = replicate_prediction_failure(prediction) + raw_status: Final = replicate_prediction_string(prediction, "status") or "starting" + prediction_input: Final = replicate_prediction_input(prediction) + duration: Final = prediction_input.get("duration") + resolution: Final = prediction_input.get("resolution") + return VideoObject( + id=encode_video_id_with_provider(prediction_id, custom_llm_provider or _REPLICATE_PROVIDER, model), + object="video", + status="failed" if failure else _STATUS_MAP.get(raw_status, "queued"), + created_at=created_at, + model=model, + seconds=str(duration) if isinstance(duration, (int, float)) and not isinstance(duration, bool) else None, + size=resolution if isinstance(resolution, str) else None, + error={"code": "replicate_error", "message": failure} if failure else None, + usage=dict(_usage_from_input(prediction_input)), + ) + + +def _owned_prediction(raw_response: httpx.Response) -> Mapping[str, object]: + prediction: Final = replicate_prediction(raw_response) + requested: Final = raw_response.request.url.fragment + if not requested or requested.startswith("deployments/"): + return prediction + requested_path, _, requested_version = requested.partition(":") + actual_path: Final = replicate_prediction_string(prediction, "model") + actual_version: Final = replicate_prediction_string(prediction, "version") + if (actual_path is not None and actual_path != requested_path) or ( + requested_version and actual_version is not None and actual_version != requested_version + ): + raise ReplicateError( + status_code=404, + message=f"Replicate prediction {replicate_prediction_string(prediction, 'id')} does not belong to {requested}", + headers=raw_response.headers, + ) + return prediction + + +def _get_replicate_async_httpx_client() -> AsyncHTTPHandler: + return get_async_httpx_client(llm_provider=LlmProviders.REPLICATE) + + +class ReplicateVideoConfig(BaseVideoConfig): + def __init__( + self, + sync_client_factory: Callable[[], HTTPHandler] = get_httpx_client, + async_client_factory: Callable[[], AsyncHTTPHandler] = _get_replicate_async_httpx_client, + ) -> None: + super().__init__() + self._sync_client_factory: Final = sync_client_factory + self._async_client_factory: Final = async_client_factory + + def get_supported_openai_params( + self, model: str + ) -> list[str]: # mutable-ok: base class contract returns a plain list + return ["model", "prompt", "input_reference", "seconds", "size", "extra_headers"] + + def map_openai_params( + self, + video_create_optional_params: VideoCreateOptionalRequestParams, + model: str, + drop_params: bool, + ) -> _VideoParams: + unsupported: Final = unsupported_video_fields( + video_create_optional_params, + model, + _REPLICATE_PROVIDER, + drop_params, + frozenset({"generate_audio", "aspect_ratio", "seed", "duration"}), + ) + supported: Final = frozenset(self.get_supported_openai_params(model)) + profile: Final = _profile_for_model(model) + input_reference: Final = video_create_optional_params.get("input_reference") + if "input_reference" in video_create_optional_params and not isinstance(input_reference, str): + raise ValueError("Replicate needs a public image URL for input_reference") + reference_params: Final[Mapping[str, object]] = ( + MappingProxyType({profile.reference_key: input_reference}) + if isinstance(input_reference, str) + else MappingProxyType({}) + ) + duration_params: Final[Mapping[str, object]] = ( + MappingProxyType({"duration": _duration_value(video_create_optional_params["seconds"])}) + if "seconds" in video_create_optional_params + else MappingProxyType({}) + ) + size_params: Final[Mapping[str, object]] = ( + _size_params(video_create_optional_params["size"], profile) + if "size" in video_create_optional_params + else MappingProxyType({}) + ) + return { + **reference_params, + **duration_params, + **size_params, + **{ + key: value + for key, value in video_create_optional_params.items() + if key not in supported and key not in unsupported and key != "user" + }, + } + + def validate_environment( + self, + headers: _VideoHeaders, + model: str, + api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, + ) -> _VideoHeaders: + return replicate_auth_headers( + headers, api_key or (litellm_params.api_key if litellm_params is not None else None) + ) + + def get_complete_url(self, model: str, api_base: str | None, litellm_params: _VideoParams) -> str: + return replicate_api_base(api_base) + + def transform_video_create_request( + self, + model: str, + prompt: str, + api_base: str, + video_create_optional_request_params: _VideoParams, + litellm_params: GenericLiteLLMParams, + headers: _VideoHeaders, + ) -> tuple[_VideoParams, RequestFiles, str]: + prediction_input: Final = { + **{key: value for key, value in video_create_optional_request_params.items() if key != "model"}, + "prompt": prompt, + } + request: Final = replicate_prediction_request(api_base, model, prediction_input) + return dict(request.body), [], request.url + + def transform_video_create_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + request_data: Mapping[str, object] | None = None, + ) -> VideoObject: + return _video_object(replicate_prediction(raw_response), model, custom_llm_provider, int(time.time())) + + @staticmethod + def _decode_video_id(video_id: str) -> tuple[str, str | None]: + decoded: Final = decode_video_id_with_provider(video_id) + return decoded.get("video_id", video_id), decoded.get("model_id") + + @classmethod + def _prediction_url_for_video(cls, video_id: str, api_base: str) -> str: + prediction_id, model_id = cls._decode_video_id(video_id) + prediction_url: Final = httpx.URL(replicate_prediction_url(api_base, prediction_id)) + return str(prediction_url.copy_with(fragment=model_id) if model_id else prediction_url) + + def transform_video_status_retrieve_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: + return self._prediction_url_for_video(video_id, api_base), {} + + def transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + client: HTTPHandler | None = None, + ) -> VideoObject: + prediction: Final = _owned_prediction(raw_response) + model: Final = raw_response.request.url.fragment or replicate_prediction_string(prediction, "model") + return _video_object(prediction, model, custom_llm_provider, 0) + + def transform_video_content_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: _VideoHeaders, + variant: str | None = None, + ) -> tuple[str, _VideoStringParams]: + return self._prediction_url_for_video(video_id, api_base), {} + + @staticmethod + def _video_url(raw_response: httpx.Response) -> str: + prediction: Final = _owned_prediction(raw_response) + failure: Final = replicate_prediction_failure(prediction) + if failure is not None: + raise ReplicateError(status_code=raw_response.status_code, message=failure, headers=raw_response.headers) + urls: Final = replicate_output_urls(prediction.get("output")) + if not urls: + raise ReplicateError( + status_code=raw_response.status_code, + message=f"Replicate prediction has no video output yet (status {replicate_prediction_string(prediction, 'status')})", + headers=raw_response.headers, + ) + return urls[0] + + def transform_video_content_response( + self, + raw_response: httpx.Response, + logging_obj: object, + client: "HTTPHandler | None" = None, + ) -> bytes: + download_client: Final = client if client is not None else self._sync_client_factory() + return replicate_download(download_client, self._video_url(raw_response)).content + + async def async_transform_video_content_response( + self, + raw_response: httpx.Response, + logging_obj: object, + client: "AsyncHTTPHandler | None" = None, + ) -> bytes: + download_client: Final = client if client is not None else self._async_client_factory() + return (await async_replicate_download(download_client, self._video_url(raw_response))).content + + def transform_video_remix_request( + self, + video_id: str, + prompt: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: _VideoHeaders, + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, _VideoParams]: + raise NotImplementedError("video remix is not supported for Replicate") + + def transform_video_remix_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("video remix is not supported for Replicate") + + def transform_video_list_request( + self, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: _VideoHeaders, + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_query: Mapping[str, object] | None = None, + ) -> tuple[str, _VideoParams]: + raise NotImplementedError("video listing is not supported for Replicate") + + def transform_video_list_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> _VideoStringParams: + raise NotImplementedError("video listing is not supported for Replicate") + + def transform_video_delete_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: + raise NotImplementedError("video delete is not supported for Replicate") + + def transform_video_delete_response(self, raw_response: httpx.Response, logging_obj: object) -> VideoObject: + raise NotImplementedError("video delete is not supported for Replicate") + + def transform_video_create_character_request( + self, + name: str, + video: object, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: _VideoHeaders, + ) -> tuple[str, _VideoFiles]: + raise NotImplementedError("video character creation is not supported for Replicate") + + def transform_video_create_character_response( + self, + raw_response: httpx.Response, + logging_obj: object, + ) -> CharacterObject: + raise NotImplementedError("video character creation is not supported for Replicate") + + def transform_video_get_character_request( + self, + character_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: + raise NotImplementedError("video character retrieval is not supported for Replicate") + + def transform_video_get_character_response( + self, + raw_response: httpx.Response, + logging_obj: object, + ) -> CharacterObject: + raise NotImplementedError("video character retrieval is not supported for Replicate") + + def transform_video_edit_request( + self, + prompt: str, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: _VideoHeaders, + video_file: FileContent | None = None, + extra_body: Mapping[str, object] | None = None, + prefetched_source_data: Mapping[str, object] | None = None, + ) -> tuple[str, Mapping[str, object], RequestFiles | None]: + raise NotImplementedError("video edit is not supported for Replicate") + + def transform_video_edit_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + request_data: Mapping[str, object] | None = None, + ) -> VideoObject: + raise NotImplementedError("video edit is not supported for Replicate") + + def transform_video_extension_request( + self, + prompt: str, + video_id: str, + seconds: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: _VideoHeaders, + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, _VideoParams]: + raise NotImplementedError("video extension is not supported for Replicate") + + def transform_video_extension_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("video extension is not supported for Replicate") + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: _VideoHeaders | httpx.Headers, + ) -> BaseLLMException: + return ReplicateError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/main.py b/litellm/main.py index e4110455e24..35473e63852 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -9083,7 +9083,7 @@ def speech( custom_llm_provider=custom_llm_provider, ) response: HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent] | None = None - if custom_llm_provider in ("edenai", "dashscope") and text_to_speech_provider_config is not None: + if custom_llm_provider in ("edenai", "dashscope", "replicate") and text_to_speech_provider_config is not None: litellm_params_dict["api_base"] = caller_api_base if custom_llm_provider == "dashscope" else api_base response = base_llm_http_handler.text_to_speech_handler( model=model, diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 268bce1cf6f..daabd6122cd 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -145,10 +145,10 @@ class HttpxBinaryResponseContent(_HttpxBinaryResponseContent, Generic[_ResponseT def hidden_params(self, hidden_params: dict[str, builtins.object]) -> None: # mutable-ok: API requires mutation self._hidden_params = hidden_params - def __init__(self, response: _ResponseT) -> None: + def __init__(self, response: _ResponseT, usage: "Usage | None" = None) -> None: super().__init__(response) # pyright: ignore[reportArgumentType] # SDK accepts both backends at runtime self._hidden_params = {} - self.usage = None + self.usage = usage @property def content_type(self) -> str | None: diff --git a/litellm/utils.py b/litellm/utils.py index 19b0eb80c41..c0eec2d55fb 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -11073,6 +11073,10 @@ class ProviderConfigManager: ) return get_fal_ai_image_generation_config(model) + elif LlmProviders.REPLICATE == provider: + from litellm.llms.replicate.image_generation.transformation import ReplicateImageGenerationConfig + + return ReplicateImageGenerationConfig() elif LlmProviders.STABILITY == provider: from litellm.llms.stability.image_generation import ( get_stability_image_generation_config, @@ -11174,6 +11178,10 @@ class ProviderConfigManager: from litellm.llms.fal_ai.videos.transformation import FalAIVideoConfig return FalAIVideoConfig() + elif LlmProviders.REPLICATE == provider: + from litellm.llms.replicate.videos.transformation import ReplicateVideoConfig + + return ReplicateVideoConfig() elif LlmProviders.HOSTED_VLLM == provider: from litellm.llms.hosted_vllm.videos import get_hosted_vllm_video_config @@ -11558,6 +11566,10 @@ class ProviderConfigManager: ) return MistralTextToSpeechConfig() + elif litellm.LlmProviders.REPLICATE == provider: + from litellm.llms.replicate.text_to_speech.transformation import ReplicateTextToSpeechConfig + + return ReplicateTextToSpeechConfig() elif litellm.LlmProviders.AWS_POLLY == provider: from litellm.llms.aws_polly.text_to_speech.transformation import ( AWSPollyTextToSpeechConfig, diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index adb6410d360..9698c4ac62f 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1,16 +1,19 @@ -import asyncio, importlib, json, time +import asyncio import gc -import io +import importlib +import json import logging import os import pathlib import ssl import threading +import time import weakref from collections.abc import AsyncIterator, Callable, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from contextvars import Context +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Final from unittest.mock import MagicMock, patch @@ -25,13 +28,17 @@ from yarl import URL import litellm from litellm.caching.llm_caching_handler import LLMClientCache -from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport +from litellm.exceptions import Timeout as LitellmTimeout +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.model_cost_table import publish from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine, default_token_exchange_engine from litellm.llms.base_llm.auth.types import MintedToken, TokenExchangeSpec +from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport from litellm.llms.custom_httpx.http_handler import ( _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER, AsyncHTTPHandler, HTTPHandler, + HTTPResponseLimitError, MaskedHTTPStatusError, get_httpx_client, get_shared_realtime_ssl_context, @@ -39,10 +46,6 @@ from litellm.llms.custom_httpx.http_handler import ( realtime_ssl_for_url, ) from litellm.types.llms.custom_http import VerifyTypes -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from litellm.exceptions import Timeout as LitellmTimeout -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.litellm_core_utils.model_cost_table import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -493,9 +496,11 @@ async def test_async_handler_with_shared_session(): async def test_get_async_httpx_client_with_shared_session(): """Test get_async_httpx_client with shared session""" from litellm.llms.custom_httpx.http_handler import ( - get_async_httpx_client, AsyncHTTPHandler as AsyncHTTPHandlerReload, ) + from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + ) from litellm.types.utils import LlmProviders # Create a mock shared session @@ -517,9 +522,11 @@ async def test_get_async_httpx_client_with_shared_session(): async def test_get_async_httpx_client_without_shared_session(): """Test get_async_httpx_client without shared session (backward compatibility)""" from litellm.llms.custom_httpx.http_handler import ( - get_async_httpx_client, AsyncHTTPHandler as AsyncHTTPHandlerReload, ) + from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + ) from litellm.types.utils import LlmProviders # Test without shared session @@ -555,6 +562,7 @@ async def test_session_reuse_chain(): def test_shared_session_parameter_in_acompletion(): """Test that acompletion function accepts shared_session parameter""" import inspect + from litellm.main import acompletion # Get the function signature @@ -572,6 +580,7 @@ def test_shared_session_parameter_in_acompletion(): def test_shared_session_parameter_in_completion(): """Test that completion function accepts shared_session parameter""" import inspect + from litellm.main import completion # Get the function signature @@ -590,9 +599,11 @@ def test_shared_session_parameter_in_completion(): async def test_session_reuse_integration(): """Integration test for session reuse functionality""" from litellm.llms.custom_httpx.http_handler import ( - get_async_httpx_client, AsyncHTTPHandler as AsyncHTTPHandlerReload, ) + from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + ) from litellm.types.utils import LlmProviders # Create a mock session @@ -2395,7 +2406,9 @@ def test_realtime_ssl_for_url_keeps_the_shared_context_for_wss_endpoints() -> No assert realtime_ssl_for_url("wss://aiplatform.us.rep.googleapis.com/ws") is shared -def test_realtime_ssl_for_url_turns_ssl_verify_false_into_an_unverified_tls_context(monkeypatch: pytest.MonkeyPatch) -> None: +def test_realtime_ssl_for_url_turns_ssl_verify_false_into_an_unverified_tls_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setattr("litellm.llms.custom_httpx.http_handler._shared_realtime_ssl_context", False) selected: Final = realtime_ssl_for_url("wss://aiplatform.us.rep.googleapis.com/ws") assert isinstance(selected, ssl.SSLContext) @@ -2411,7 +2424,10 @@ def _ollama_generate_route(respx_mock: respx.MockRouter, stream: bool) -> respx. final: Final = {"model": "llama3.1", "response": "", "done": True, "prompt_eval_count": 2, "eval_count": 1} if stream: return respx_mock.post("http://ollama.test:11434/api/generate").respond( - text=json.dumps({"model": "llama3.1", "response": "Sunny", "done": False}) + "\n" + json.dumps(final) + "\n", + text=json.dumps({"model": "llama3.1", "response": "Sunny", "done": False}) + + "\n" + + json.dumps(final) + + "\n", headers={"content-type": "application/x-ndjson"}, ) return respx_mock.post("http://ollama.test:11434/api/generate").respond(json={**final, "response": "Sunny"}) @@ -2567,8 +2583,11 @@ class _PoolDiagnosticClock: return self.now def call_later( - self, delay: float, callback: Callable[..., object], - *args: object, context: Context | None = None, + self, + delay: float, + callback: Callable[..., object], + *args: object, + context: Context | None = None, ) -> asyncio.TimerHandle: timer: Final = self.schedule(delay, callback, *args, context=context) self.timers.append(timer) @@ -2631,7 +2650,9 @@ def _pool_diagnostic_messages(caplog: pytest.LogCaptureFixture) -> tuple[str, .. @pytest.mark.parametrize("source", ["sdk", "proxy", "native-connector", "rebuilt"]) @pytest.mark.parametrize("error", [TimeoutError("private timeout"), asyncio.CancelledError("private cancellation")]) async def test_aiohttp_pool_failure_identifies_queue_wait_and_cleans_up_timer( - source: str, error: BaseException, monkeypatch: pytest.MonkeyPatch, + source: str, + error: BaseException, + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, ) -> None: pool_diagnostic_clock: Final = _pool_diagnostic_clock(monkeypatch) @@ -2663,7 +2684,8 @@ async def test_aiohttp_pool_failure_identifies_queue_wait_and_cleans_up_timer( @pytest.mark.parametrize("queued", [False, True]) async def test_aiohttp_pool_diagnostics_do_not_label_slow_upstream_as_pool_saturation( - queued: bool, monkeypatch: pytest.MonkeyPatch, + queued: bool, + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, ) -> None: pool_diagnostic_clock: Final = _pool_diagnostic_clock(monkeypatch) @@ -2685,7 +2707,8 @@ async def test_aiohttp_pool_diagnostics_do_not_label_slow_upstream_as_pool_satur async def test_aiohttp_pool_warns_while_queued_once_per_request_with_isolated_concurrent_waits( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, ) -> None: pool_diagnostic_clock: Final = _pool_diagnostic_clock(monkeypatch) caplog.set_level(logging.WARNING, logger="LiteLLM") @@ -2721,15 +2744,22 @@ async def test_aiohttp_pool_warns_while_queued_once_per_request_with_isolated_co assert all(timer.cancelled() for timer in pool_diagnostic_clock.timers) -@pytest.mark.parametrize(("header", "location", "host", "port"), [ - ("Location", "https://redirected.test:9443/private-path?token=private-query", "redirected.test", 9443), - ("Location", "/private-path?token=private-query", "initial.test", 8443), - ("URI", "https://redirected.test/private-path", "redirected.test", 443), - ("Location", "http://[invalid", None, None), -]) +@pytest.mark.parametrize( + ("header", "location", "host", "port"), + [ + ("Location", "https://redirected.test:9443/private-path?token=private-query", "redirected.test", 9443), + ("Location", "/private-path?token=private-query", "initial.test", 8443), + ("URI", "https://redirected.test/private-path", "redirected.test", 443), + ("Location", "http://[invalid", None, None), + ], +) async def test_aiohttp_pool_warning_tracks_redirect_destination_without_masking_invalid_redirects( - header: str, location: str, host: str | None, port: int | None, - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, + header: str, + location: str, + host: str | None, + port: int | None, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, ) -> None: from aiohttp.client_proto import ResponseHandler from aiohttp.helpers import TimerNoop @@ -2742,9 +2772,15 @@ async def test_aiohttp_pool_warning_tracks_redirect_destination_without_masking_ traces: Final = await _pool_diagnostic_request(session, url) loop: Final = asyncio.get_running_loop() response: Final = ClientResponse( - "POST", url, writer=None, continue100=None, timer=TimerNoop(), + "POST", + url, + writer=None, + continue100=None, + timer=TimerNoop(), request_info=RequestInfo(url, "POST", CIMultiDictProxy(CIMultiDict()), url), - traces=[], loop=loop, session=session, + traces=[], + loop=loop, + session=session, stream_writer=StreamWriter(ResponseHandler(loop), loop), ) response._headers = CIMultiDictProxy(CIMultiDict({header: location})) @@ -2768,3 +2804,48 @@ async def test_aiohttp_pool_warning_tracks_redirect_destination_without_masking_ assert "private" not in messages[0] finally: response.close() + + +def test_sync_bounded_get_returns_a_body_within_the_limit(respx_mock: respx.MockRouter): + respx_mock.get("https://example.com/small.bin").respond(200, content=b"x" * 50) + handler = HTTPHandler() + try: + response = handler.get("https://example.com/small.bin", max_response_bytes=100) + finally: + handler.close() + assert response.status_code == 200 + assert response.content == b"x" * 50 + + +@pytest.mark.parametrize("headers", [{"content-length": "101"}, {}]) +def test_sync_bounded_get_refuses_a_body_over_the_limit(respx_mock: respx.MockRouter, headers: dict[str, str]): + respx_mock.get("https://example.com/big.bin").respond(200, content=b"x" * 101, headers=headers) + handler = HTTPHandler() + try: + with pytest.raises(HTTPResponseLimitError): + handler.get("https://example.com/big.bin", max_response_bytes=100) + finally: + handler.close() + + +def test_sync_bounded_get_follows_a_redirect_with_the_limit(respx_mock: respx.MockRouter): + respx_mock.get("https://example.com/first.bin").respond(302, headers={"location": "/final.bin"}) + final = respx_mock.get("https://example.com/final.bin").respond(200, content=b"ok") + handler = HTTPHandler() + try: + response = handler.get("https://example.com/first.bin", max_response_bytes=100, follow_redirects=True) + finally: + handler.close() + assert response.content == b"ok" + assert final.calls[0].request.headers["accept-encoding"] == "identity" + + +def test_sync_bounded_get_stops_redirect_loops(respx_mock: respx.MockRouter): + route = respx_mock.get("https://example.com/loop.bin").respond(302, headers={"location": "/loop.bin"}) + handler = HTTPHandler() + try: + with pytest.raises(ValueError, match="Too many redirects"): + handler.get("https://example.com/loop.bin", max_response_bytes=100, follow_redirects=True) + finally: + handler.close() + assert route.call_count == 11 diff --git a/tests/unit/llms/replicate/image_generation/__init__.py b/tests/unit/llms/replicate/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/replicate/image_generation/test_replicate_image_generation_transformation.py b/tests/unit/llms/replicate/image_generation/test_replicate_image_generation_transformation.py new file mode 100644 index 00000000000..6ae0d364b6a --- /dev/null +++ b/tests/unit/llms/replicate/image_generation/test_replicate_image_generation_transformation.py @@ -0,0 +1,391 @@ +import base64 +import time +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError +from litellm.constants import REPLICATE_MAX_OUTPUT_DOWNLOAD_BYTES +from litellm.litellm_core_utils.url_utils import SSRFError +from litellm.llms.replicate import common_utils +from litellm.llms.replicate.common_utils import ReplicateError +from litellm.llms.replicate.image_generation.transformation import ReplicateImageGenerationConfig +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +API_BASE: Final = "https://api.replicate.com/v1" +MODEL: Final = "black-forest-labs/flux-schnell" +PREDICTIONS_URL: Final = f"{API_BASE}/models/{MODEL}/predictions" +IMAGE_URLS: Final = ("https://replicate.delivery/out/0.webp", "https://replicate.delivery/out/1.webp") + + +@pytest.fixture(autouse=True) +def _skip_url_dns_validation(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + + +class _FakeClient: + def __init__(self, *responses: httpx.Response) -> None: + self.responses = list(responses) + self.calls: list[tuple[str, dict[str, str] | None]] = [] + + def get( + self, url: str, headers: dict[str, str] | None = None, params: object = None, **kwargs: object + ) -> httpx.Response: + self.calls.append((url, headers)) + return self.responses.pop(0) + + +class _FakeAsyncClient(_FakeClient): + async def get( + self, url: str, headers: dict[str, str] | None = None, params: object = None, **kwargs: object + ) -> httpx.Response: # pyright: ignore[reportIncompatibleMethodOverride] # async twin of the sync fake + return super().get(url, headers=headers, params=params, **kwargs) + + +def _prediction_response(body: dict[str, object], status_code: int = 201) -> httpx.Response: + request = httpx.Request( + "POST", + PREDICTIONS_URL, + headers={"Authorization": "Bearer k", "Prefer": "wait", "Content-Type": "application/json"}, + ) + return httpx.Response(status_code, json=body, request=request) + + +def _download(url: str, content: bytes) -> httpx.Response: + return httpx.Response(200, content=content, request=httpx.Request("GET", url)) + + +def _config(client: _FakeClient) -> ReplicateImageGenerationConfig: + return ReplicateImageGenerationConfig(sync_client_factory=lambda _ssl_verify: client) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + + +def _transform( + config: ReplicateImageGenerationConfig, response: httpx.Response, optional_params: dict[str, object] | None = None +): + return config.transform_image_generation_response( + model=MODEL, + raw_response=response, + model_response=litellm.ImageResponse(), + logging_obj=Mock(), + request_data={}, + optional_params=optional_params or {}, + litellm_params={}, + encoding=None, + ) + + +def test_provider_config_is_registered(): + config = ProviderConfigManager.get_provider_image_generation_config(model=MODEL, provider=LlmProviders.REPLICATE) + assert isinstance(config, ReplicateImageGenerationConfig) + + +def test_map_openai_params_maps_n_and_size(): + mapped = ReplicateImageGenerationConfig().map_openai_params( + {"n": 2, "size": "1024x1024", "response_format": "b64_json"}, {}, MODEL, False + ) + assert mapped == {"num_outputs": 2, "aspect_ratio": "1:1", "response_format": "b64_json"} + assert ReplicateImageGenerationConfig().map_openai_params({"size": "1792x1024"}, {}, MODEL, False) == { + "aspect_ratio": "7:4" + } + + +def test_map_openai_params_rejects_or_drops_unsupported(): + with pytest.raises(ValueError, match="quality"): + ReplicateImageGenerationConfig().map_openai_params({"quality": "hd"}, {}, MODEL, False) + assert ReplicateImageGenerationConfig().map_openai_params({"quality": "hd", "n": 1}, {}, MODEL, True) == { + "num_outputs": 1 + } + + +def test_validate_environment_sets_bearer_and_prefer_wait(): + headers = ReplicateImageGenerationConfig().validate_environment({}, MODEL, [], {}, {}, api_key="r8_key") + assert headers["Authorization"] == "Bearer r8_key" + assert headers["Prefer"] == "wait" + + +def test_request_url_and_body(): + config = ReplicateImageGenerationConfig() + assert config.get_complete_url(None, None, MODEL, {}, {}) == PREDICTIONS_URL + assert config.get_complete_url("https://proxy.internal", None, "owner/name:v1", {}, {}) == ( + "https://proxy.internal/v1/models/owner/name/versions/v1/predictions" + ) + body = config.transform_image_generation_request( + MODEL, "a red fox", {"num_outputs": 2, "aspect_ratio": "1:1", "response_format": "url"}, {}, {} + ) + assert body == {"input": {"prompt": "a red fox", "num_outputs": 2, "aspect_ratio": "1:1"}} + versioned = config.transform_image_generation_request("owner/name:v1", "a red fox", {}, {}, {}) + assert versioned == {"input": {"prompt": "a red fox"}} + + +@pytest.mark.parametrize("output", [list(IMAGE_URLS), IMAGE_URLS[0]]) +def test_finished_prediction_returns_urls_without_downloading(output: object): + client = _FakeClient() + result = _transform(_config(client), _prediction_response({"id": "p1", "status": "succeeded", "output": output})) + assert [image.url for image in result.data] == (list(IMAGE_URLS) if isinstance(output, list) else [IMAGE_URLS[0]]) + assert all(image.b64_json is None for image in result.data) + assert client.calls == [] + + +def test_b64_json_downloads_each_output(): + client = _FakeClient(_download(IMAGE_URLS[0], b"img0"), _download(IMAGE_URLS[1], b"img1")) + result = _transform( + _config(client), + _prediction_response({"id": "p1", "status": "succeeded", "output": list(IMAGE_URLS)}), + {"response_format": "b64_json"}, + ) + assert [image.b64_json for image in result.data] == [ + base64.b64encode(b"img0").decode(), + base64.b64encode(b"img1").decode(), + ] + assert all(image.url is None for image in result.data) + assert [url for url, _ in client.calls] == list(IMAGE_URLS) + + +def test_unfinished_prediction_is_polled_until_output(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(time, "sleep", lambda _: None) + poll_url = f"{API_BASE}/predictions/p1" + client = _FakeClient( + httpx.Response( + 200, json={"id": "p1", "status": "processing", "output": None}, request=httpx.Request("GET", poll_url) + ), + httpx.Response( + 200, + json={"id": "p1", "status": "succeeded", "output": [IMAGE_URLS[0]]}, + request=httpx.Request("GET", poll_url), + ), + ) + pending = _prediction_response({"id": "p1", "status": "starting", "output": None, "urls": {"get": poll_url}}) + result = _transform(_config(client), pending) + assert [image.url for image in result.data] == [IMAGE_URLS[0]] + assert [url for url, _ in client.calls] == [poll_url, poll_url] + poll_headers = client.calls[0][1] + assert poll_headers is not None + assert poll_headers["authorization"] == "Bearer k" + assert "prefer" not in poll_headers + + +def test_partial_output_keeps_polling_until_terminal_status(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(time, "sleep", lambda _: None) + poll_url = f"{API_BASE}/predictions/p1" + client = _FakeClient( + httpx.Response( + 200, + json={"id": "p1", "status": "succeeded", "output": list(IMAGE_URLS)}, + request=httpx.Request("GET", poll_url), + ), + ) + partial = _prediction_response( + {"id": "p1", "status": "processing", "output": [IMAGE_URLS[0]], "urls": {"get": poll_url}} + ) + result = _transform(_config(client), partial) + assert [image.url for image in result.data] == list(IMAGE_URLS) + + +def test_polling_gives_up_after_retry_budget(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(time, "sleep", lambda _: None) + monkeypatch.setattr(common_utils, "DEFAULT_REPLICATE_SYNC_POLLING_RETRIES", 2) + poll_url = f"{API_BASE}/predictions/p1" + still_running = {"id": "p1", "status": "processing", "output": None} + client = _FakeClient( + *(httpx.Response(200, json=still_running, request=httpx.Request("GET", poll_url)) for _ in range(2)) + ) + pending = _prediction_response({"id": "p1", "status": "starting", "output": None, "urls": {"get": poll_url}}) + with pytest.raises(ReplicateError, match="did not finish after 2 polls"): + _transform(_config(client), pending) + + +def test_poll_http_error_raises_instead_of_retrying(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(time, "sleep", lambda _: None) + poll_url = f"{API_BASE}/predictions/p1" + client = _FakeClient( + httpx.Response(401, json={"detail": "Unauthenticated"}, request=httpx.Request("GET", poll_url)), + httpx.Response(200, json={"id": "p1", "status": "succeeded", "output": IMAGE_URLS[0]}), + ) + pending = _prediction_response({"id": "p1", "status": "starting", "output": None, "urls": {"get": poll_url}}) + with pytest.raises(ReplicateError, match="Unauthenticated") as excinfo: + _transform(_config(client), pending) + assert excinfo.value.status_code == 401 + assert len(client.calls) == 1 + + +def test_polling_client_gets_deployment_ssl_verify(): + seen: list[object] = [] + + def factory(ssl_verify: object) -> _FakeClient: + seen.append(ssl_verify) + return _FakeClient() + + config = ReplicateImageGenerationConfig(sync_client_factory=factory) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + config.transform_image_generation_response( + model=MODEL, + raw_response=_prediction_response({"id": "p1", "status": "succeeded", "output": [IMAGE_URLS[0]]}), + model_response=litellm.ImageResponse(), + logging_obj=Mock(), + request_data={}, + optional_params={}, + litellm_params={"ssl_verify": "/etc/ssl/gateway.pem"}, + encoding=None, + ) + assert seen == ["/etc/ssl/gateway.pem"] + + +def test_output_urls_keep_nested_order_without_recursion(): + nested = [[IMAGE_URLS[0]], "not a url", {"files": [IMAGE_URLS[1], {"deep": "https://cdn.example/3.png"}]}, 7] + assert common_utils.replicate_output_urls(nested) == (IMAGE_URLS[0], IMAGE_URLS[1], "https://cdn.example/3.png") + assert common_utils.replicate_output_urls("plain text") == () + + +def test_request_flattens_extra_body_into_native_input(): + body = ReplicateImageGenerationConfig().transform_image_generation_request( + MODEL, + "a boat", + { + "num_outputs": 1, + "extra_body": {"seed": 7, "num_outputs": 3, "prompt": "unaudited"}, + "response_format": "url", + }, + {}, + {}, + ) + assert body == {"input": {"prompt": "a boat", "num_outputs": 3, "seed": 7}} + + +def test_poll_url_must_share_the_api_origin(): + pending = _prediction_response( + {"id": "p1", "status": "starting", "output": None, "urls": {"get": "https://evil.example/predictions/p1"}} + ) + with pytest.raises(SSRFError, match="Origin mismatch"): + _transform(_config(_FakeClient()), pending) + + +def test_b64_download_blocks_internal_urls(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "user_url_validation", True) + done = _prediction_response({"id": "p1", "status": "succeeded", "output": ["http://10.0.0.8/secret.png"]}) + with pytest.raises(SSRFError): + _transform(_config(_FakeClient()), done, {"response_format": "b64_json"}) + + +def test_failed_prediction_raises_with_replicate_error_text(): + with pytest.raises(ReplicateError, match="NSFW content detected"): + _transform( + _config(_FakeClient()), + _prediction_response({"id": "p1", "status": "failed", "error": "NSFW content detected"}), + ) + + +@pytest.mark.asyncio +async def test_async_transform_downloads_b64_json_with_async_client(): + client = _FakeAsyncClient(_download(IMAGE_URLS[0], b"img0")) + config = ReplicateImageGenerationConfig(async_client_factory=lambda _ssl_verify: client) # pyright: ignore[reportArgumentType] # fake stands in for AsyncHTTPHandler + result = await config.async_transform_image_generation_response( + model=MODEL, + raw_response=_prediction_response({"id": "p1", "status": "succeeded", "output": [IMAGE_URLS[0]]}), + model_response=litellm.ImageResponse(), + logging_obj=Mock(), + request_data={}, + optional_params={"response_format": "b64_json"}, + litellm_params={}, + encoding=None, + ) + assert [image.b64_json for image in result.data] == [base64.b64encode(b"img0").decode()] + + +def test_native_width_and_height_become_the_response_size(): + result = _transform( + _config(_FakeClient()), + _prediction_response( + { + "id": "p1", + "status": "succeeded", + "output": IMAGE_URLS[0], + "input": {"prompt": "x", "aspect_ratio": "custom", "width": 2048, "height": 1536}, + } + ), + ) + assert result.size == "2048x1536" + + +@pytest.mark.parametrize( + "prediction_input", + [ + {"prompt": "x"}, + {"prompt": "x", "width": 1024}, + {"prompt": "x", "width": "1024", "height": "768"}, + {"width": 0, "height": 768}, + ], +) +def test_response_size_stays_unset_without_both_native_dimensions(prediction_input: dict[str, object]): + result = _transform( + _config(_FakeClient()), + _prediction_response({"id": "p1", "status": "succeeded", "output": IMAGE_URLS[0], "input": prediction_input}), + ) + assert result.size is None + + +def test_b64_download_is_bounded_by_the_output_byte_limit(): + seen: list[object] = [] + + class _LimitAwareClient(_FakeClient): + def get( + self, url: str, headers: dict[str, str] | None = None, params: object = None, **kwargs: object + ) -> httpx.Response: + seen.append(kwargs.get("max_response_bytes")) + return super().get(url, headers=headers, params=params, **kwargs) + + client = _LimitAwareClient(_download(IMAGE_URLS[0], b"png")) + result = _transform( + _config(client), + _prediction_response({"id": "p1", "status": "succeeded", "output": IMAGE_URLS[0]}), + {"response_format": "b64_json"}, + ) + assert result.data[0].b64_json == base64.b64encode(b"png").decode() + assert seen == [REPLICATE_MAX_OUTPUT_DOWNLOAD_BYTES] + + +def test_oversized_b64_download_is_a_replicate_error(): + class _OverflowingClient(_FakeClient): + def get( + self, url: str, headers: dict[str, str] | None = None, params: object = None, **kwargs: object + ) -> httpx.Response: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + + with pytest.raises(ReplicateError, match="REPLICATE_MAX_OUTPUT_DOWNLOAD_BYTES") as refused: + _transform( + _config(_OverflowingClient()), + _prediction_response({"id": "p1", "status": "succeeded", "output": IMAGE_URLS[0]}), + {"response_format": "b64_json"}, + ) + assert refused.value.status_code == 502 + + +@pytest.mark.parametrize( + "model", + [ + "black-forest-labs/../xai/grok-imagine-image", + "deployments/acme/../../models/xai/grok-imagine-image", + "black-forest-labs//flux-schnell", + "./flux-schnell", + ], +) +def test_model_paths_with_dot_or_empty_segments_are_rejected(model: str): + with pytest.raises(ReplicateError) as refused: + common_utils.replicate_prediction_request("https://api.replicate.com/v1", model, {}) + assert refused.value.status_code == 400 + + +def test_model_path_segments_are_url_encoded(): + request = common_utils.replicate_prediction_request("https://api.replicate.com/v1", "acme/odd name?x", {}) + assert request.url == "https://api.replicate.com/v1/models/acme/odd%20name%3Fx/predictions" + + +def test_a_pinned_version_stays_under_the_authorized_model_path(): + request = common_utils.replicate_prediction_request("https://api.replicate.com/v1", "acme/fox:ab12", {}) + assert request.url == "https://api.replicate.com/v1/models/acme/fox/versions/ab12/predictions" + assert "version" not in request.body + with pytest.raises(ReplicateError) as refused: + common_utils.replicate_prediction_request("https://api.replicate.com/v1", "acme/fox:", {}) + assert refused.value.status_code == 400 diff --git a/tests/unit/llms/replicate/text_to_speech/__init__.py b/tests/unit/llms/replicate/text_to_speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/replicate/text_to_speech/test_replicate_text_to_speech_transformation.py b/tests/unit/llms/replicate/text_to_speech/test_replicate_text_to_speech_transformation.py new file mode 100644 index 00000000000..364986d34d4 --- /dev/null +++ b/tests/unit/llms/replicate/text_to_speech/test_replicate_text_to_speech_transformation.py @@ -0,0 +1,219 @@ +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest + +import litellm +from litellm.llms.replicate.common_utils import ReplicateError +from litellm.llms.replicate.text_to_speech.transformation import ReplicateTextToSpeechConfig +from litellm.types.llms.openai import HttpxBinaryResponseContent +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +API_BASE: Final = "https://api.replicate.com/v1" +MODEL: Final = "xai/grok-text-to-speech" +PREDICTIONS_URL: Final = f"{API_BASE}/models/{MODEL}/predictions" +AUDIO_URL: Final = "https://replicate.delivery/out/speech.mp3" + + +@pytest.fixture(autouse=True) +def _skip_url_dns_validation(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + + +class _FakeClient: + def __init__(self, *responses: httpx.Response) -> None: + self.responses = list(responses) + self.calls: list[str] = [] + + def get( + self, url: str, headers: dict[str, str] | None = None, params: object = None, **kwargs: object + ) -> httpx.Response: + self.calls.append(url) + return self.responses.pop(0) + + +class _FakeAsyncClient(_FakeClient): + async def get( + self, url: str, headers: dict[str, str] | None = None, params: object = None, **kwargs: object + ) -> httpx.Response: # pyright: ignore[reportIncompatibleMethodOverride] # async twin of the sync fake + return super().get(url, headers=headers, params=params, **kwargs) + + +def _prediction_response(body: dict[str, object]) -> httpx.Response: + request = httpx.Request("POST", PREDICTIONS_URL, headers={"Authorization": "Bearer k", "Prefer": "wait"}) + return httpx.Response(201, json=body, request=request) + + +def _audio(content: bytes, content_type: str = "audio/mpeg") -> httpx.Response: + return httpx.Response( + 200, content=content, headers={"content-type": content_type}, request=httpx.Request("GET", AUDIO_URL) + ) + + +def test_provider_config_is_registered(): + config = ProviderConfigManager.get_provider_text_to_speech_config(model=MODEL, provider=LlmProviders.REPLICATE) + assert isinstance(config, ReplicateTextToSpeechConfig) + + +def test_map_openai_params_keeps_voice_and_speed_only(): + voice, params = ReplicateTextToSpeechConfig().map_openai_params( + MODEL, {"speed": 1.2, "response_format": "mp3", "instructions": "cheerful"}, voice="ara" + ) + assert voice == "ara" + assert params == {"speed": 1.2} + assert ReplicateTextToSpeechConfig().map_openai_params(MODEL, {}, voice=None) == (None, {}) + + +def test_map_openai_params_passes_extra_body_as_native_input(): + _, params = ReplicateTextToSpeechConfig().map_openai_params( + MODEL, {"speed": 0.9}, voice="ara", kwargs={"extra_body": {"output_format": "wav", "speed": 2}} + ) + assert params == {"output_format": "wav", "speed": 0.9} + _, ignored = ReplicateTextToSpeechConfig().map_openai_params(MODEL, {}, voice=None, kwargs={"extra_body": "x"}) + assert ignored == {} + + +def test_validate_environment_sets_bearer_and_prefer_wait(): + headers = ReplicateTextToSpeechConfig().validate_environment({}, MODEL, api_key="r8_key") + assert headers["Authorization"] == "Bearer r8_key" + assert headers["Prefer"] == "wait" + + +def test_get_complete_url(): + config = ReplicateTextToSpeechConfig() + assert config.get_complete_url(MODEL, None, {}) == PREDICTIONS_URL + assert ( + config.get_complete_url(MODEL, "https://proxy.internal", {}) + == f"https://proxy.internal/v1/models/{MODEL}/predictions" + ) + + +@pytest.mark.parametrize( + ("model", "voice_key"), + [ + (MODEL, "voice"), + ("minimax/speech-02-turbo", "voice_id"), + ("inworld/realtime-tts-2", "voice_id"), + ("qwen/qwen3-tts", "speaker"), + ("deployments/minimax/my-tts", "voice_id"), + ("qwen/qwen3-tts:0123abcd", "speaker"), + ], +) +def test_request_body_uses_model_family_voice_key(model: str, voice_key: str): + request = ReplicateTextToSpeechConfig().transform_text_to_speech_request( + model, "hello there", "Wise_Woman", {"speed": 1.1}, {}, {} + ) + assert request["dict_body"]["input"] == {"text": "hello there", voice_key: "Wise_Woman", "speed": 1.1} + + +def test_request_body_without_voice(): + request = ReplicateTextToSpeechConfig().transform_text_to_speech_request(MODEL, "hello", None, {}, {}, {}) + assert request == {"dict_body": {"input": {"text": "hello"}}} + + +def test_request_input_wins_over_extra_body_text(): + _, params = ReplicateTextToSpeechConfig().map_openai_params( + MODEL, {}, voice="ara", kwargs={"extra_body": {"text": "a much longer text that is not billed"}} + ) + request = ReplicateTextToSpeechConfig().transform_text_to_speech_request(MODEL, "hello", "ara", params, {}, {}) + assert request["dict_body"]["input"] == {"text": "hello", "voice": "ara"} + + +def test_response_downloads_audio_from_output_url(): + client = _FakeClient(_audio(b"mp3-bytes")) + config = ReplicateTextToSpeechConfig(sync_client_factory=lambda _ssl: client) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + result = config.transform_text_to_speech_response( + MODEL, _prediction_response({"id": "p1", "status": "succeeded", "output": AUDIO_URL}), Mock() + ) + assert isinstance(result, HttpxBinaryResponseContent) + assert result.content == b"mp3-bytes" + assert result.response.headers["content-type"] == "audio/mpeg" + assert client.calls == [AUDIO_URL] + + +@pytest.mark.asyncio +async def test_async_response_downloads_audio_with_async_client(): + client = _FakeAsyncClient(_audio(b"wav-bytes", "audio/wav")) + config = ReplicateTextToSpeechConfig(async_client_factory=lambda _ssl: client) # pyright: ignore[reportArgumentType] # fake stands in for AsyncHTTPHandler + result = await config.async_transform_text_to_speech_response( + MODEL, _prediction_response({"id": "p1", "status": "succeeded", "output": [AUDIO_URL]}), Mock() + ) + assert result.content == b"wav-bytes" + assert result.response.headers["content-type"] == "audio/wav" + + +def test_response_client_uses_the_request_ssl_verify(): + seen: list[object] = [] + + def factory(ssl_verify: object) -> _FakeClient: + seen.append(ssl_verify) + return _FakeClient(_audio(b"mp3-bytes")) + + config = ReplicateTextToSpeechConfig(sync_client_factory=factory) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + logging_obj = Mock() + logging_obj.litellm_params = {"ssl_verify": "/certs/gateway.pem"} + config.transform_text_to_speech_response( + MODEL, _prediction_response({"id": "p1", "status": "succeeded", "output": AUDIO_URL}), logging_obj + ) + assert seen == ["/certs/gateway.pem"] + + +def test_failed_prediction_raises(): + config = ReplicateTextToSpeechConfig(sync_client_factory=lambda _ssl: _FakeClient()) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + with pytest.raises(ReplicateError, match="text too long"): + config.transform_text_to_speech_response( + MODEL, _prediction_response({"id": "p1", "status": "failed", "error": "text too long"}), Mock() + ) + with pytest.raises(ReplicateError, match="without audio output"): + config.transform_text_to_speech_response( + MODEL, _prediction_response({"id": "p1", "status": "succeeded", "output": None}), Mock() + ) + + +def test_response_usage_comes_from_prediction_token_metrics(): + config = ReplicateTextToSpeechConfig(sync_client_factory=lambda _ssl: _FakeClient(_audio(b"mp3-bytes"))) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + result = config.transform_text_to_speech_response( + MODEL, + _prediction_response( + { + "id": "p1", + "status": "succeeded", + "output": AUDIO_URL, + "metrics": {"predict_time": 2.3, "token_input_count": 8, "token_output_count": 32}, + } + ), + Mock(), + ) + assert result.usage is not None + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (8, 32, 40) + assert result.logging_summary()["usage"]["completion_tokens"] == 32 + + +@pytest.mark.parametrize( + "metrics", + [ + {"predict_time": 2.3, "token_input_count": 8, "token_output_count": 32}, + {"predict_time": 2.3, "input_token_count": 8, "output_token_count": 32}, + ], +) +def test_response_usage_accepts_both_replicate_metric_spellings(metrics: dict[str, object]): + config = ReplicateTextToSpeechConfig(sync_client_factory=lambda _ssl: _FakeClient(_audio(b"mp3-bytes"))) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + result = config.transform_text_to_speech_response( + MODEL, + _prediction_response({"id": "p1", "status": "succeeded", "output": AUDIO_URL, "metrics": metrics}), + Mock(), + ) + assert result.usage is not None + assert (result.usage.prompt_tokens, result.usage.completion_tokens) == (8, 32) + + +@pytest.mark.parametrize("metrics", [None, {"predict_time": 2.3}, {"token_input_count": 0, "token_output_count": -1}]) +def test_response_without_token_metrics_has_no_usage(metrics: dict[str, object] | None): + config = ReplicateTextToSpeechConfig(sync_client_factory=lambda _ssl: _FakeClient(_audio(b"mp3-bytes"))) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + body: dict[str, object] = {"id": "p1", "status": "succeeded", "output": AUDIO_URL} + result = config.transform_text_to_speech_response( + MODEL, _prediction_response(body if metrics is None else {**body, "metrics": metrics}), Mock() + ) + assert result.usage is None diff --git a/tests/unit/llms/replicate/videos/__init__.py b/tests/unit/llms/replicate/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/replicate/videos/test_replicate_video_transformation.py b/tests/unit/llms/replicate/videos/test_replicate_video_transformation.py new file mode 100644 index 00000000000..b46860499bd --- /dev/null +++ b/tests/unit/llms/replicate/videos/test_replicate_video_transformation.py @@ -0,0 +1,350 @@ +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest + +import litellm +from litellm.cost_calculator import default_video_cost_calculator +from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.url_utils import SSRFError +from litellm.llms.replicate.common_utils import ReplicateError +from litellm.llms.replicate.videos.transformation import ReplicateVideoConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.types.videos.utils import decode_video_id_with_provider +from litellm.utils import ProviderConfigManager + +API_BASE: Final = "https://api.replicate.com/v1" +MODEL: Final = "xai/grok-imagine-video" +MINIMAX_MODEL: Final = "minimax/h3-lite" +IMAGE_URL: Final = "https://example.com/frame.png" +VIDEO_URL: Final = "https://replicate.delivery/out/video.mp4" + + +@pytest.fixture(autouse=True) +def _skip_url_dns_validation(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + + +class _FakeClient: + def __init__(self, *responses: httpx.Response) -> None: + self.responses = list(responses) + self.calls: list[str] = [] + + def get( + self, url: str, headers: dict[str, str] | None = None, params: object = None, **kwargs: object + ) -> httpx.Response: + self.calls.append(url) + return self.responses.pop(0) + + +class _FakeAsyncClient(_FakeClient): + async def get( + self, url: str, headers: dict[str, str] | None = None, params: object = None, **kwargs: object + ) -> httpx.Response: # pyright: ignore[reportIncompatibleMethodOverride] # async twin of the sync fake + return super().get(url, headers=headers, params=params, **kwargs) + + +def _prediction_response(body: dict[str, object], url: str = f"{API_BASE}/predictions/abc") -> httpx.Response: + return httpx.Response(200, json=body, request=httpx.Request("GET", url, headers={"Authorization": "Bearer k"})) + + +def test_provider_config_is_registered(): + config = ProviderConfigManager.get_provider_video_config(model=MODEL, provider=LlmProviders.REPLICATE) + assert isinstance(config, ReplicateVideoConfig) + assert isinstance(config.get_error_class("bad key", 401, {}), ReplicateError) + + +def test_map_openai_params_generic_model(): + mapped = ReplicateVideoConfig().map_openai_params( + {"seconds": "6", "size": "1280x720", "input_reference": IMAGE_URL, "user": "u1", "seed": 7}, + MODEL, + False, + ) + assert mapped == {"image": IMAGE_URL, "duration": 6, "resolution": "720p", "aspect_ratio": "16:9", "seed": 7} + assert isinstance(mapped["duration"], int) + + +@pytest.mark.parametrize("model", [MINIMAX_MODEL, f"deployments/{MINIMAX_MODEL}", f"{MINIMAX_MODEL}:0123abcd"]) +def test_map_openai_params_minimax_profile(model: str): + mapped = ReplicateVideoConfig().map_openai_params( + {"seconds": 5, "size": "1024x768", "input_reference": IMAGE_URL}, model, False + ) + assert mapped == {"first_frame_image": IMAGE_URL, "duration": 5, "resolution": "768P", "ratio": "4:3"} + + +@pytest.mark.parametrize( + ("size", "expected"), + [("720p", {"resolution": "720p"}), ("1080P", {"resolution": "1080p"}), ("4k", {"resolution": "4k"})], +) +def test_map_openai_params_named_sizes(size: str, expected: dict[str, str]): + assert ReplicateVideoConfig().map_openai_params({"size": size}, MODEL, False) == expected + + +def test_map_openai_params_fractional_seconds_and_bad_values(): + assert ReplicateVideoConfig().map_openai_params({"seconds": "2.5"}, MODEL, False) == {"duration": 2.5} + with pytest.raises(ValueError, match="numeric"): + ReplicateVideoConfig().map_openai_params({"seconds": "auto"}, MODEL, False) + with pytest.raises(ValueError, match="public image URL"): + ReplicateVideoConfig().map_openai_params({"input_reference": b"bytes"}, MODEL, False) + + +def test_map_openai_params_unsupported_native_fields(): + with pytest.raises(UnsupportedParamsError, match="frame_images"): + ReplicateVideoConfig().map_openai_params({"frame_images": [{"url": IMAGE_URL}]}, MODEL, False) + assert ReplicateVideoConfig().map_openai_params( + {"frame_images": [{"url": IMAGE_URL}], "seconds": 5}, MODEL, True + ) == {"duration": 5} + + +def test_validate_environment_requires_key_and_sends_bearer(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("REPLICATE_API_KEY", raising=False) + monkeypatch.delenv("REPLICATE_API_TOKEN", raising=False) + monkeypatch.setattr(litellm, "replicate_key", None) + with pytest.raises(ReplicateError, match="REPLICATE_API_KEY"): + ReplicateVideoConfig().validate_environment({}, MODEL) + headers = ReplicateVideoConfig().validate_environment({"X-Trace": "1"}, MODEL, api_key="r8_key") + assert headers["Authorization"] == "Bearer r8_key" + assert headers["X-Trace"] == "1" + via_params = ReplicateVideoConfig().validate_environment( + {}, MODEL, litellm_params=GenericLiteLLMParams(api_key="r8_p") + ) + assert via_params["Authorization"] == "Bearer r8_p" + + +@pytest.mark.parametrize( + ("model", "url", "extra_body"), + [ + (MODEL, f"{API_BASE}/models/{MODEL}/predictions", {}), + ("owner/name:abc123", f"{API_BASE}/models/owner/name/versions/abc123/predictions", {}), + ("deployments/acme/fast-video", f"{API_BASE}/deployments/acme/fast-video/predictions", {}), + ], +) +def test_transform_video_create_request(model: str, url: str, extra_body: dict[str, str]): + body, files, request_url = ReplicateVideoConfig().transform_video_create_request( + model=model, + prompt="a quiet ocean", + api_base=API_BASE, + video_create_optional_request_params={"duration": 5, "resolution": "720p", "model": model}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert request_url == url + assert files == [] + assert body == {**extra_body, "input": {"prompt": "a quiet ocean", "duration": 5, "resolution": "720p"}} + + +def test_transform_video_create_request_keeps_the_inspected_prompt(): + body, _, _ = ReplicateVideoConfig().transform_video_create_request( + model="minimax/h3-lite", + prompt="a quiet ocean", + api_base=API_BASE, + video_create_optional_request_params={"duration": 5, "prompt": "unaudited"}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert body["input"] == {"duration": 5, "prompt": "a quiet ocean"} + + +def test_get_complete_url_defaults_and_override(): + config = ReplicateVideoConfig() + assert config.get_complete_url(MODEL, None, {}) == API_BASE + assert config.get_complete_url(MODEL, "https://proxy.internal/", {}) == "https://proxy.internal/v1" + assert config.get_complete_url(MODEL, "https://proxy.internal/v1", {}) == "https://proxy.internal/v1" + + +def test_transform_video_create_response_encodes_id_and_usage(): + response = _prediction_response( + {"id": "abc", "status": "starting", "input": {"prompt": "x", "duration": 6, "resolution": "480p"}} + ) + video = ReplicateVideoConfig().transform_video_create_response(MODEL, response, Mock(), "replicate") + decoded = decode_video_id_with_provider(video.id) + assert decoded["custom_llm_provider"] == "replicate" + assert decoded["model_id"] == MODEL + assert decoded["video_id"] == "abc" + assert video.status == "queued" + assert video.model == MODEL + assert video.seconds == "6" + assert video.size == "480p" + assert video.usage == {"duration_seconds": 6.0, "video_resolution": "480p"} + + +def test_status_request_uses_prediction_id_from_encoded_video_id(): + create = _prediction_response({"id": "abc", "status": "starting"}) + video = ReplicateVideoConfig().transform_video_create_response(MODEL, create, Mock()) + url, data = ReplicateVideoConfig().transform_video_status_retrieve_request( + video.id, API_BASE, GenericLiteLLMParams(), {} + ) + assert url == f"{API_BASE}/predictions/abc#{MODEL}" + assert httpx.URL(url).raw_path == b"/v1/predictions/abc" + assert data == {} + content_url, _ = ReplicateVideoConfig().transform_video_content_request( + video.id, API_BASE, GenericLiteLLMParams(), {} + ) + assert content_url == url + + +@pytest.mark.parametrize("model", [f"{MODEL}:0123abcd", "deployments/acme/fast-video"]) +def test_status_response_keeps_the_requested_model_and_video_id(model: str): + create = _prediction_response({"id": "abc", "status": "starting", "model": "xai/grok-imagine-video"}) + video = ReplicateVideoConfig().transform_video_create_response(model, create, Mock()) + url, _ = ReplicateVideoConfig().transform_video_status_retrieve_request( + video.id, API_BASE, GenericLiteLLMParams(), {} + ) + status = httpx.Response( + 200, + json={"id": "abc", "status": "succeeded", "model": "xai/grok-imagine-video", "output": VIDEO_URL}, + request=httpx.Request("GET", url), + ) + polled = ReplicateVideoConfig().transform_video_status_retrieve_response(status, Mock()) + assert polled.id == video.id + assert polled.model == model + + +@pytest.mark.parametrize( + ("replicate_status", "expected"), + [("starting", "queued"), ("processing", "in_progress"), ("succeeded", "completed")], +) +def test_status_response_maps_statuses(replicate_status: str, expected: str): + response = _prediction_response( + {"id": "abc", "status": replicate_status, "model": MODEL, "input": {"duration": 5, "resolution": "720p"}} + ) + video = ReplicateVideoConfig().transform_video_status_retrieve_response(response, Mock(), "replicate") + assert video.status == expected + assert video.error is None + assert video.model == MODEL + assert video.usage == {"duration_seconds": 5.0, "video_resolution": "720p"} + assert decode_video_id_with_provider(video.id)["video_id"] == "abc" + + +@pytest.mark.parametrize("replicate_status", ["failed", "canceled"]) +def test_status_response_surfaces_failures(replicate_status: str): + response = _prediction_response({"id": "abc", "status": replicate_status, "error": "NSFW content detected"}) + video = ReplicateVideoConfig().transform_video_status_retrieve_response(response, Mock()) + assert video.status == "failed" + assert video.error == {"code": "replicate_error", "message": "NSFW content detected"} + + +def test_status_response_without_duration_has_no_usage(): + response = _prediction_response({"id": "abc", "status": "succeeded", "input": {"prompt": "x"}}) + assert ReplicateVideoConfig().transform_video_status_retrieve_response(response, Mock()).usage == {} + + +@pytest.mark.parametrize("output", [VIDEO_URL, [VIDEO_URL], {"video": VIDEO_URL}]) +def test_content_response_downloads_first_output(output: object): + download = httpx.Response(200, content=b"mp4-bytes", request=httpx.Request("GET", VIDEO_URL)) + client = _FakeClient(download) + config = ReplicateVideoConfig(sync_client_factory=lambda: client) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + response = _prediction_response({"id": "abc", "status": "succeeded", "output": output}) + assert config.transform_video_content_response(response, Mock()) == b"mp4-bytes" + assert client.calls == [VIDEO_URL] + + +@pytest.mark.asyncio +async def test_async_content_response_downloads_first_output(): + download = httpx.Response(200, content=b"mp4-bytes", request=httpx.Request("GET", VIDEO_URL)) + client = _FakeAsyncClient(download) + config = ReplicateVideoConfig(async_client_factory=lambda: client) # pyright: ignore[reportArgumentType] # fake stands in for AsyncHTTPHandler + response = _prediction_response({"id": "abc", "status": "succeeded", "output": [VIDEO_URL]}) + assert await config.async_transform_video_content_response(response, Mock()) == b"mp4-bytes" + + +def test_content_response_prefers_the_handler_client(): + download = httpx.Response(200, content=b"mp4-bytes", request=httpx.Request("GET", VIDEO_URL)) + handler_client = _FakeClient(download) + config = ReplicateVideoConfig(sync_client_factory=lambda: _FakeClient()) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + response = _prediction_response({"id": "abc", "status": "succeeded", "output": VIDEO_URL}) + assert config.transform_video_content_response(response, Mock(), client=handler_client) == b"mp4-bytes" # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + assert handler_client.calls == [VIDEO_URL] + + +@pytest.mark.parametrize( + ("requested_model", "prediction"), + [ + ( + "acme/other-video", + {"id": "abc", "status": "succeeded", "model": "xai/grok-imagine-video", "output": VIDEO_URL}, + ), + (f"{MODEL}:v2", {"id": "abc", "status": "succeeded", "model": MODEL, "version": "v1", "output": VIDEO_URL}), + ], +) +def test_retrieval_rejects_a_prediction_from_another_model(requested_model: str, prediction: dict[str, object]): + create = _prediction_response({"id": "abc", "status": "starting"}) + video = ReplicateVideoConfig().transform_video_create_response(requested_model, create, Mock()) + config = ReplicateVideoConfig(sync_client_factory=lambda: _FakeClient()) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + status_url, _ = config.transform_video_status_retrieve_request(video.id, API_BASE, GenericLiteLLMParams(), {}) + foreign = httpx.Response(200, json=prediction, request=httpx.Request("GET", status_url)) + with pytest.raises(ReplicateError, match="does not belong to") as excinfo: + config.transform_video_status_retrieve_response(foreign, Mock()) + assert excinfo.value.status_code == 404 + with pytest.raises(ReplicateError, match="does not belong to"): + config.transform_video_content_response(foreign, Mock()) + + +def test_retrieval_accepts_deployment_predictions_without_a_model_check(): + create = _prediction_response({"id": "abc", "status": "starting"}) + video = ReplicateVideoConfig().transform_video_create_response("deployments/acme/fast", create, Mock()) + status_url, _ = ReplicateVideoConfig().transform_video_status_retrieve_request( + video.id, API_BASE, GenericLiteLLMParams(), {} + ) + status = httpx.Response( + 200, + json={"id": "abc", "status": "succeeded", "model": "xai/grok-imagine-video"}, + request=httpx.Request("GET", status_url), + ) + assert ( + ReplicateVideoConfig().transform_video_status_retrieve_response(status, Mock()).model == "deployments/acme/fast" + ) + + +def test_content_download_blocks_internal_urls(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "user_url_validation", True) + config = ReplicateVideoConfig(sync_client_factory=lambda: _FakeClient()) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + response = _prediction_response( + {"id": "abc", "status": "succeeded", "output": "http://169.254.169.254/latest/meta-data"} + ) + with pytest.raises(SSRFError): + config.transform_video_content_response(response, Mock()) + + +def test_content_response_errors_without_output(): + config = ReplicateVideoConfig(sync_client_factory=lambda: _FakeClient()) # pyright: ignore[reportArgumentType] # fake stands in for HTTPHandler + with pytest.raises(ReplicateError, match="no video output yet"): + config.transform_video_content_response(_prediction_response({"id": "abc", "status": "processing"}), Mock()) + with pytest.raises(ReplicateError, match="boom"): + config.transform_video_content_response( + _prediction_response({"id": "abc", "status": "failed", "error": "boom"}), Mock() + ) + + +def test_completed_status_usage_prices_through_cost_map(local_model_cost_map: None): + row = litellm.model_cost[f"replicate/{MINIMAX_MODEL}"] + response = _prediction_response( + {"id": "abc", "status": "succeeded", "model": MINIMAX_MODEL, "input": {"duration": 6, "resolution": "480P"}} + ) + usage = ReplicateVideoConfig().transform_video_status_retrieve_response(response, Mock()).usage + assert usage is not None + cost = default_video_cost_calculator( + f"replicate/{MINIMAX_MODEL}", + usage["duration_seconds"], + "replicate", + video_resolution=usage["video_resolution"], + ) + assert cost == pytest.approx(6 * row["output_cost_per_second_480p"]) + + +def test_video_cost_uses_tiered_replicate_rows(local_model_cost_map: None): + rows = { + model: row + for model, row in litellm.model_cost.items() + if row.get("litellm_provider") == "replicate" and row.get("mode") == "video_generation" + } + assert rows + for model, row in rows.items(): + for key, value in row.items(): + if key.startswith("output_cost_per_second_") and value is not None: + tier = key.removeprefix("output_cost_per_second_") + assert default_video_cost_calculator(model, 5, "replicate", video_resolution=tier) == pytest.approx( + 5 * value + )