mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(hosted_vllm): inline Omni media URLs through SSRF-safe fetch
Fetch image, video, and audio reference HTTP URLs with safe_get and replace them with data URLs so Omni never requests user-controlled hosts
This commit is contained in:
mubashir1osmani
parent
fe03df4551
commit
a5779151e0
2 changed files with 205 additions and 44 deletions
|
|
@ -1,14 +1,19 @@
|
|||
"""Video generation for Hosted VLLM (vLLM-Omni OpenAI-compatible /v1/videos)."""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import mimetypes
|
||||
from collections.abc import Mapping
|
||||
from io import BufferedReader
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from httpx._types import FileTypes, RequestFiles
|
||||
|
||||
from litellm.images.utils import ImageEditRequestUtils
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, safe_get
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
|
@ -55,17 +60,74 @@ _VLLM_OMNI_VIDEO_PARAMS: Final = (
|
|||
"aspect_ratio",
|
||||
)
|
||||
|
||||
_REFERENCE_URL_KEYS: Final = MappingProxyType(
|
||||
{
|
||||
"image_reference": "image_url",
|
||||
"video_reference": "video_url",
|
||||
"audio_reference": "audio_url",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _serialize_form_value(value: object) -> str:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, bool):
|
||||
return "true" if value else "false"
|
||||
if isinstance(value, (Mapping, list)):
|
||||
if isinstance(value, (Mapping, list, tuple)):
|
||||
return json.dumps(value)
|
||||
return str(value)
|
||||
|
||||
|
||||
def _maybe_json(value: object) -> object:
|
||||
if not isinstance(value, str):
|
||||
return value
|
||||
stripped: Final = value.strip()
|
||||
if not stripped or stripped[0] not in "{[":
|
||||
return value
|
||||
return json.loads(stripped)
|
||||
|
||||
|
||||
def _content_type_from_headers(headers: Mapping[str, object]) -> str:
|
||||
raw: Final = headers.get("content-type", "application/octet-stream")
|
||||
if not isinstance(raw, str) or not raw.strip():
|
||||
return "application/octet-stream"
|
||||
return raw.split(";", 1)[0].strip() or "application/octet-stream"
|
||||
|
||||
|
||||
def _fetch_url_as_data_url(url: str, client: HTTPHandler) -> str:
|
||||
response: Final = safe_get(client, url)
|
||||
response.raise_for_status()
|
||||
encoded: Final = base64.b64encode(response.content).decode("ascii")
|
||||
return f"data:{_content_type_from_headers(response.headers)};base64,{encoded}"
|
||||
|
||||
|
||||
def _inline_reference_item(url_key: str, item: object, client: HTTPHandler) -> object:
|
||||
if not isinstance(item, Mapping):
|
||||
return item
|
||||
url: Final = item.get(url_key)
|
||||
if not isinstance(url, str):
|
||||
return item
|
||||
scheme: Final = urlparse(url).scheme.lower()
|
||||
if scheme in ("", "data"):
|
||||
return item
|
||||
if scheme not in ("http", "https"):
|
||||
raise SSRFError(f"URL scheme '{scheme}' is not allowed")
|
||||
return {**item, url_key: _fetch_url_as_data_url(url, client)} # mutable-ok: JSON form field serialized immediately
|
||||
|
||||
|
||||
def _inline_media_reference(field_name: str, value: object, client: HTTPHandler) -> object:
|
||||
url_key: Final = _REFERENCE_URL_KEYS.get(field_name)
|
||||
if url_key is None:
|
||||
return value
|
||||
parsed: Final = _maybe_json(value)
|
||||
if isinstance(parsed, list):
|
||||
return tuple(_inline_reference_item(url_key, item, client) for item in parsed)
|
||||
if isinstance(parsed, Mapping):
|
||||
return _inline_reference_item(url_key, parsed, client)
|
||||
return value
|
||||
|
||||
|
||||
def _input_reference_file(reference: object) -> tuple[str, FileTypes]:
|
||||
if isinstance(reference, BufferedReader):
|
||||
reader_name: Final = reference.name
|
||||
|
|
@ -88,6 +150,10 @@ class HostedVLLMVideoConfig(OpenAIVideoConfig):
|
|||
https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/videos_api/
|
||||
"""
|
||||
|
||||
def __init__(self, media_http_client: HTTPHandler | None = None) -> None:
|
||||
super().__init__()
|
||||
self._media_http_client = media_http_client
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseVideoConfig contract
|
||||
return [ # mutable-ok: BaseVideoConfig returns list
|
||||
*super().get_supported_openai_params(model),
|
||||
|
|
@ -145,9 +211,10 @@ class HostedVLLMVideoConfig(OpenAIVideoConfig):
|
|||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict, # mutable-ok: BaseVideoConfig contract
|
||||
) -> tuple[dict, RequestFiles, str]: # mutable-ok: BaseVideoConfig contract
|
||||
media_client: Final = self._media_http_client or HTTPHandler(concurrent_limit=1)
|
||||
input_reference: Final = video_create_optional_request_params.get("input_reference")
|
||||
form_files: Final = tuple(
|
||||
(key, (None, _serialize_form_value(value)))
|
||||
(key, (None, _serialize_form_value(_inline_media_reference(key, value, media_client))))
|
||||
for key, value in video_create_optional_request_params.items()
|
||||
if key not in _EXCLUDED_FORM_KEYS and value is not None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,12 +1,15 @@
|
|||
"""Tests for hosted_vllm video generation (vLLM-Omni /v1/videos)."""
|
||||
|
||||
import base64
|
||||
import json
|
||||
from io import BytesIO
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.hosted_vllm.videos import get_hosted_vllm_video_config
|
||||
from litellm.llms.hosted_vllm.videos.transformation import (
|
||||
HostedVLLMVideoConfig,
|
||||
|
|
@ -205,53 +208,144 @@ def test_get_supported_openai_params_includes_omni_extensions():
|
|||
assert "audio_reference" in supported
|
||||
|
||||
|
||||
def _mock_http_client(response_body: dict) -> MagicMock:
|
||||
mock_client = MagicMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {"content-type": "application/json"}
|
||||
mock_response.json.return_value = response_body
|
||||
mock_response.text = json.dumps(response_body)
|
||||
mock_client.post.return_value = mock_response
|
||||
return mock_client
|
||||
def _http_handler_for(handler) -> HTTPHandler:
|
||||
return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler)))
|
||||
|
||||
|
||||
def test_video_generation_posts_multipart_not_json():
|
||||
mock_client = _mock_http_client(
|
||||
{
|
||||
"id": "video-123",
|
||||
"object": "video",
|
||||
"status": "queued",
|
||||
"created_at": 1701234567,
|
||||
}
|
||||
)
|
||||
captured: list[httpx.Request] = []
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.llm_http_handler._get_httpx_client",
|
||||
return_value=mock_client,
|
||||
):
|
||||
response = litellm.video_generation(
|
||||
model="hosted_vllm/MiniMax-H3",
|
||||
prompt="three cats march into a bedroom playing tiny brass instruments",
|
||||
api_base="http://localhost:8091",
|
||||
api_key="test-key",
|
||||
extra_body={
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"fps": 24,
|
||||
"extra_params": {"task": "t2va", "duration": 10.0},
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured.append(request)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "video-123",
|
||||
"object": "video",
|
||||
"status": "queued",
|
||||
"created_at": 1701234567,
|
||||
},
|
||||
)
|
||||
|
||||
response = litellm.video_generation(
|
||||
model="hosted_vllm/MiniMax-H3",
|
||||
prompt="three cats march into a bedroom playing tiny brass instruments",
|
||||
api_base="http://localhost:8091",
|
||||
api_key="test-key",
|
||||
client=_http_handler_for(handler),
|
||||
extra_body={
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"fps": 24,
|
||||
"extra_params": {"task": "t2va", "duration": 10.0},
|
||||
},
|
||||
)
|
||||
|
||||
assert isinstance(response, VideoObject)
|
||||
assert response.status == "queued"
|
||||
mock_client.post.assert_called_once()
|
||||
post_kwargs = mock_client.post.call_args.kwargs
|
||||
assert post_kwargs["url"] == "http://localhost:8091/v1/videos"
|
||||
assert post_kwargs.get("json") is None
|
||||
assert post_kwargs["files"]
|
||||
fields = _form_fields(post_kwargs["files"])
|
||||
assert fields["prompt"] == "three cats march into a bedroom playing tiny brass instruments"
|
||||
assert fields["width"] == "1280"
|
||||
assert json.loads(fields["extra_params"]) == {"task": "t2va", "duration": 10.0}
|
||||
assert post_kwargs["headers"]["Authorization"] == "Bearer test-key"
|
||||
assert len(captured) == 1
|
||||
request = captured[0]
|
||||
assert str(request.url) == "http://localhost:8091/v1/videos"
|
||||
assert request.headers["authorization"] == "Bearer test-key"
|
||||
body = request.content
|
||||
assert b'name="prompt"' in body
|
||||
assert b"three cats march into a bedroom playing tiny brass instruments" in body
|
||||
assert b'name="width"' in body
|
||||
assert b"1280" in body
|
||||
assert b'name="extra_params"' in body
|
||||
assert b"t2va" in body
|
||||
assert request.headers.get("content-type", "").startswith("multipart/form-data")
|
||||
|
||||
|
||||
def test_http_image_reference_is_inlined_as_data_url():
|
||||
png_bytes = b"fake-png"
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.url.host == "1.1.1.1"
|
||||
return httpx.Response(200, content=png_bytes, headers={"content-type": "image/png"})
|
||||
|
||||
config = HostedVLLMVideoConfig(media_http_client=_http_handler_for(handler))
|
||||
_, files, _ = config.transform_video_create_request(
|
||||
model="MiniMax-H3",
|
||||
prompt="a person singing",
|
||||
api_base="http://localhost:8091/v1/videos",
|
||||
video_create_optional_request_params={
|
||||
"image_reference": {"image_url": "http://1.1.1.1/face.png"},
|
||||
},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
|
||||
payload = json.loads(_form_fields(files)["image_reference"])
|
||||
assert payload["image_url"] == "data:image/png;base64," + base64.b64encode(png_bytes).decode("ascii")
|
||||
|
||||
|
||||
def test_http_audio_reference_json_string_is_inlined_as_data_url():
|
||||
audio_bytes = b"fake-mp3"
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, content=audio_bytes, headers={"content-type": "audio/mpeg"})
|
||||
|
||||
config = HostedVLLMVideoConfig(media_http_client=_http_handler_for(handler))
|
||||
_, files, _ = config.transform_video_create_request(
|
||||
model="MiniMax-H3",
|
||||
prompt="a person singing",
|
||||
api_base="http://localhost:8091/v1/videos",
|
||||
video_create_optional_request_params={
|
||||
"audio_reference": '{"audio_url": "http://1.1.1.1/speech.mp3"}',
|
||||
},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
|
||||
payload = json.loads(_form_fields(files)["audio_reference"])
|
||||
assert payload["audio_url"].startswith("data:audio/mpeg;base64,")
|
||||
assert base64.b64decode(payload["audio_url"].split(",", 1)[1]) == audio_bytes
|
||||
|
||||
|
||||
def test_data_url_image_reference_is_not_fetched():
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
raise AssertionError(f"unexpected fetch of {request.url}")
|
||||
|
||||
data_url = "data:image/png;base64,AAAA"
|
||||
config = HostedVLLMVideoConfig(media_http_client=_http_handler_for(handler))
|
||||
_, files, _ = config.transform_video_create_request(
|
||||
model="MiniMax-H3",
|
||||
prompt="a person singing",
|
||||
api_base="http://localhost:8091/v1/videos",
|
||||
video_create_optional_request_params={"image_reference": {"image_url": data_url}},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert json.loads(_form_fields(files)["image_reference"])["image_url"] == data_url
|
||||
|
||||
|
||||
def test_metadata_url_in_image_reference_is_rejected():
|
||||
config = HostedVLLMVideoConfig()
|
||||
with pytest.raises(SSRFError, match="blocked address"):
|
||||
config.transform_video_create_request(
|
||||
model="MiniMax-H3",
|
||||
prompt="a person singing",
|
||||
api_base="http://localhost:8091/v1/videos",
|
||||
video_create_optional_request_params={
|
||||
"image_reference": {"image_url": "http://169.254.169.254/latest/meta-data/"},
|
||||
},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
|
||||
|
||||
def test_file_scheme_media_reference_is_rejected():
|
||||
config = HostedVLLMVideoConfig()
|
||||
with pytest.raises(SSRFError, match="scheme"):
|
||||
config.transform_video_create_request(
|
||||
model="MiniMax-H3",
|
||||
prompt="a person singing",
|
||||
api_base="http://localhost:8091/v1/videos",
|
||||
video_create_optional_request_params={
|
||||
"video_reference": {"video_url": "file:///etc/passwd"},
|
||||
},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue