mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
feat(replicate): add video generation, image generation and text to speech through the predictions API (#51360)
* feat(replicate): add video generation, image generation and text to speech through the predictions API * fix(replicate): surface poll HTTP errors, keep deployment ssl_verify and native speech input, flatten outputs without recursion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(replicate): keep model family mapping for deployments and versions, preserve video id on status, poll until terminal status Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(replicate): validate output downloads and poll origin, bind retrieval to the requested model, flatten image extra_body, keep speech input authoritative, use the async image path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(replicate): keep the audited prompt authoritative over image extra_body Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(replicate): cover extra_body.prompt not replacing the audited prompt Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(replicate): honor the request ssl_verify when polling and downloading speech output Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(replicate): build the video usage into VideoObject instead of assigning it afterwards Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(replicate): keep the inspected video prompt authoritative over native extra_body prompt Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(replicate): bill native image dimensions and token-metered speech output Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(replicate): bound output downloads, accept both token metric spellings and build speech usage at construction Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(http_handler): walk bounded-download redirects iteratively Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(replicate): reject dot and empty model path segments and keep the bounded redirect walk annotated Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(http_handler): drop the unused rebind suppression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(http_handler): put the rebind suppression on the redirect hop binding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(replicate): run pinned versions through the model's own versions endpoint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu
GitHub
parent
482c8eb1e5
commit
e3b32534dc
21 changed files with 2373 additions and 37 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
228
litellm/llms/replicate/image_generation/transformation.py
Normal file
228
litellm/llms/replicate/image_generation/transformation.py
Normal file
|
|
@ -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)
|
||||
204
litellm/llms/replicate/text_to_speech/transformation.py
Normal file
204
litellm/llms/replicate/text_to_speech/transformation.py
Normal file
|
|
@ -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)
|
||||
486
litellm/llms/replicate/videos/transformation.py
Normal file
486
litellm/llms/replicate/videos/transformation.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
0
tests/unit/llms/replicate/image_generation/__init__.py
Normal file
0
tests/unit/llms/replicate/image_generation/__init__.py
Normal file
|
|
@ -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
|
||||
0
tests/unit/llms/replicate/text_to_speech/__init__.py
Normal file
0
tests/unit/llms/replicate/text_to_speech/__init__.py
Normal file
|
|
@ -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
|
||||
0
tests/unit/llms/replicate/videos/__init__.py
Normal file
0
tests/unit/llms/replicate/videos/__init__.py
Normal file
|
|
@ -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
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue