mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge pull request #39839 from BerriAI/litellm_async_remote_image_fetch
fix(async): move remote image fetches off the event loop for Snowflake, Bedrock invoke Claude, Mantle and Gemini
This commit is contained in:
Mateo Wang
GitHub
commit
9700f666d0
21 changed files with 1552 additions and 217 deletions
|
|
@ -2,7 +2,11 @@
|
|||
Helper functions to handle images passed in messages
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from httpx import Response
|
||||
|
|
@ -11,9 +15,11 @@ import litellm
|
|||
from litellm import verbose_logger
|
||||
from litellm.caching.caching import InMemoryCache
|
||||
from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB
|
||||
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
MAX_IMGS_IN_MEMORY: Final = 10
|
||||
MAX_CONCURRENT_REMOTE_MEDIA_FETCHES: Final = 20
|
||||
|
||||
in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY)
|
||||
|
||||
|
|
@ -72,6 +78,14 @@ def _process_image_response(response: Response, url: str) -> str:
|
|||
return result
|
||||
|
||||
|
||||
def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError":
|
||||
verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict)
|
||||
return litellm.ImageFetchError(
|
||||
"Error: Unable to fetch image from URL. The proxy could not resolve this host or its URL policy rejected it; "
|
||||
f"an admin can check the proxy log and `user_url_allowed_hosts` in general_settings. url={url}"
|
||||
)
|
||||
|
||||
|
||||
async def async_convert_url_to_base64(url: str) -> str:
|
||||
if url.startswith("data:") and ";base64," in url:
|
||||
return url
|
||||
|
|
@ -93,6 +107,8 @@ async def async_convert_url_to_base64(url: str) -> str:
|
|||
return _process_image_response(response, url)
|
||||
except litellm.ImageFetchError:
|
||||
raise
|
||||
except SSRFError as e:
|
||||
raise _rejected_image_fetch(url, e) from e
|
||||
except Exception:
|
||||
pass
|
||||
raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}")
|
||||
|
|
@ -119,8 +135,192 @@ def convert_url_to_base64(url: str) -> str:
|
|||
return _process_image_response(response, url)
|
||||
except litellm.ImageFetchError:
|
||||
raise
|
||||
except SSRFError as e:
|
||||
raise _rejected_image_fetch(url, e) from e
|
||||
except Exception as e:
|
||||
verbose_logger.exception(e)
|
||||
raise litellm.ImageFetchError(
|
||||
f"Error: Unable to fetch image from URL after 3 attempts. url={url}",
|
||||
)
|
||||
|
||||
|
||||
_REMOTE_URL_PREFIXES: Final = ("http://", "https://")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RemoteImage:
|
||||
part: Mapping[str, object]
|
||||
image_url: Mapping[str, object] | None
|
||||
url: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RemoteFile:
|
||||
part: Mapping[str, object]
|
||||
file: Mapping[str, object]
|
||||
url: str
|
||||
|
||||
|
||||
def _as_mapping(value: object) -> Mapping[str, object] | None:
|
||||
return value if isinstance(value, Mapping) else None # pyright: ignore[reportUnknownVariableType] # fields are parsed one by one
|
||||
|
||||
|
||||
def _remote_url(candidate: object) -> str | None:
|
||||
return candidate if isinstance(candidate, str) and candidate.startswith(_REMOTE_URL_PREFIXES) else None
|
||||
|
||||
|
||||
_ANTHROPIC_MEDIA_BLOCK_TYPES: Final = frozenset({"document", "image"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RemoteSource:
|
||||
part: Mapping[str, object]
|
||||
source: Mapping[str, object]
|
||||
url: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RemoteMedia:
|
||||
url: str
|
||||
fields: Mapping[str, object]
|
||||
|
||||
|
||||
_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def inline_every_remote_url(_media: RemoteMedia) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None:
|
||||
if fields.get("type") != "image_url":
|
||||
return None
|
||||
image_url: Final = fields.get("image_url")
|
||||
image_url_fields: Final = _as_mapping(image_url)
|
||||
url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url)
|
||||
return _RemoteImage(fields, image_url_fields, url) if url is not None else None
|
||||
|
||||
|
||||
def _parse_remote_file(fields: Mapping[str, object]) -> _RemoteFile | None:
|
||||
file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None
|
||||
url: Final = _remote_url(file.get("file_id")) if file is not None else None
|
||||
return _RemoteFile(fields, file, url) if file is not None and url is not None else None
|
||||
|
||||
|
||||
def _parse_remote_source(fields: Mapping[str, object]) -> _RemoteSource | None:
|
||||
source: Final = _as_mapping(fields.get("source")) if fields.get("type") in _ANTHROPIC_MEDIA_BLOCK_TYPES else None
|
||||
url: Final = _remote_url(source.get("url")) if source is not None and source.get("type") == "url" else None
|
||||
return _RemoteSource(fields, source, url) if source is not None and url is not None else None
|
||||
|
||||
|
||||
def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSource | None:
|
||||
fields: Final = _as_mapping(part)
|
||||
if fields is None:
|
||||
return None
|
||||
return _parse_remote_image(fields) or _parse_remote_file(fields) or _parse_remote_source(fields)
|
||||
|
||||
|
||||
def _remote_media(remote: _RemoteImage | _RemoteFile | _RemoteSource) -> RemoteMedia:
|
||||
match remote:
|
||||
case _RemoteImage(_, image_url, url):
|
||||
return RemoteMedia(url, image_url if image_url is not None else _NO_FIELDS)
|
||||
case _RemoteFile(_, file, url):
|
||||
return RemoteMedia(url, file)
|
||||
case _RemoteSource(_, source, url):
|
||||
return RemoteMedia(url, source)
|
||||
|
||||
|
||||
_PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"})
|
||||
|
||||
|
||||
def _inferred_format(file: Mapping[str, object], url: str) -> Mapping[str, str]:
|
||||
return _PDF_FORMAT if "format" not in file and url.lower().endswith(".pdf") else MappingProxyType({})
|
||||
|
||||
|
||||
def _inlined_image_url(image_url: Mapping[str, object] | None, data_url: str) -> Mapping[str, object] | str:
|
||||
return {**image_url, "url": data_url} if image_url is not None else data_url # mutable-ok: json-serialized part
|
||||
|
||||
|
||||
def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mapping[str, object]:
|
||||
kept: Final = {k: v for k, v in file.items() if k != "file_id"} # mutable-ok: json-serialized message part
|
||||
return {**kept, **_inferred_format(file, url), "file_data": data_url} # mutable-ok: json-serialized part
|
||||
|
||||
|
||||
def _base64_source(url: str, data_url: str) -> Mapping[str, str]:
|
||||
fetched_media_type, data = data_url.removeprefix("data:").split(";base64,", 1)
|
||||
media_type: Final = "application/pdf" if url.lower().endswith(".pdf") else fetched_media_type
|
||||
return {"type": "base64", "media_type": media_type, "data": data} # mutable-ok: json-serialized message part
|
||||
|
||||
|
||||
def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) -> Mapping[str, object]:
|
||||
match remote:
|
||||
case _RemoteImage(part, image_url, _):
|
||||
return {**part, "image_url": _inlined_image_url(image_url, data_url)} # mutable-ok: json-serialized part
|
||||
case _RemoteFile(part, file, url):
|
||||
return {**part, "file": _inlined_file(file, url, data_url)} # mutable-ok: json-serialized message part
|
||||
case _RemoteSource(part, _, url):
|
||||
return {**part, "source": _base64_source(url, data_url)} # mutable-ok: json-serialized message part
|
||||
|
||||
|
||||
def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]:
|
||||
content: Final = message.get("content")
|
||||
return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one
|
||||
|
||||
|
||||
def _inline_part(part: object, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]) -> object:
|
||||
remote: Final = _parse_remote_part(part)
|
||||
if remote is None or not should_inline(_remote_media(remote)):
|
||||
return part
|
||||
data_url: Final = data_urls.get(remote.url)
|
||||
return _inline(remote, data_url) if data_url is not None else part
|
||||
|
||||
|
||||
def _inline_message(
|
||||
message: AllMessageValues, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]
|
||||
) -> AllMessageValues:
|
||||
parts: Final = _content_parts(message)
|
||||
if not parts:
|
||||
return message
|
||||
inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks
|
||||
_inline_part(part, data_urls, should_inline) for part in parts
|
||||
]
|
||||
inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message
|
||||
return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined
|
||||
|
||||
|
||||
async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str:
|
||||
async with in_flight:
|
||||
return await async_convert_url_to_base64(url)
|
||||
|
||||
|
||||
async def _fetch_data_urls(remote_urls: tuple[str, ...]) -> tuple[str, ...]:
|
||||
in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES)
|
||||
fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight)) for url in remote_urls)
|
||||
try:
|
||||
return tuple(await asyncio.gather(*fetches))
|
||||
except BaseException:
|
||||
for fetch in fetches:
|
||||
fetch.cancel()
|
||||
await asyncio.gather(*fetches, return_exceptions=True)
|
||||
raise
|
||||
|
||||
|
||||
async def async_inline_remote_media(
|
||||
messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues]
|
||||
should_inline: Callable[[RemoteMedia], bool] = inline_every_remote_url,
|
||||
) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues]
|
||||
remote_urls: Final = tuple(
|
||||
dict.fromkeys(
|
||||
remote.url
|
||||
for message in messages
|
||||
for part in _content_parts(message)
|
||||
if (remote := _parse_remote_part(part)) is not None and should_inline(_remote_media(remote))
|
||||
)
|
||||
)
|
||||
if not remote_urls:
|
||||
return messages
|
||||
data_urls: Final = await _fetch_data_urls(remote_urls)
|
||||
inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True)))
|
||||
return [ # mutable-ok: transform_request takes a list
|
||||
_inline_message(message, inlined, should_inline) for message in messages
|
||||
]
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config):
|
|||
check but still resolve DNS and still rewrite HTTP to the resolved IP.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import socket
|
||||
from ipaddress import ip_address, ip_network
|
||||
from typing import Any, Final, Protocol
|
||||
|
|
@ -471,7 +472,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response
|
|||
kwargs.pop("follow_redirects", None)
|
||||
headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})}
|
||||
for _ in range(_MAX_REDIRECTS):
|
||||
validated_url, original_host = validate_url(url)
|
||||
validated_url, original_host = await asyncio.to_thread(validate_url, url)
|
||||
response = await fetcher.get(
|
||||
validated_url,
|
||||
headers={**headers_view["headers"], "Host": original_host},
|
||||
|
|
|
|||
|
|
@ -411,6 +411,10 @@ class BaseConfig(ABC):
|
|||
def has_custom_stream_wrapper(self) -> bool:
|
||||
return False
|
||||
|
||||
@property
|
||||
def uses_async_transform_request(self) -> bool:
|
||||
return False
|
||||
|
||||
@property
|
||||
def supports_stream_param_in_request_body(self) -> bool:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import types
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import httpx
|
||||
|
|
@ -102,6 +103,24 @@ class BaseImageEditConfig(ABC):
|
|||
) -> tuple[dict, RequestFiles]:
|
||||
pass
|
||||
|
||||
async def async_transform_image_edit_request(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str | None,
|
||||
image: FileTypes | None,
|
||||
image_edit_optional_request_params: Mapping[str, object],
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: Mapping[str, str],
|
||||
) -> tuple[dict, RequestFiles]:
|
||||
return self.transform_image_edit_request(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
image=image,
|
||||
image_edit_optional_request_params=dict(image_edit_optional_request_params),
|
||||
litellm_params=litellm_params,
|
||||
headers=dict(headers),
|
||||
)
|
||||
|
||||
def finalize_image_edit_request_data(self, data: dict, resolved_request_url: str) -> dict:
|
||||
"""
|
||||
Last pass on the request dict after ``transform_image_edit_request``, using the
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
convert_to_anthropic_image_obj,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
||||
async_convert_url_to_base64,
|
||||
async_inline_remote_media,
|
||||
convert_url_to_base64,
|
||||
)
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
|
@ -172,6 +172,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
|
||||
return _anthropic_request
|
||||
|
||||
@property
|
||||
def uses_async_transform_request(self) -> bool:
|
||||
return True
|
||||
|
||||
async def async_transform_request(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -180,26 +184,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
_anthropic_request: Final = self._build_bedrock_anthropic_request_base(
|
||||
return self.transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
messages=await async_inline_remote_media(messages),
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
await self._async_convert_document_url_sources_to_base64(_anthropic_request)
|
||||
beta_list: Final = self._compute_bedrock_invoke_beta_headers(
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
headers=headers,
|
||||
)
|
||||
if beta_list:
|
||||
_anthropic_request["anthropic_beta"] = beta_list
|
||||
|
||||
return _anthropic_request
|
||||
|
||||
def _build_bedrock_anthropic_request_base(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -321,45 +313,6 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
"data": image_chunk["data"],
|
||||
}
|
||||
|
||||
async def _async_convert_document_url_sources_to_base64(self, anthropic_request: dict) -> None:
|
||||
"""
|
||||
Async version of document URL conversion for async completion paths.
|
||||
"""
|
||||
messages: Final = anthropic_request.get("messages")
|
||||
if not isinstance(messages, list):
|
||||
return
|
||||
|
||||
for message in messages:
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
content = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
|
||||
for block in content:
|
||||
if not isinstance(block, dict) or block.get("type") != "document":
|
||||
continue
|
||||
source = block.get("source")
|
||||
if not isinstance(source, dict) or source.get("type") != "url":
|
||||
continue
|
||||
source_url = source.get("url")
|
||||
if not isinstance(source_url, str):
|
||||
continue
|
||||
|
||||
inferred_format: str | None = None
|
||||
if source_url.lower().endswith(".pdf"):
|
||||
inferred_format = "application/pdf"
|
||||
base64_url = await async_convert_url_to_base64(url=source_url)
|
||||
image_chunk = convert_to_anthropic_image_obj(
|
||||
openai_image_url=base64_url,
|
||||
format=inferred_format,
|
||||
)
|
||||
block["source"] = {
|
||||
"type": "base64",
|
||||
"media_type": image_chunk["media_type"],
|
||||
"data": image_chunk["data"],
|
||||
}
|
||||
|
||||
def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict:
|
||||
"""
|
||||
Convert tool search entries to the format supported by the Bedrock Invoke API.
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth.
|
|||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
|
||||
AmazonAnthropicClaudeConfig,
|
||||
)
|
||||
|
|
@ -110,21 +111,13 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig):
|
|||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
model_id: Final = model.replace("mantle/", "", 1)
|
||||
|
||||
request: Final = self._build_bedrock_anthropic_request_base(
|
||||
model=model_id,
|
||||
messages=messages,
|
||||
return self.transform_request(
|
||||
model=model,
|
||||
messages=await async_inline_remote_media(messages),
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
await self._async_convert_document_url_sources_to_base64(request)
|
||||
return self._restore_mantle_body_fields(
|
||||
request=request,
|
||||
model_id=model_id,
|
||||
optional_params=optional_params,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _restore_mantle_body_fields(request: dict, model_id: str, optional_params: dict) -> dict:
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ API Reference: https://docs.bfl.ai/
|
|||
|
||||
import base64
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -16,7 +17,7 @@ from httpx._types import RequestFiles
|
|||
|
||||
import litellm
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
from litellm.litellm_core_utils.url_utils import safe_get
|
||||
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
|
||||
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.images.main import ImageEditOptionalRequestParams
|
||||
|
|
@ -37,6 +38,22 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
_BFL_REQUEST_PARAMS: Final = (
|
||||
"seed",
|
||||
"output_format",
|
||||
"safety_tolerance",
|
||||
"prompt_upsampling",
|
||||
"aspect_ratio",
|
||||
"steps",
|
||||
"guidance",
|
||||
"grow_mask",
|
||||
"top",
|
||||
"bottom",
|
||||
"left",
|
||||
"right",
|
||||
)
|
||||
|
||||
|
||||
class BlackForestLabsImageEditConfig(BaseImageEditConfig):
|
||||
"""
|
||||
Configuration for Black Forest Labs image editing.
|
||||
|
|
@ -85,34 +102,10 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
|
|||
BFL-specific params are passed through directly.
|
||||
"""
|
||||
optional_params: Final[dict[str, object]] = {}
|
||||
|
||||
# Pass through BFL-specific params
|
||||
bfl_params: Final = [
|
||||
"seed",
|
||||
"output_format",
|
||||
"safety_tolerance",
|
||||
"prompt_upsampling",
|
||||
# Kontext-specific
|
||||
"aspect_ratio",
|
||||
# Fill/Inpaint-specific
|
||||
"steps",
|
||||
"guidance",
|
||||
"grow_mask",
|
||||
# Expand-specific
|
||||
"top",
|
||||
"bottom",
|
||||
"left",
|
||||
"right",
|
||||
]
|
||||
|
||||
# Convert TypedDict to regular dict for access
|
||||
params_dict: Final = dict(image_edit_optional_params)
|
||||
|
||||
for param in bfl_params:
|
||||
if param in params_dict:
|
||||
value = params_dict[param]
|
||||
if value is not None:
|
||||
optional_params[param] = value
|
||||
params: Final[Mapping[str, object]] = image_edit_optional_params
|
||||
for param in _BFL_REQUEST_PARAMS:
|
||||
if (value := params.get(param)) is not None:
|
||||
optional_params[param] = value
|
||||
|
||||
# Set default output format
|
||||
if "output_format" not in optional_params:
|
||||
|
|
@ -251,23 +244,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
|
|||
"input_image": b64_image,
|
||||
}
|
||||
|
||||
# Add optional params (only BFL-recognized parameters)
|
||||
bfl_request_params: Final = [
|
||||
"seed",
|
||||
"output_format",
|
||||
"safety_tolerance",
|
||||
"prompt_upsampling",
|
||||
"aspect_ratio",
|
||||
"steps",
|
||||
"guidance",
|
||||
"grow_mask",
|
||||
"top",
|
||||
"bottom",
|
||||
"left",
|
||||
"right",
|
||||
]
|
||||
for key, value in image_edit_optional_request_params.items():
|
||||
if key in bfl_request_params and value is not None:
|
||||
if key in _BFL_REQUEST_PARAMS and value is not None:
|
||||
request_body[key] = value
|
||||
|
||||
# Handle mask if provided (for inpainting)
|
||||
|
|
@ -277,7 +255,39 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
|
|||
request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8")
|
||||
|
||||
# BFL uses JSON, not multipart - return empty files
|
||||
return request_body, []
|
||||
return request_body, ()
|
||||
|
||||
async def async_transform_image_edit_request(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str | None,
|
||||
image: FileTypes | None,
|
||||
image_edit_optional_request_params: Mapping[str, object],
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: Mapping[str, str],
|
||||
) -> tuple[dict, RequestFiles]:
|
||||
downloaded_image: Final = await self._fetch_remote_image(image)
|
||||
downloaded_mask: Final = await self._fetch_remote_image(image_edit_optional_request_params.get("mask"))
|
||||
return self.transform_image_edit_request(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
image=image if downloaded_image is None else downloaded_image,
|
||||
image_edit_optional_request_params=(
|
||||
dict(image_edit_optional_request_params)
|
||||
if downloaded_mask is None
|
||||
else {**image_edit_optional_request_params, "mask": downloaded_mask}
|
||||
),
|
||||
litellm_params=litellm_params,
|
||||
headers=dict(headers),
|
||||
)
|
||||
|
||||
async def _fetch_remote_image(self, image: object) -> bytes | None:
|
||||
candidate: Final = image[0] if isinstance(image, list) and image else image
|
||||
if not isinstance(candidate, str) or not candidate.startswith(("http://", "https://")):
|
||||
return None
|
||||
response: Final = await async_safe_get(litellm.module_level_aclient, candidate, timeout=60.0)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
|
||||
def transform_image_edit_response(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -116,6 +116,7 @@ from litellm.types.llms.anthropic_skills import (
|
|||
Skill,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
CreateBatchRequest,
|
||||
CreateFileRequest,
|
||||
FileContentRequest,
|
||||
|
|
@ -485,7 +486,7 @@ class BaseLLMHTTPHandler:
|
|||
def completion(
|
||||
self,
|
||||
model: str,
|
||||
messages: list,
|
||||
messages: list[AllMessageValues],
|
||||
api_base: str | None,
|
||||
custom_llm_provider: str,
|
||||
model_response: ModelResponse,
|
||||
|
|
@ -504,7 +505,7 @@ class BaseLLMHTTPHandler:
|
|||
shared_session: Optional["ClientSession"] = None,
|
||||
):
|
||||
json_mode: Final[bool] = optional_params.pop("json_mode", False)
|
||||
extra_body: Final[dict | None] = optional_params.pop("extra_body", None)
|
||||
extra_body: Final[Mapping[str, object] | None] = optional_params.pop("extra_body", None)
|
||||
|
||||
provider_config = provider_config or ProviderConfigManager.get_provider_chat_config(
|
||||
model=model, provider=litellm.LlmProviders(custom_llm_provider)
|
||||
|
|
@ -519,14 +520,17 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
# get config from model, custom llm provider
|
||||
headers = provider_config.validate_environment(
|
||||
api_key=api_key,
|
||||
headers=headers or {},
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
request_headers: Final = cast( # cast-ok: validate_environment is declared as a bare dict
|
||||
"dict[str, object]",
|
||||
provider_config.validate_environment(
|
||||
api_key=api_key,
|
||||
headers=headers or {},
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
),
|
||||
)
|
||||
|
||||
api_base = provider_config.get_complete_url(
|
||||
|
|
@ -538,93 +542,117 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
data: dict[str, object] = provider_config.transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
if extra_body is not None:
|
||||
data = {**data, **extra_body}
|
||||
|
||||
headers, signed_json_body = provider_config.sign_request(
|
||||
headers=headers,
|
||||
optional_params={
|
||||
**optional_params,
|
||||
**_aws_signing_overrides(optional_params, litellm_params),
|
||||
},
|
||||
request_data=data,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
stream=stream,
|
||||
fake_stream=fake_stream,
|
||||
model=model,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key=api_key,
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
"api_base": api_base,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
# Check if stream was converted for WebSearch interception
|
||||
# This is set by the async_pre_request_hook in WebSearchInterceptionLogger
|
||||
if litellm_params.get("_websearch_interception_converted_stream", False):
|
||||
logging_obj.model_call_details["websearch_interception_converted_stream"] = True
|
||||
|
||||
if acompletion is True:
|
||||
if stream is True:
|
||||
data = self._add_stream_param_to_request_body(
|
||||
data=data,
|
||||
provider_config=provider_config,
|
||||
def sign_and_log(
|
||||
transformed: dict[str, object], # mutable-ok: async_completion takes dict
|
||||
) -> tuple[dict[str, object], dict[str, object], bytes | None]: # mutable-ok: async_completion takes dict
|
||||
data: Final = {**transformed, **extra_body} if extra_body is not None else transformed
|
||||
signed: Final = cast( # cast-ok: sign_request is declared as a bare dict
|
||||
"tuple[dict[str, object], bytes | None]",
|
||||
provider_config.sign_request(
|
||||
headers=request_headers,
|
||||
optional_params={
|
||||
**optional_params,
|
||||
**_aws_signing_overrides(optional_params, litellm_params),
|
||||
},
|
||||
request_data=data,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
stream=stream,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
model=model,
|
||||
),
|
||||
)
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key=api_key,
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
"api_base": api_base,
|
||||
"headers": signed[0],
|
||||
},
|
||||
)
|
||||
if litellm_params.get("_websearch_interception_converted_stream", False):
|
||||
logging_obj.model_call_details["websearch_interception_converted_stream"] = True
|
||||
return data, signed[0], signed[1]
|
||||
|
||||
def dispatch_async(
|
||||
data: dict[str, object], # mutable-ok: async_completion takes dict
|
||||
signed_headers: dict[str, object], # mutable-ok: async_completion takes dict
|
||||
signed_json_body: bytes | None,
|
||||
):
|
||||
async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None
|
||||
if stream is True:
|
||||
return self.acompletion_stream_function(
|
||||
model=model,
|
||||
messages=messages,
|
||||
api_base=api_base,
|
||||
headers=headers,
|
||||
headers=signed_headers,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
provider_config=provider_config,
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
data=data,
|
||||
data=self._add_stream_param_to_request_body(
|
||||
data=data,
|
||||
provider_config=provider_config,
|
||||
fake_stream=fake_stream,
|
||||
),
|
||||
fake_stream=fake_stream,
|
||||
client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
|
||||
client=async_client,
|
||||
litellm_params=litellm_params,
|
||||
json_mode=json_mode,
|
||||
optional_params=optional_params,
|
||||
signed_json_body=signed_json_body,
|
||||
)
|
||||
return self.async_completion(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
provider_config=provider_config,
|
||||
api_base=api_base,
|
||||
headers=signed_headers,
|
||||
data=data,
|
||||
timeout=timeout,
|
||||
model=model,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
api_key=api_key,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
encoding=encoding,
|
||||
client=async_client,
|
||||
json_mode=json_mode,
|
||||
signed_json_body=signed_json_body,
|
||||
shared_session=shared_session,
|
||||
)
|
||||
|
||||
else:
|
||||
return self.async_completion(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
provider_config=provider_config,
|
||||
api_base=api_base,
|
||||
headers=headers,
|
||||
data=data,
|
||||
timeout=timeout,
|
||||
model=model,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
api_key=api_key,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
encoding=encoding,
|
||||
client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
|
||||
json_mode=json_mode,
|
||||
signed_json_body=signed_json_body,
|
||||
shared_session=shared_session,
|
||||
if acompletion is True and provider_config.uses_async_transform_request:
|
||||
|
||||
async def transform_then_dispatch():
|
||||
transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict
|
||||
"dict[str, object]",
|
||||
await provider_config.async_transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=request_headers,
|
||||
),
|
||||
)
|
||||
return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed))
|
||||
|
||||
return transform_then_dispatch()
|
||||
|
||||
data, signed_headers, signed_json_body = sign_and_log(
|
||||
provider_config.transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=request_headers,
|
||||
)
|
||||
)
|
||||
|
||||
if acompletion is True:
|
||||
return dispatch_async(data, signed_headers, signed_json_body)
|
||||
|
||||
if stream is True:
|
||||
data = self._add_stream_param_to_request_body(
|
||||
|
|
@ -638,7 +666,7 @@ class BaseLLMHTTPHandler:
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging_obj,
|
||||
api_base=api_base,
|
||||
headers=headers,
|
||||
headers=signed_headers,
|
||||
data=data,
|
||||
signed_json_body=signed_json_body,
|
||||
messages=messages,
|
||||
|
|
@ -648,7 +676,7 @@ class BaseLLMHTTPHandler:
|
|||
completion_stream, headers = self.make_sync_call(
|
||||
provider_config=provider_config,
|
||||
api_base=api_base,
|
||||
headers=headers,
|
||||
headers=signed_headers,
|
||||
data=data,
|
||||
signed_json_body=signed_json_body,
|
||||
original_data=data,
|
||||
|
|
@ -681,7 +709,7 @@ class BaseLLMHTTPHandler:
|
|||
sync_httpx_client=sync_httpx_client,
|
||||
provider_config=provider_config,
|
||||
api_base=api_base,
|
||||
headers=headers,
|
||||
headers=signed_headers,
|
||||
data=data,
|
||||
signed_json_body=signed_json_body,
|
||||
timeout=timeout,
|
||||
|
|
@ -6754,7 +6782,7 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params=dict(litellm_params),
|
||||
)
|
||||
|
||||
data, files = image_edit_provider_config.transform_image_edit_request(
|
||||
data, files = await image_edit_provider_config.async_transform_image_edit_request(
|
||||
model=model,
|
||||
image=image,
|
||||
prompt=prompt,
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionFileObject
|
|||
from litellm.types.llms.vertex_ai import ContentType, PartType
|
||||
from litellm.utils import supports_reasoning
|
||||
|
||||
from ...vertex_ai.gemini.transformation import _gemini_convert_messages_with_history
|
||||
from ...vertex_ai.gemini.transformation import GEMINI_FILES_API_URI_PREFIX, _gemini_convert_messages_with_history
|
||||
from ...vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
|
||||
|
||||
|
||||
|
|
@ -127,7 +127,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
|
|||
if element.get("type") == "image_url":
|
||||
img_element = cast(ChatCompletionImageObject, element) # cast-ok: runtime type tag checked
|
||||
_image_url, format, detail = _image_url_fields(img_element)
|
||||
if _image_url and "https://" in _image_url:
|
||||
if (
|
||||
_image_url
|
||||
and "https://" in _image_url
|
||||
and not _image_url.startswith(GEMINI_FILES_API_URI_PREFIX)
|
||||
):
|
||||
image_obj = convert_to_anthropic_image_obj(_image_url, format=format)
|
||||
converted_image_url = convert_generic_image_chunk_to_openai_image_obj(image_obj)
|
||||
if detail is not None:
|
||||
|
|
@ -147,7 +151,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
|
|||
llm_provider="gemini",
|
||||
)
|
||||
file_id = _file_field.get("file_id")
|
||||
if file_id and ("http://" in file_id or "https://" in file_id):
|
||||
if (
|
||||
file_id
|
||||
and ("http://" in file_id or "https://" in file_id)
|
||||
and not file_id.startswith(GEMINI_FILES_API_URI_PREFIX)
|
||||
):
|
||||
# Convert HTTP/HTTPS file URL to base64 data
|
||||
try:
|
||||
base64_data = convert_url_to_base64(file_id)
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
create_anthropic_image_param,
|
||||
select_anthropic_content_block_type_for_file,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
|
||||
from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload
|
||||
|
|
@ -421,6 +422,21 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body)
|
||||
return self._transform_request_openai(model, messages, optional_params, stream, extra_body)
|
||||
|
||||
@property
|
||||
def uses_async_transform_request(self) -> bool:
|
||||
return True
|
||||
|
||||
async def async_transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
|
||||
optional_params: dict[str, object], # mutable-ok: BaseConfig signature
|
||||
litellm_params: dict[str, object], # mutable-ok: BaseConfig signature
|
||||
headers: dict[str, object], # mutable-ok: BaseConfig signature
|
||||
) -> dict[str, object]: # mutable-ok: BaseConfig signature
|
||||
inlined_messages: Final = await async_inline_remote_media(messages) if _is_claude_model(model) else messages
|
||||
return self.transform_request(model, inlined_messages, optional_params, litellm_params, headers)
|
||||
|
||||
def _transform_request_openai(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ Why separate file? Make it easy to see how transformation works
|
|||
import json
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, cast
|
||||
from urllib.parse import quote
|
||||
|
||||
|
|
@ -27,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
convert_to_gemini_tool_call_result,
|
||||
response_schema_prompt,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, async_inline_remote_media
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
|
||||
from litellm.types.files import (
|
||||
|
|
@ -68,6 +70,7 @@ _GCS_METADATA_VERTEX_BASE: Any | None = None
|
|||
# Shared sync client for GCS JSON API metadata reads so proxy/SSL settings
|
||||
# from litellm's HTTP stack apply (see Greptile review on PR #27278).
|
||||
_GCS_METADATA_HTTP_HANDLER: HTTPHandler | None = None
|
||||
GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/"
|
||||
_GEMINI_MIME_TYPE_ALIASES: Final[dict[str, str]] = {
|
||||
"image/jpg": "image/jpeg",
|
||||
}
|
||||
|
|
@ -556,7 +559,7 @@ def _process_gemini_media(
|
|||
file_data = FileDataType(mime_type=mime_type, file_uri=image_url)
|
||||
part: PartType = {"file_data": file_data}
|
||||
return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata)
|
||||
elif image_url.startswith("https://generativelanguage.googleapis.com/v1beta/files/"):
|
||||
elif image_url.startswith(GEMINI_FILES_API_URI_PREFIX):
|
||||
# Gemini Files API URIs — the file is already uploaded to Google's
|
||||
# servers; pass the URI through as file_data without fetching it.
|
||||
# These URLs return 403 when accessed directly, so we must not try
|
||||
|
|
@ -1307,6 +1310,23 @@ def sync_transform_request_body(
|
|||
)
|
||||
|
||||
|
||||
def _explicit_mime_type(fields: Mapping[str, object]) -> str | None:
|
||||
hint: Final = fields.get("format") or fields.get("mime_type") or fields.get("content_type")
|
||||
return hint if isinstance(hint, str) else None
|
||||
|
||||
|
||||
def _ai_studio_inlines(media: RemoteMedia) -> bool:
|
||||
return not media.url.startswith(GEMINI_FILES_API_URI_PREFIX)
|
||||
|
||||
|
||||
def _vertex_inlines(media: RemoteMedia) -> bool:
|
||||
if media.url.startswith(GEMINI_FILES_API_URI_PREFIX):
|
||||
return False
|
||||
return media.url.startswith("http://") or (
|
||||
_explicit_mime_type(media.fields) is None and _get_image_mime_type_from_url(media.url) is None
|
||||
)
|
||||
|
||||
|
||||
async def async_transform_request_body(
|
||||
gemini_api_key: str | None,
|
||||
messages: list[AllMessageValues],
|
||||
|
|
@ -1348,13 +1368,17 @@ async def async_transform_request_body(
|
|||
vertex_auth_header=vertex_auth_header,
|
||||
)
|
||||
|
||||
if _openai_messages_may_need_sync_gcs_metadata_fetch(messages):
|
||||
inlined_messages: Final = await async_inline_remote_media(
|
||||
messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines
|
||||
)
|
||||
|
||||
if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages):
|
||||
# _transform_request_body may issue a sync httpx.get (up to 5s timeout)
|
||||
# via _get_gcs_object_content_type to fetch GCS object metadata. Run the
|
||||
# whole sync transformation on a worker thread so it does not block the
|
||||
# async event loop.
|
||||
return await asyncify(_transform_request_body)(
|
||||
messages=messages,
|
||||
messages=inlined_messages,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_params=litellm_params,
|
||||
|
|
@ -1363,7 +1387,7 @@ async def async_transform_request_body(
|
|||
)
|
||||
|
||||
return _transform_request_body(
|
||||
messages=messages,
|
||||
messages=inlined_messages,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_params=litellm_params,
|
||||
|
|
|
|||
|
|
@ -4474,7 +4474,7 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR
|
|||
acompletion: Final = ctx.acompletion
|
||||
api_base: Final = ctx.api_base
|
||||
api_key: Final = ctx.api_key
|
||||
client = _dispatch_client_http(ctx)
|
||||
injected_client: Final = _dispatch_client_http(ctx)
|
||||
custom_llm_provider: Final = ctx.custom_llm_provider
|
||||
headers: Final = ctx.headers
|
||||
litellm_params: Final = ctx.litellm_params
|
||||
|
|
@ -4486,11 +4486,11 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR
|
|||
shared_session: Final = ctx.shared_session
|
||||
stream: Final = ctx.stream
|
||||
timeout: Final = ctx.timeout
|
||||
client: Final = (
|
||||
injected_client if injected_client is not None else (HTTPHandler(timeout=timeout) if stream is False else None)
|
||||
) # Keep this here, otherwise, the httpx.client closes and streaming is impossible
|
||||
|
||||
try:
|
||||
client = (
|
||||
HTTPHandler(timeout=timeout) if stream is False else None
|
||||
) # Keep this here, otherwise, the httpx.client closes and streaming is impossible
|
||||
response: Final = base_llm_http_handler.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
|
|
|
|||
|
|
@ -7,9 +7,12 @@
|
|||
# 4. Added proper cleanup in fixtures
|
||||
# 5. Added worker-specific isolation for parallel execution
|
||||
|
||||
import base64
|
||||
import importlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import asyncio
|
||||
|
|
@ -595,3 +598,43 @@ def pytest_sessionfinish(session, exitstatus):
|
|||
_close_handler_if_needed(getattr(litellm, "aclient", None))
|
||||
_close_handler_if_needed(getattr(litellm, "client", None))
|
||||
_run_coroutine_if_needed(close_litellm_async_clients())
|
||||
|
||||
|
||||
ONE_PIXEL_PNG = base64.b64decode(
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_only_image_fetch(monkeypatch):
|
||||
from litellm.litellm_core_utils.prompt_templates import factory, image_handling
|
||||
from litellm.llms.gemini.chat import transformation as gemini_chat_transformation
|
||||
|
||||
fetch = SimpleNamespace(
|
||||
fetched=[],
|
||||
base64_png=base64.b64encode(ONE_PIXEL_PNG).decode(),
|
||||
data_url="data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode(),
|
||||
)
|
||||
|
||||
def forbid_sync_fetch(client, url, **kwargs):
|
||||
raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}")
|
||||
|
||||
async def serve_png(client, url, **kwargs):
|
||||
fetch.fetched.append(url)
|
||||
return httpx.Response(
|
||||
200,
|
||||
content=ONE_PIXEL_PNG,
|
||||
headers={"content-type": "image/png"},
|
||||
request=httpx.Request("GET", url),
|
||||
)
|
||||
|
||||
def forbid_sync_convert(url, *args, **kwargs):
|
||||
if url.startswith(("http://", "https://")):
|
||||
raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}")
|
||||
return url
|
||||
|
||||
monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch)
|
||||
monkeypatch.setattr(image_handling, "async_safe_get", serve_png)
|
||||
for module in (image_handling, factory, gemini_chat_transformation):
|
||||
monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert)
|
||||
return fetch
|
||||
|
|
|
|||
|
|
@ -1,3 +1,7 @@
|
|||
import asyncio
|
||||
import copy
|
||||
import time
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -7,9 +11,13 @@ import litellm
|
|||
from litellm import constants
|
||||
from litellm.litellm_core_utils.prompt_templates import image_handling
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
||||
MAX_CONCURRENT_REMOTE_MEDIA_FETCHES,
|
||||
RemoteMedia,
|
||||
async_convert_url_to_base64,
|
||||
async_inline_remote_media,
|
||||
convert_url_to_base64,
|
||||
)
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
|
|
@ -107,9 +115,7 @@ class StreamingLargeImageClient:
|
|||
request=Request("GET", url),
|
||||
)
|
||||
# Mock the iter_bytes method to return our generator
|
||||
response.iter_bytes = lambda chunk_size=8192: generate_chunks(
|
||||
size_bytes, chunk_size
|
||||
)
|
||||
response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size)
|
||||
return response
|
||||
|
||||
|
||||
|
|
@ -207,9 +213,7 @@ def test_streaming_download_handles_petabyte_file(monkeypatch):
|
|||
"""
|
||||
# Simulate a 1 petabyte file (1,000,000 GB)
|
||||
# Without streaming protection, this would cause OOM or hang indefinitely
|
||||
client = StreamingLargeImageClient(
|
||||
size_mb=1_000_000_000, include_content_length=False
|
||||
)
|
||||
client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False)
|
||||
monkeypatch.setattr(litellm, "module_level_client", client)
|
||||
|
||||
with pytest.raises(litellm.ImageFetchError) as excinfo:
|
||||
|
|
@ -268,3 +272,259 @@ def test_image_size_limit_disabled(monkeypatch):
|
|||
|
||||
assert "Image URL download is disabled" in str(excinfo.value)
|
||||
assert "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0" in str(excinfo.value)
|
||||
|
||||
|
||||
async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_only_image_fetch):
|
||||
image_url = f"http://img.example/{uuid.uuid4()}.png"
|
||||
pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf"
|
||||
messages = [
|
||||
{"role": "system", "content": "be terse"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "what is this?"},
|
||||
{"type": "image_url", "image_url": {"url": image_url, "detail": "low"}},
|
||||
{"type": "image_url", "image_url": image_url},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
|
||||
{"type": "file", "file": {"file_id": pdf_url}},
|
||||
{"type": "file", "file": {"file_id": image_url, "format": "image/png"}},
|
||||
{"type": "document", "source": {"type": "url", "url": pdf_url}, "title": "the doc"},
|
||||
{"type": "image", "source": {"type": "url", "url": image_url}},
|
||||
{"type": "document", "source": {"type": "file", "file_id": "file_abc"}},
|
||||
],
|
||||
},
|
||||
]
|
||||
snapshot = copy.deepcopy(messages)
|
||||
|
||||
inlined = await async_inline_remote_media(messages)
|
||||
|
||||
data_url = async_only_image_fetch.data_url
|
||||
base64_png = async_only_image_fetch.base64_png
|
||||
assert inlined[0] == {"role": "system", "content": "be terse"}
|
||||
assert inlined[1]["content"] == [
|
||||
{"type": "text", "text": "what is this?"},
|
||||
{"type": "image_url", "image_url": {"url": data_url, "detail": "low"}},
|
||||
{"type": "image_url", "image_url": data_url},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
|
||||
{"type": "file", "file": {"format": "application/pdf", "file_data": data_url}},
|
||||
{"type": "file", "file": {"format": "image/png", "file_data": data_url}},
|
||||
{
|
||||
"type": "document",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": base64_png},
|
||||
"title": "the doc",
|
||||
},
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": base64_png}},
|
||||
{"type": "document", "source": {"type": "file", "file_id": "file_abc"}},
|
||||
]
|
||||
assert sorted(async_only_image_fetch.fetched) == sorted([image_url, pdf_url])
|
||||
assert messages == snapshot
|
||||
|
||||
|
||||
async def test_async_inline_remote_media_inlines_only_the_parts_the_predicate_accepts(async_only_image_fetch):
|
||||
files_api_prefix = "https://generativelanguage.googleapis.com/v1beta/files/"
|
||||
files_api_pdf = f"{files_api_prefix}{uuid.uuid4().hex}"
|
||||
hinted_image = f"https://img.example/{uuid.uuid4()}.png"
|
||||
plain_image = f"https://img.example/{uuid.uuid4()}.png"
|
||||
hinted_document = f"https://docs.example/{uuid.uuid4()}.pdf"
|
||||
seen = []
|
||||
|
||||
def inline_unhinted_outside_files_api(media: RemoteMedia) -> bool:
|
||||
seen.append(media)
|
||||
return not media.url.startswith(files_api_prefix) and "format" not in media.fields
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "file", "file": {"file_id": files_api_pdf}},
|
||||
{"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}},
|
||||
{"type": "image_url", "image_url": {"url": plain_image}},
|
||||
{"type": "image_url", "image_url": plain_image},
|
||||
{"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
snapshot = copy.deepcopy(messages)
|
||||
|
||||
inlined = await async_inline_remote_media(messages, should_inline=inline_unhinted_outside_files_api)
|
||||
|
||||
assert inlined[0]["content"] == [
|
||||
{"type": "file", "file": {"file_id": files_api_pdf}},
|
||||
{"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}},
|
||||
{"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}},
|
||||
{"type": "image_url", "image_url": async_only_image_fetch.data_url},
|
||||
{"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}},
|
||||
]
|
||||
assert async_only_image_fetch.fetched == [plain_image]
|
||||
assert [(media.url, dict(media.fields)) for media in seen[:5]] == [
|
||||
(files_api_pdf, {"file_id": files_api_pdf}),
|
||||
(hinted_image, {"url": hinted_image, "format": "image/png"}),
|
||||
(plain_image, {"url": plain_image}),
|
||||
(plain_image, {}),
|
||||
(hinted_document, {"type": "url", "url": hinted_document, "format": "application/pdf"}),
|
||||
]
|
||||
assert messages == snapshot
|
||||
|
||||
|
||||
async def test_async_inline_remote_media_inlines_a_shared_url_only_where_the_predicate_accepts_it(
|
||||
async_only_image_fetch,
|
||||
):
|
||||
shared = f"https://img.example/{uuid.uuid4()}.png"
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": shared, "format": "image/png"}},
|
||||
{"type": "image_url", "image_url": {"url": shared}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
inlined = await async_inline_remote_media(messages, should_inline=lambda media: "format" not in media.fields)
|
||||
|
||||
assert inlined[0]["content"] == [
|
||||
{"type": "image_url", "image_url": {"url": shared, "format": "image/png"}},
|
||||
{"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}},
|
||||
]
|
||||
assert async_only_image_fetch.fetched == [shared]
|
||||
|
||||
|
||||
async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fails(monkeypatch):
|
||||
missing = f"http://img.example/{uuid.uuid4()}-missing.png"
|
||||
slow = f"http://img.example/{uuid.uuid4()}-slow.png"
|
||||
slow_fetch_outcomes = []
|
||||
|
||||
async def serve(client, url, **kwargs):
|
||||
if url == missing:
|
||||
return Response(404, request=Request("GET", url))
|
||||
try:
|
||||
await asyncio.sleep(5)
|
||||
except asyncio.CancelledError:
|
||||
slow_fetch_outcomes.append("cancelled")
|
||||
raise
|
||||
slow_fetch_outcomes.append("finished")
|
||||
return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url))
|
||||
|
||||
monkeypatch.setattr(image_handling, "async_safe_get", serve)
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": missing}},
|
||||
{"type": "image_url", "image_url": {"url": slow}},
|
||||
],
|
||||
}
|
||||
]
|
||||
started = time.perf_counter()
|
||||
|
||||
with pytest.raises(litellm.ImageFetchError, match="Status code: 404"):
|
||||
await async_inline_remote_media(messages)
|
||||
|
||||
assert slow_fetch_outcomes == ["cancelled"]
|
||||
assert time.perf_counter() - started < 1
|
||||
|
||||
|
||||
_SSRF_VERDICTS = (
|
||||
SSRFError(
|
||||
"URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, "
|
||||
"add the host to `user_url_allowed_hosts` in general_settings."
|
||||
),
|
||||
SSRFError("DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known"),
|
||||
SSRFError("No addresses found for 'internal.example'"),
|
||||
)
|
||||
|
||||
|
||||
def _assert_verdict_free_messages(messages, url):
|
||||
assert len(messages) == len(_SSRF_VERDICTS)
|
||||
assert len(set(messages)) == 1, "a caller must not be able to tell a blocked host from one that does not resolve"
|
||||
message = messages[0]
|
||||
assert "The proxy could not resolve this host or its URL policy rejected it" in message
|
||||
assert "user_url_allowed_hosts" in message
|
||||
assert url in message
|
||||
assert "10.0.0.8" not in message
|
||||
assert "DNS" not in message
|
||||
assert "No addresses" not in message
|
||||
|
||||
|
||||
async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch):
|
||||
attempts = []
|
||||
messages = []
|
||||
url = f"http://internal.example/{uuid.uuid4()}.png"
|
||||
|
||||
for verdict in _SSRF_VERDICTS:
|
||||
|
||||
async def block(client, fetched_url, verdict=verdict, **kwargs):
|
||||
attempts.append(fetched_url)
|
||||
raise verdict
|
||||
|
||||
monkeypatch.setattr(image_handling, "async_safe_get", block)
|
||||
with pytest.raises(litellm.ImageFetchError) as raised:
|
||||
await async_convert_url_to_base64(url)
|
||||
messages.append(raised.value.message)
|
||||
|
||||
assert attempts == [url] * len(_SSRF_VERDICTS)
|
||||
_assert_verdict_free_messages(messages, url)
|
||||
|
||||
|
||||
def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch):
|
||||
attempts = []
|
||||
messages = []
|
||||
url = f"http://internal.example/{uuid.uuid4()}.png"
|
||||
|
||||
for verdict in _SSRF_VERDICTS:
|
||||
|
||||
def block(client, fetched_url, verdict=verdict, **kwargs):
|
||||
attempts.append(fetched_url)
|
||||
raise verdict
|
||||
|
||||
monkeypatch.setattr(image_handling, "safe_get", block)
|
||||
with pytest.raises(litellm.ImageFetchError) as raised:
|
||||
convert_url_to_base64(url)
|
||||
messages.append(raised.value.message)
|
||||
|
||||
assert attempts == [url] * len(_SSRF_VERDICTS)
|
||||
_assert_verdict_free_messages(messages, url)
|
||||
|
||||
|
||||
async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch):
|
||||
in_flight = {"now": 0, "peak": 0}
|
||||
|
||||
async def serve_png_slowly(client, url, **kwargs):
|
||||
in_flight["now"] += 1
|
||||
in_flight["peak"] = max(in_flight["peak"], in_flight["now"])
|
||||
await asyncio.sleep(0.01)
|
||||
in_flight["now"] -= 1
|
||||
return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url))
|
||||
|
||||
monkeypatch.setattr(image_handling, "async_safe_get", serve_png_slowly)
|
||||
urls = [f"https://img.example/{uuid.uuid4()}.png" for _ in range(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES + 5)]
|
||||
messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}} for url in urls]}]
|
||||
|
||||
inlined = await async_inline_remote_media(messages)
|
||||
|
||||
assert in_flight["peak"] == MAX_CONCURRENT_REMOTE_MEDIA_FETCHES
|
||||
assert all(part["image_url"]["url"].startswith("data:image/png;base64,") for part in inlined[0]["content"])
|
||||
|
||||
|
||||
async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch):
|
||||
messages = [
|
||||
{"role": "user", "content": [{"type": "text", "text": "hi"}]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}],
|
||||
},
|
||||
]
|
||||
|
||||
assert await async_inline_remote_media(messages) is messages
|
||||
assert async_only_image_fetch.fetched == []
|
||||
|
||||
|
||||
async def test_async_inline_remote_media_raises_image_fetch_error_when_the_fetch_fails(monkeypatch):
|
||||
async def serve_404(client, url, **kwargs):
|
||||
return Response(404, request=Request("GET", url))
|
||||
|
||||
monkeypatch.setattr(image_handling, "async_safe_get", serve_404)
|
||||
url = f"http://img.example/{uuid.uuid4()}.png"
|
||||
|
||||
with pytest.raises(litellm.ImageFetchError, match="Status code: 404"):
|
||||
await async_inline_remote_media([{"role": "user", "content": [{"type": "image_url", "image_url": url}]}])
|
||||
|
|
|
|||
|
|
@ -1,5 +1,9 @@
|
|||
import asyncio
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
|
|
@ -535,3 +539,34 @@ def test_assert_same_origin_error_message_does_not_leak_hostnames():
|
|||
detail = str(exc.value)
|
||||
assert "attacker.example.com" not in detail
|
||||
assert "api.internal-corp.example" not in detail
|
||||
|
||||
|
||||
async def test_async_safe_get_resolves_dns_off_the_event_loop(monkeypatch):
|
||||
loop_thread = threading.current_thread()
|
||||
resolver_threads = []
|
||||
|
||||
def slow_getaddrinfo(host, port, *args, **kwargs):
|
||||
resolver_threads.append(threading.current_thread())
|
||||
time.sleep(0.4)
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port or 443))]
|
||||
|
||||
monkeypatch.setattr(url_utils.socket, "getaddrinfo", slow_getaddrinfo)
|
||||
|
||||
class FakeClient:
|
||||
async def get(self, url, **kwargs):
|
||||
return httpx.Response(200, request=httpx.Request("GET", url))
|
||||
|
||||
ticks = [time.perf_counter()]
|
||||
|
||||
async def heartbeat():
|
||||
while True:
|
||||
await asyncio.sleep(0.01)
|
||||
ticks.append(time.perf_counter())
|
||||
|
||||
beating = asyncio.create_task(heartbeat())
|
||||
response = await url_utils.async_safe_get(FakeClient(), "https://img.example/a.png")
|
||||
beating.cancel()
|
||||
|
||||
assert response.status_code == 200
|
||||
assert resolver_threads and all(thread is not loop_thread for thread in resolver_threads)
|
||||
assert max(b - a for a, b in zip(ticks, ticks[1:])) < 0.2
|
||||
|
|
|
|||
|
|
@ -1,15 +1,19 @@
|
|||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
# Ensure the project root is on the import path so `litellm` can be imported when
|
||||
# tests are executed from any working directory.
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
|
||||
AmazonAnthropicClaudeConfig,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
|
||||
def test_get_supported_params_thinking():
|
||||
|
|
@ -714,3 +718,95 @@ def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking
|
|||
|
||||
assert result["thinking"] == {"type": "adaptive"}
|
||||
assert result["output_config"] == {"effort": "high"}
|
||||
|
||||
|
||||
async def test_bedrock_invoke_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch):
|
||||
image_url = f"http://img.example/{uuid.uuid4()}.png"
|
||||
captured = {}
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = request.content.decode()
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "us.anthropic.claude-sonnet-5",
|
||||
"content": [{"type": "text", "text": "Green"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/invoke/us.anthropic.claude-sonnet-5",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What colour is this?"},
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
aws_access_key_id="AKIAEXAMPLE",
|
||||
aws_secret_access_key="fake-secret",
|
||||
aws_region_name="us-east-1",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "Green"
|
||||
assert async_only_image_fetch.fetched == [image_url]
|
||||
assert image_url not in captured["body"]
|
||||
assert async_only_image_fetch.base64_png in captured["body"]
|
||||
|
||||
|
||||
async def test_bedrock_invoke_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch):
|
||||
pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf"
|
||||
captured = {}
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = json.loads(request.content)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "us.anthropic.claude-sonnet-5",
|
||||
"content": [{"type": "text", "text": "A lease"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/invoke/us.anthropic.claude-sonnet-5",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is this document?"},
|
||||
{"type": "document", "source": {"type": "url", "url": pdf_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
aws_access_key_id="AKIAEXAMPLE",
|
||||
aws_secret_access_key="fake-secret",
|
||||
aws_region_name="us-east-1",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "A lease"
|
||||
assert async_only_image_fetch.fetched == [pdf_url]
|
||||
assert {
|
||||
"type": "document",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png},
|
||||
} in captured["body"]["messages"][0]["content"]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,99 @@
|
|||
import json
|
||||
import uuid
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
|
||||
async def test_bedrock_mantle_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch):
|
||||
image_url = f"http://img.example/{uuid.uuid4()}.png"
|
||||
captured = {}
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = request.content.decode()
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "us.anthropic.claude-sonnet-5",
|
||||
"content": [{"type": "text", "text": "Green"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/mantle/us.anthropic.claude-sonnet-5",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What colour is this?"},
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
aws_access_key_id="AKIAEXAMPLE",
|
||||
aws_secret_access_key="fake-secret",
|
||||
aws_region_name="us-east-1",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "Green"
|
||||
assert async_only_image_fetch.fetched == [image_url]
|
||||
assert image_url not in captured["body"]
|
||||
assert async_only_image_fetch.base64_png in captured["body"]
|
||||
|
||||
|
||||
async def test_bedrock_mantle_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch):
|
||||
pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf"
|
||||
captured = {}
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = json.loads(request.content)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "us.anthropic.claude-sonnet-5",
|
||||
"content": [{"type": "text", "text": "A lease"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/mantle/us.anthropic.claude-sonnet-5",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is this document?"},
|
||||
{"type": "document", "source": {"type": "url", "url": pdf_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
aws_access_key_id="AKIAEXAMPLE",
|
||||
aws_secret_access_key="fake-secret",
|
||||
aws_region_name="us-east-1",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "A lease"
|
||||
assert async_only_image_fetch.fetched == [pdf_url]
|
||||
assert {
|
||||
"type": "document",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png},
|
||||
} in captured["body"]["messages"][0]["content"]
|
||||
|
|
@ -16,6 +16,9 @@ import httpx
|
|||
import pytest
|
||||
|
||||
|
||||
from litellm.llms.black_forest_labs.image_edit import (
|
||||
transformation as bfl_transformation,
|
||||
)
|
||||
from litellm.llms.black_forest_labs.image_edit.transformation import (
|
||||
BlackForestLabsImageEditConfig,
|
||||
)
|
||||
|
|
@ -186,7 +189,7 @@ class TestBlackForestLabsImageEditTransformation:
|
|||
assert data["output_format"] == "jpeg"
|
||||
|
||||
# BFL uses JSON, not multipart - files should be empty
|
||||
assert files == []
|
||||
assert files == ()
|
||||
|
||||
def test_transform_image_edit_request_with_mask(self):
|
||||
"""Test request transformation with mask for inpainting."""
|
||||
|
|
@ -299,3 +302,76 @@ class TestBlackForestLabsImageEditTransformation:
|
|||
def test_use_multipart_form_data_returns_false(self):
|
||||
"""Test that use_multipart_form_data returns False for BFL."""
|
||||
assert self.config.use_multipart_form_data() is False
|
||||
|
||||
|
||||
async def test_async_transform_image_edit_request_downloads_url_images_with_the_async_fetcher(monkeypatch):
|
||||
served = b"png-bytes-from-cdn"
|
||||
fetched = []
|
||||
|
||||
def forbid_sync_fetch(client, url, **kwargs):
|
||||
raise AssertionError(f"sync image fetch ran on the event loop: {url}")
|
||||
|
||||
async def serve(client, url, **kwargs):
|
||||
fetched.append((url, kwargs.get("timeout")))
|
||||
return httpx.Response(200, content=served, request=httpx.Request("GET", url))
|
||||
|
||||
monkeypatch.setattr(bfl_transformation, "safe_get", forbid_sync_fetch)
|
||||
monkeypatch.setattr(bfl_transformation, "async_safe_get", serve)
|
||||
|
||||
data, files = await BlackForestLabsImageEditConfig().async_transform_image_edit_request(
|
||||
model="flux-kontext-pro",
|
||||
prompt="Add a red hat",
|
||||
image="https://cdn.example/photo.png",
|
||||
image_edit_optional_request_params={"mask": "https://cdn.example/mask.png", "seed": 7},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert base64.b64decode(data["input_image"]) == served
|
||||
assert base64.b64decode(data["mask"]) == served
|
||||
assert data["seed"] == 7
|
||||
assert files == ()
|
||||
assert fetched == [("https://cdn.example/photo.png", 60.0), ("https://cdn.example/mask.png", 60.0)]
|
||||
|
||||
|
||||
async def test_async_transform_image_edit_request_never_fetches_for_local_images(monkeypatch):
|
||||
def refuse(*args, **kwargs):
|
||||
raise AssertionError("no network fetch expected for local image bytes")
|
||||
|
||||
monkeypatch.setattr(bfl_transformation, "safe_get", refuse)
|
||||
monkeypatch.setattr(bfl_transformation, "async_safe_get", refuse)
|
||||
|
||||
data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request(
|
||||
model="flux-kontext-pro",
|
||||
prompt="Add a red hat",
|
||||
image=[BytesIO(b"first"), BytesIO(b"other")],
|
||||
image_edit_optional_request_params={"mask": b"mask-bytes"},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert base64.b64decode(data["input_image"]) == b"first"
|
||||
assert base64.b64decode(data["mask"]) == b"mask-bytes"
|
||||
|
||||
|
||||
async def test_async_transform_image_edit_request_downloads_only_the_first_url_of_a_list(monkeypatch):
|
||||
fetched = []
|
||||
|
||||
async def serve(client, url, **kwargs):
|
||||
fetched.append(url)
|
||||
return httpx.Response(200, content=b"first-bytes", request=httpx.Request("GET", url))
|
||||
|
||||
monkeypatch.setattr(bfl_transformation, "safe_get", lambda *args, **kwargs: pytest.fail("sync fetch ran"))
|
||||
monkeypatch.setattr(bfl_transformation, "async_safe_get", serve)
|
||||
|
||||
data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request(
|
||||
model="flux-kontext-pro",
|
||||
prompt="Add a red hat",
|
||||
image=["https://cdn.example/a.png", "https://cdn.example/b.png"],
|
||||
image_edit_optional_request_params={},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert fetched == ["https://cdn.example/a.png"]
|
||||
assert base64.b64decode(data["input_image"]) == b"first-bytes"
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
|
|
@ -17,7 +18,8 @@ from litellm.llms.base_llm.audio_transcription.transformation import (
|
|||
AudioTranscriptionRequestData,
|
||||
BaseAudioTranscriptionConfig,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import (
|
||||
BaseLLMHTTPHandler,
|
||||
|
|
@ -30,7 +32,7 @@ from litellm.llms.azure.videos.transformation import AzureVideoConfig
|
|||
from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import TranscriptionResponse
|
||||
from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse
|
||||
|
||||
_ACTIVE_KEY = "_code_interpreter_interception_active"
|
||||
_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
|
||||
|
|
@ -3184,3 +3186,288 @@ async def test_async_container_list_handler_transforms_success_response():
|
|||
|
||||
assert [container.id for container in response.data] == ["cntr_a"]
|
||||
assert response.has_more is True
|
||||
|
||||
|
||||
class _TransformRecordingConfig(BaseConfig):
|
||||
def __init__(self, transform_async: bool):
|
||||
self.transform_async = transform_async
|
||||
self.transform_calls = []
|
||||
self.sign_threads = []
|
||||
|
||||
@property
|
||||
def uses_async_transform_request(self) -> bool:
|
||||
return self.transform_async
|
||||
|
||||
def get_supported_openai_params(self, model):
|
||||
return []
|
||||
|
||||
def map_openai_params(self, non_default_params, optional_params, model, drop_params):
|
||||
return optional_params
|
||||
|
||||
def validate_environment(
|
||||
self, headers, model, messages, optional_params, litellm_params, api_key=None, api_base=None
|
||||
):
|
||||
return {}
|
||||
|
||||
def transform_request(self, model, messages, optional_params, litellm_params, headers):
|
||||
self.transform_calls.append("sync")
|
||||
return {"transformed_by": "sync"}
|
||||
|
||||
async def async_transform_request(self, model, messages, optional_params, litellm_params, headers):
|
||||
self.transform_calls.append("async")
|
||||
return {"transformed_by": "async"}
|
||||
|
||||
def sign_request(
|
||||
self, headers, optional_params, request_data, api_base, api_key=None, model=None, stream=None, fake_stream=None
|
||||
):
|
||||
self.sign_threads.append(threading.current_thread())
|
||||
return headers, None
|
||||
|
||||
def transform_response(
|
||||
self,
|
||||
model,
|
||||
raw_response,
|
||||
model_response,
|
||||
logging_obj,
|
||||
request_data,
|
||||
messages,
|
||||
optional_params,
|
||||
litellm_params,
|
||||
encoding,
|
||||
api_key=None,
|
||||
json_mode=None,
|
||||
):
|
||||
model_response.choices[0].message.content = raw_response.json()["transformed_by"]
|
||||
return model_response
|
||||
|
||||
def get_error_class(self, error_message, status_code, headers):
|
||||
return BaseLLMException(status_code=status_code, message=error_message, headers=headers)
|
||||
|
||||
def get_model_response_iterator(self, streaming_response, sync_stream, json_mode=False):
|
||||
return litellm.OpenAIGPTConfig().get_model_response_iterator(
|
||||
streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode
|
||||
)
|
||||
|
||||
|
||||
def _start_async_completion(config, logging_obj=None):
|
||||
captured = {}
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = json.loads(request.content)
|
||||
return httpx.Response(200, json=captured["body"])
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
pending = BaseLLMHTTPHandler().completion(
|
||||
model="stub-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
api_base="https://llm.example/v1/chat",
|
||||
custom_llm_provider="openai",
|
||||
model_response=ModelResponse(),
|
||||
encoding=None,
|
||||
logging_obj=logging_obj if logging_obj is not None else Mock(dynamic_success_callbacks=None, model_call_details={}),
|
||||
optional_params={},
|
||||
timeout=10.0,
|
||||
litellm_params={},
|
||||
acompletion=True,
|
||||
client=client,
|
||||
provider_config=config,
|
||||
)
|
||||
return pending, captured
|
||||
|
||||
|
||||
async def test_completion_awaits_async_transform_request_when_config_opts_in():
|
||||
config = _TransformRecordingConfig(transform_async=True)
|
||||
|
||||
pending, captured = _start_async_completion(config)
|
||||
assert config.transform_calls == []
|
||||
|
||||
response = await pending
|
||||
|
||||
assert config.transform_calls == ["async"]
|
||||
assert captured["body"] == {"transformed_by": "async"}
|
||||
assert response.choices[0].message.content == "async"
|
||||
|
||||
|
||||
async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_transform():
|
||||
config = _TransformRecordingConfig(transform_async=True)
|
||||
loop_thread = threading.current_thread()
|
||||
pre_call_threads = []
|
||||
logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={})
|
||||
logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread())
|
||||
|
||||
pending, captured = _start_async_completion(config, logging_obj)
|
||||
response = await pending
|
||||
|
||||
assert response.choices[0].message.content == "async"
|
||||
assert captured["body"] == {"transformed_by": "async"}
|
||||
assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads)
|
||||
assert pre_call_threads and all(thread is not loop_thread for thread in pre_call_threads)
|
||||
|
||||
|
||||
async def test_completion_keeps_sync_transform_request_before_returning_by_default():
|
||||
config = _TransformRecordingConfig(transform_async=False)
|
||||
|
||||
pending, captured = _start_async_completion(config)
|
||||
assert config.transform_calls == ["sync"]
|
||||
|
||||
response = await pending
|
||||
|
||||
assert config.transform_calls == ["sync"]
|
||||
assert captured["body"] == {"transformed_by": "sync"}
|
||||
assert response.choices[0].message.content == "sync"
|
||||
|
||||
|
||||
def _sse_echoing_transformed_by(request):
|
||||
transformed_by = json.loads(request.content)["transformed_by"]
|
||||
chunk = {
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "stub-model",
|
||||
"choices": [{"index": 0, "delta": {"content": transformed_by}, "finish_reason": None}],
|
||||
}
|
||||
return httpx.Response(
|
||||
200,
|
||||
content=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode(),
|
||||
headers={"content-type": "text/event-stream"},
|
||||
request=request,
|
||||
)
|
||||
|
||||
|
||||
def _streaming_logging_obj():
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
logging_obj = Logging(
|
||||
model="stub-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="acompletion",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="async-transform-stream",
|
||||
function_id="f",
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="stub-model", user="", optional_params={}, litellm_params={}, custom_llm_provider="openai"
|
||||
)
|
||||
return logging_obj
|
||||
|
||||
|
||||
async def test_completion_streams_after_the_async_transform_request():
|
||||
config = _TransformRecordingConfig(transform_async=True)
|
||||
loop_thread = threading.current_thread()
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(_sse_echoing_transformed_by))
|
||||
|
||||
stream = await BaseLLMHTTPHandler().completion(
|
||||
model="stub-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
api_base="https://llm.example/v1/chat",
|
||||
custom_llm_provider="openai",
|
||||
model_response=ModelResponse(),
|
||||
encoding=None,
|
||||
logging_obj=_streaming_logging_obj(),
|
||||
optional_params={},
|
||||
timeout=10.0,
|
||||
litellm_params={},
|
||||
acompletion=True,
|
||||
stream=True,
|
||||
client=client,
|
||||
provider_config=config,
|
||||
)
|
||||
collected = [chunk async for chunk in stream]
|
||||
|
||||
assert config.transform_calls == ["async"]
|
||||
assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads)
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in collected) == "async"
|
||||
|
||||
|
||||
class _ImageEditRecordingConfig(BaseImageEditConfig):
|
||||
def __init__(self):
|
||||
self.transform_calls = []
|
||||
|
||||
def get_supported_openai_params(self, model):
|
||||
return []
|
||||
|
||||
def map_openai_params(self, image_edit_optional_params, model, drop_params):
|
||||
return dict(image_edit_optional_params)
|
||||
|
||||
def validate_environment(self, headers, model, api_key=None, litellm_params=None, api_base=None):
|
||||
return {}
|
||||
|
||||
def get_complete_url(self, model, api_base, litellm_params):
|
||||
return "https://images.example/v1/edits"
|
||||
|
||||
def use_multipart_form_data(self):
|
||||
return False
|
||||
|
||||
def transform_image_edit_request(
|
||||
self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers
|
||||
):
|
||||
self.transform_calls.append("sync")
|
||||
return {"transformed_by": "sync"}, []
|
||||
|
||||
async def async_transform_image_edit_request(
|
||||
self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers
|
||||
):
|
||||
self.transform_calls.append("async")
|
||||
return {"transformed_by": "async"}, []
|
||||
|
||||
def transform_image_edit_response(self, model, raw_response, logging_obj):
|
||||
return ImageResponse(data=[ImageObject(b64_json=raw_response.json()["transformed_by"])])
|
||||
|
||||
|
||||
def _echo_json_transport(captured):
|
||||
def handle(request):
|
||||
captured["body"] = json.loads(request.content)
|
||||
return httpx.Response(200, json=captured["body"])
|
||||
|
||||
return httpx.MockTransport(handle)
|
||||
|
||||
|
||||
async def test_async_image_edit_handler_awaits_the_async_transform():
|
||||
config = _ImageEditRecordingConfig()
|
||||
captured = {}
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=_echo_json_transport(captured))
|
||||
|
||||
response = await BaseLLMHTTPHandler().async_image_edit_handler(
|
||||
model="edit-model",
|
||||
image=b"raw-image",
|
||||
prompt="add a hat",
|
||||
image_edit_provider_config=config,
|
||||
image_edit_optional_request_params={},
|
||||
custom_llm_provider="openai",
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
logging_obj=Mock(),
|
||||
timeout=10.0,
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert config.transform_calls == ["async"]
|
||||
assert captured["body"] == {"transformed_by": "async"}
|
||||
assert response.data[0].b64_json == "async"
|
||||
|
||||
|
||||
def test_image_edit_handler_keeps_the_sync_transform():
|
||||
config = _ImageEditRecordingConfig()
|
||||
captured = {}
|
||||
client = HTTPHandler()
|
||||
client.client = httpx.Client(transport=_echo_json_transport(captured))
|
||||
|
||||
response = BaseLLMHTTPHandler().image_edit_handler(
|
||||
model="edit-model",
|
||||
image=b"raw-image",
|
||||
prompt="add a hat",
|
||||
image_edit_provider_config=config,
|
||||
image_edit_optional_request_params={},
|
||||
custom_llm_provider="openai",
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
logging_obj=Mock(),
|
||||
timeout=10.0,
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert config.transform_calls == ["sync"]
|
||||
assert captured["body"] == {"transformed_by": "sync"}
|
||||
assert response.data[0].b64_json == "sync"
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ Tests tool calling request/response transformations and chat completions
|
|||
import asyncio
|
||||
import os
|
||||
import copy
|
||||
import uuid
|
||||
import json
|
||||
from typing import Any, Dict, List
|
||||
|
||||
|
|
@ -945,3 +946,48 @@ class TestSnowflakeChatCompletion:
|
|||
|
||||
assert len(chunks_received) > 0
|
||||
content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content)
|
||||
|
||||
|
||||
async def test_snowflake_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch):
|
||||
image_url = f"http://img.example/{uuid.uuid4()}.png"
|
||||
captured = {}
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = request.content.decode()
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": [{"type": "text", "text": "Green"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model="snowflake/claude-sonnet-4-6",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What colour is this?"},
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
api_key="fake-jwt",
|
||||
account_id="FAKE-ACCOUNT",
|
||||
api_base=FAKE_API_BASE,
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "Green"
|
||||
assert async_only_image_fetch.fetched == [image_url]
|
||||
assert image_url not in captured["body"]
|
||||
assert async_only_image_fetch.base64_png in captured["body"]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,13 @@
|
|||
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import Mock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.vertex_ai.gemini import transformation
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
|
|
@ -338,3 +345,133 @@ def test_map_function_enterprise_web_search_snake_case():
|
|||
|
||||
assert len(result) == 1
|
||||
assert "enterpriseWebSearch" in result[0]
|
||||
|
||||
|
||||
async def test_gemini_ai_studio_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch):
|
||||
image_url = f"http://img.example/{uuid.uuid4()}.png"
|
||||
captured = {}
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = request.content.decode()
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"candidates": [{"content": {"parts": [{"text": "Green"}], "role": "model"}, "finishReason": "STOP"}],
|
||||
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model="gemini/gemini-3.8-flash",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What colour is this?"},
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
api_key="fake-gemini-key",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "Green"
|
||||
assert async_only_image_fetch.fetched == [image_url]
|
||||
assert image_url not in captured["body"]
|
||||
assert async_only_image_fetch.base64_png in captured["body"]
|
||||
|
||||
|
||||
async def test_gemini_ai_studio_async_completion_passes_files_api_uris_through_unfetched(async_only_image_fetch):
|
||||
files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}"
|
||||
files_api_image = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}"
|
||||
captured = {}
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = json.loads(request.content)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"candidates": [{"content": {"parts": [{"text": "A report"}], "role": "model"}, "finishReason": "STOP"}],
|
||||
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model="gemini/gemini-3.8-flash",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Summarize these"},
|
||||
{"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}},
|
||||
{"type": "image_url", "image_url": {"url": files_api_image, "format": "image/png"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
api_key="fake-gemini-key",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "A report"
|
||||
assert async_only_image_fetch.fetched == []
|
||||
file_parts = [part["file_data"] for part in captured["body"]["contents"][0]["parts"] if "file_data" in part]
|
||||
assert file_parts == [
|
||||
{"mime_type": "application/pdf", "file_uri": files_api_pdf},
|
||||
{"mime_type": "image/png", "file_uri": files_api_image},
|
||||
]
|
||||
|
||||
|
||||
async def test_vertex_ai_async_transform_inlines_only_the_urls_gemini_cannot_fetch_itself(async_only_image_fetch):
|
||||
plain_http_png = f"http://img.example/{uuid.uuid4()}.png"
|
||||
extensionless_https = f"https://cdn.example/files/{uuid.uuid4().hex}"
|
||||
https_png = f"https://img.example/{uuid.uuid4()}.png"
|
||||
hinted_extensionless = f"https://cdn.example/files/{uuid.uuid4().hex}"
|
||||
files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}"
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe these"},
|
||||
{"type": "image_url", "image_url": {"url": plain_http_png}},
|
||||
{"type": "image_url", "image_url": {"url": extensionless_https}},
|
||||
{"type": "image_url", "image_url": {"url": https_png}},
|
||||
{"type": "image_url", "image_url": {"url": hinted_extensionless, "mime_type": "image/webp"}},
|
||||
{"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
body = await transformation.async_transform_request_body(
|
||||
gemini_api_key=None,
|
||||
messages=messages,
|
||||
api_base=None,
|
||||
model="gemini-3.8-flash",
|
||||
client=None,
|
||||
timeout=None,
|
||||
extra_headers=None,
|
||||
optional_params={},
|
||||
logging_obj=Mock(),
|
||||
custom_llm_provider="vertex_ai",
|
||||
litellm_params={},
|
||||
vertex_project="qa-project",
|
||||
vertex_location="us-central1",
|
||||
vertex_auth_header=None,
|
||||
)
|
||||
|
||||
inlined = {"inline_data": {"mime_type": "image/png", "data": async_only_image_fetch.base64_png}}
|
||||
assert body["contents"][0]["parts"] == [
|
||||
{"text": "Describe these"},
|
||||
inlined,
|
||||
inlined,
|
||||
{"file_data": {"mime_type": "image/png", "file_uri": https_png}},
|
||||
{"file_data": {"mime_type": "image/webp", "file_uri": hinted_extensionless}},
|
||||
{"file_data": {"mime_type": "application/pdf", "file_uri": files_api_pdf}},
|
||||
]
|
||||
assert sorted(async_only_image_fetch.fetched) == sorted([plain_http_png, extensionless_https])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue