mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
503 lines
20 KiB
Python
503 lines
20 KiB
Python
import asyncio
|
|
import copy
|
|
import logging
|
|
from collections.abc import Iterator, Mapping
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import orjson
|
|
import pytest
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.testclient import TestClient
|
|
from starlette.requests import Request
|
|
from starlette.responses import Response
|
|
|
|
from litellm._logging import verbose_proxy_logger
|
|
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
|
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|
from litellm.proxy.image_endpoints import endpoints
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_generation_prompt_rerouting(monkeypatch):
|
|
"""Ensure image prompts are exposed to guardrails and restored afterwards."""
|
|
|
|
async def fake_add_litellm_data_to_request(**kwargs):
|
|
return kwargs["data"]
|
|
|
|
async def fake_update_request_status(**_: Any) -> None:
|
|
await asyncio.sleep(0)
|
|
|
|
proxy_logger_calls: dict[str, Any] = {}
|
|
|
|
async def fake_pre_call_hook(*, user_api_key_dict, data, call_type): # type: ignore[override]
|
|
proxy_logger_calls["pre_call_input"] = copy.deepcopy(data)
|
|
modified = {
|
|
**data,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "sanitized prompt",
|
|
}
|
|
],
|
|
}
|
|
return modified
|
|
|
|
async def fake_post_call_failure_hook(**_: Any) -> None:
|
|
return None
|
|
|
|
async def fake_post_call_success_hook(*, data, user_api_key_dict, response):
|
|
return response
|
|
|
|
async def fake_post_call_response_headers_hook(**kwargs):
|
|
return {"x-callback-test": "value"}
|
|
|
|
fake_proxy_logger = SimpleNamespace(
|
|
pre_call_hook=fake_pre_call_hook,
|
|
update_request_status=fake_update_request_status,
|
|
post_call_failure_hook=fake_post_call_failure_hook,
|
|
post_call_success_hook=fake_post_call_success_hook,
|
|
post_call_response_headers_hook=fake_post_call_response_headers_hook,
|
|
)
|
|
|
|
captured_route_request_data: dict[str, Any] = {}
|
|
|
|
async def fake_route_request(*, data, **kwargs): # type: ignore[override]
|
|
captured_route_request_data.update(data)
|
|
|
|
async def _inner():
|
|
class FakeResponse(dict):
|
|
_hidden_params = {}
|
|
|
|
return FakeResponse(result="ok")
|
|
|
|
return _inner()
|
|
|
|
scope = {
|
|
"type": "http",
|
|
"method": "POST",
|
|
"path": "/v1/images/generations",
|
|
"headers": [],
|
|
}
|
|
body = orjson.dumps({"prompt": "original prompt"})
|
|
|
|
async def receive():
|
|
return {"type": "http.request", "body": body, "more_body": False}
|
|
|
|
request = Request(scope, receive)
|
|
response = Response()
|
|
user_api_key = UserAPIKeyAuth()
|
|
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.proxy_server.add_litellm_data_to_request",
|
|
fake_add_litellm_data_to_request,
|
|
)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {})
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", fake_proxy_logger)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version")
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.get_custom_headers",
|
|
classmethod(lambda *args, **kwargs: {}),
|
|
)
|
|
monkeypatch.setattr("litellm.proxy.image_endpoints.endpoints.route_request", fake_route_request)
|
|
|
|
result = await endpoints.image_generation(
|
|
request=request,
|
|
fastapi_response=response,
|
|
user_api_key_dict=user_api_key,
|
|
)
|
|
await asyncio.sleep(0)
|
|
|
|
assert result == {"result": "ok"}
|
|
pre_call_input = proxy_logger_calls["pre_call_input"]
|
|
assert pre_call_input["messages"][0]["content"] == "original prompt"
|
|
assert captured_route_request_data["prompt"] == "sanitized prompt"
|
|
assert "messages" not in captured_route_request_data
|
|
logging_obj = captured_route_request_data.get("litellm_logging_obj")
|
|
assert logging_obj is not None
|
|
assert logging_obj.messages == "sanitized prompt"
|
|
assert logging_obj.model_call_details["prompt"] == "sanitized prompt"
|
|
assert response.headers.get("x-callback-test") == "value"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_generation_router_rejection_dispatches_failure_callback(monkeypatch):
|
|
"""Ensure router rejection (budget, cooldown, etc.) triggers failure handlers on litellm_logging_obj."""
|
|
from unittest.mock import AsyncMock
|
|
|
|
async def fake_add_litellm_data_to_request(**kwargs):
|
|
return kwargs["data"]
|
|
|
|
fake_proxy_logger = SimpleNamespace(
|
|
pre_call_hook=AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data),
|
|
post_call_failure_hook=AsyncMock(),
|
|
)
|
|
|
|
class RejectingRouter:
|
|
model_names = ["dall-e-3"]
|
|
team_public_model_names = []
|
|
|
|
def is_recognized_model(self, model: str) -> bool:
|
|
return True
|
|
|
|
async def aimage_generation(self, *args, **kwargs):
|
|
logging_obj = kwargs.get("litellm_logging_obj")
|
|
if logging_obj is not None:
|
|
# Simulate Router.async_get_available_deployment failure handling
|
|
await logging_obj.dispatch_failure_handlers(
|
|
exception=Exception("No deployment available"),
|
|
traceback_exception="Traceback...",
|
|
prefer_async_handlers=True,
|
|
)
|
|
raise Exception("No deployment available")
|
|
|
|
rejecting_router = RejectingRouter()
|
|
|
|
scope = {
|
|
"type": "http",
|
|
"method": "POST",
|
|
"path": "/v1/images/generations",
|
|
"headers": [],
|
|
}
|
|
body = orjson.dumps({"prompt": "generate a sunset", "model": "dall-e-3"})
|
|
|
|
async def receive():
|
|
return {"type": "http.request", "body": body, "more_body": False}
|
|
|
|
request = Request(scope, receive)
|
|
response = Response()
|
|
user_api_key = UserAPIKeyAuth()
|
|
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.proxy_server.add_litellm_data_to_request",
|
|
fake_add_litellm_data_to_request,
|
|
)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", rejecting_router)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {})
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", fake_proxy_logger)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version")
|
|
|
|
# Spy on dispatch_failure_handlers in litellm.litellm_core_utils.litellm_logging.Logging
|
|
dispatched_exceptions = []
|
|
orig_function_setup = endpoints.litellm.utils.function_setup
|
|
|
|
def spy_function_setup(*args, **kwargs):
|
|
logging_obj, data = orig_function_setup(*args, **kwargs)
|
|
orig_dispatch = logging_obj.dispatch_failure_handlers
|
|
|
|
async def spy_dispatch(*d_args, **d_kwargs):
|
|
dispatched_exceptions.append(d_kwargs.get("exception"))
|
|
return await orig_dispatch(*d_args, **d_kwargs)
|
|
|
|
logging_obj.dispatch_failure_handlers = spy_dispatch
|
|
return logging_obj, data
|
|
|
|
monkeypatch.setattr("litellm.utils.function_setup", spy_function_setup)
|
|
|
|
with pytest.raises(ProxyException, match="No deployment available"):
|
|
await endpoints.image_generation(
|
|
request=request,
|
|
fastapi_response=response,
|
|
user_api_key_dict=user_api_key,
|
|
)
|
|
|
|
assert len(dispatched_exceptions) == 1
|
|
assert str(dispatched_exceptions[0]) == "No deployment available"
|
|
|
|
|
|
def _image_edit_client(monkeypatch, captured: dict[str, Any]) -> TestClient:
|
|
class CaptureProcessing:
|
|
def __init__(self, data: dict[str, Any]) -> None:
|
|
captured.update(data)
|
|
|
|
async def base_process_llm_request(self, **_: Any) -> dict[str, Any]:
|
|
return {"data": [{"b64_json": "aGk="}]}
|
|
|
|
monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", CaptureProcessing)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
|
|
|
|
app = FastAPI()
|
|
app.include_router(endpoints.router)
|
|
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth()
|
|
return TestClient(app)
|
|
|
|
|
|
def test_image_edit_image_array_alias_is_not_forwarded(monkeypatch):
|
|
"""The documented `image[]` alias must reach the provider only as `image`."""
|
|
captured: dict[str, Any] = {}
|
|
|
|
response = _image_edit_client(monkeypatch, captured).post(
|
|
"/v1/images/edits",
|
|
files={"image[]": ("tree.png", b"\x89PNG\r\n\x1a\ntree", "image/png")},
|
|
data={"model": "gpt-image-1", "prompt": "add a hat"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert "image[]" not in captured
|
|
assert [buffer.getvalue() for buffer in captured["image"]] == [b"\x89PNG\r\n\x1a\ntree"]
|
|
assert [buffer.name for buffer in captured["image"]] == ["tree.png"]
|
|
|
|
|
|
def test_image_edit_mask_array_alias_is_not_forwarded(monkeypatch):
|
|
"""`mask[]` has the same shape as `image[]` and must be dropped the same way."""
|
|
captured: dict[str, Any] = {}
|
|
|
|
response = _image_edit_client(monkeypatch, captured).post(
|
|
"/v1/images/edits",
|
|
files={
|
|
"image": ("tree.png", b"\x89PNG\r\n\x1a\ntree", "image/png"),
|
|
"mask[]": ("mask.png", b"\x89PNG\r\n\x1a\nmask", "image/png"),
|
|
},
|
|
data={"model": "gpt-image-1", "prompt": "add a hat"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert "mask[]" not in captured
|
|
assert [buffer.getvalue() for buffer in captured["mask"]] == [b"\x89PNG\r\n\x1a\nmask"]
|
|
assert [buffer.getvalue() for buffer in captured["image"]] == [b"\x89PNG\r\n\x1a\ntree"]
|
|
|
|
|
|
def test_image_edit_canonical_file_fields_still_reach_the_provider(monkeypatch):
|
|
"""Dropping the bracketed aliases must not touch the canonical fields."""
|
|
captured: dict[str, Any] = {}
|
|
|
|
response = _image_edit_client(monkeypatch, captured).post(
|
|
"/v1/images/edits",
|
|
files={
|
|
"image": ("tree.png", b"\x89PNG\r\n\x1a\ntree", "image/png"),
|
|
"mask": ("mask.png", b"\x89PNG\r\n\x1a\nmask", "image/png"),
|
|
},
|
|
data={"model": "gpt-image-1", "prompt": "add a hat"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert [buffer.getvalue() for buffer in captured["image"]] == [b"\x89PNG\r\n\x1a\ntree"]
|
|
assert [buffer.getvalue() for buffer in captured["mask"]] == [b"\x89PNG\r\n\x1a\nmask"]
|
|
assert captured["prompt"] == "add a hat"
|
|
|
|
|
|
def test_image_edit_multipart_n_reaches_the_provider_as_an_int(monkeypatch):
|
|
"""A multipart `n` must not arrive as the string Starlette parsed it into."""
|
|
captured: dict[str, Any] = {}
|
|
|
|
response = _image_edit_client(monkeypatch, captured).post(
|
|
"/v1/images/edits",
|
|
files={"image": ("tree.png", b"\x89PNG\r\n\x1a\n", "image/png")},
|
|
data={"model": "nova-canvas", "prompt": "add a hat", "n": "2", "size": "1024x1024"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert captured["n"] == 2
|
|
assert isinstance(captured["n"], int)
|
|
assert captured["size"] == "1024x1024"
|
|
assert captured["prompt"] == "add a hat"
|
|
|
|
|
|
def test_image_edit_multipart_n_that_is_not_a_number_is_left_alone(monkeypatch):
|
|
"""An unparseable `n` still reaches the provider, which rejects it as before."""
|
|
captured: dict[str, Any] = {}
|
|
|
|
response = _image_edit_client(monkeypatch, captured).post(
|
|
"/v1/images/edits",
|
|
files={"image": ("tree.png", b"\x89PNG\r\n\x1a\n", "image/png")},
|
|
data={"model": "nova-canvas", "prompt": "add a hat", "n": "two"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert captured["n"] == "two"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"files, form, missing",
|
|
[
|
|
({}, {"model": "stability.stable-style-transfer-v1:0", "prompt": "oil painting"}, "image"),
|
|
(
|
|
{"image": ("tree.png", b"\x89PNG\r\n\x1a\n", "image/png")},
|
|
{"model": "stability.stable-image-remove-background-v1:0"},
|
|
"prompt",
|
|
),
|
|
],
|
|
)
|
|
def test_image_edit_without_an_optional_field_reaches_the_provider_with_it_set_to_none(
|
|
monkeypatch, files, form, missing
|
|
):
|
|
captured: dict[str, Any] = {}
|
|
|
|
response = _image_edit_client(monkeypatch, captured).post("/v1/images/edits", files=files or None, data=form)
|
|
|
|
assert response.status_code == 200, response.text
|
|
assert missing in captured and captured[missing] is None, captured
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(monkeypatch: pytest.MonkeyPatch):
|
|
"""A bare HTTPException carries no type or param, so the tail used to ship the
|
|
literal string "None" in both fields."""
|
|
|
|
async def fake_add_litellm_data_to_request(**kwargs: object) -> object:
|
|
return kwargs["data"]
|
|
|
|
async def fake_pre_call_hook(
|
|
*, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str
|
|
) -> dict[str, object]:
|
|
return data
|
|
|
|
async def fake_post_call_failure_hook(**_: object) -> None:
|
|
return None
|
|
|
|
async def failing_route_request(**_: object) -> None:
|
|
raise HTTPException(
|
|
status_code=404, detail={"error": "image_generation: Invalid model name passed in model=dall-e-3"}
|
|
)
|
|
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {})
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.proxy_server.proxy_logging_obj",
|
|
SimpleNamespace(pre_call_hook=fake_pre_call_hook, post_call_failure_hook=fake_post_call_failure_hook),
|
|
)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version")
|
|
monkeypatch.setattr("litellm.proxy.image_endpoints.endpoints.route_request", failing_route_request)
|
|
|
|
body = orjson.dumps({"model": "dall-e-3", "prompt": "a lighthouse at dusk"})
|
|
|
|
async def receive() -> dict[str, object]:
|
|
return {"type": "http.request", "body": body, "more_body": False}
|
|
|
|
request = Request({"type": "http", "method": "POST", "path": "/v1/images/generations", "headers": []}, receive)
|
|
|
|
with pytest.raises(ProxyException) as raised:
|
|
await endpoints.image_generation(
|
|
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
|
|
)
|
|
|
|
assert (raised.value.type, raised.value.param, raised.value.code) == ("invalid_request_error", None, "404")
|
|
|
|
|
|
@pytest.fixture
|
|
def propagating_proxy_logger() -> Iterator[None]:
|
|
verbose_proxy_logger.propagate = True
|
|
try:
|
|
yield
|
|
finally:
|
|
verbose_proxy_logger.propagate = False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failure_log_carries_the_callers_litellm_call_id(
|
|
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, propagating_proxy_logger: None
|
|
) -> None:
|
|
"""LIT-7836: the /v1/images/generations error line must carry the litellm_call_id
|
|
the client sent, both rendered in the message and as a structured record field."""
|
|
call_id = "images-call-7836"
|
|
|
|
async def fake_add_litellm_data_to_request(**kwargs: object) -> object:
|
|
return kwargs["data"]
|
|
|
|
async def fake_pre_call_hook(
|
|
*, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str
|
|
) -> dict[str, object]:
|
|
return data
|
|
|
|
async def fake_post_call_failure_hook(**_: object) -> None:
|
|
return None
|
|
|
|
async def failing_route_request(**_: object) -> None:
|
|
raise HTTPException(status_code=401, detail={"error": "invalid api key"})
|
|
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {})
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.proxy_server.proxy_logging_obj",
|
|
SimpleNamespace(pre_call_hook=fake_pre_call_hook, post_call_failure_hook=fake_post_call_failure_hook),
|
|
)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version")
|
|
monkeypatch.setattr("litellm.proxy.image_endpoints.endpoints.route_request", failing_route_request)
|
|
|
|
body = orjson.dumps({"model": "dall-e-3", "prompt": "a lighthouse at dusk"})
|
|
|
|
async def receive() -> dict[str, object]:
|
|
return {"type": "http.request", "body": body, "more_body": False}
|
|
|
|
request = Request(
|
|
{
|
|
"type": "http",
|
|
"method": "POST",
|
|
"path": "/v1/images/generations",
|
|
"headers": [(b"x-litellm-call-id", call_id.encode())],
|
|
},
|
|
receive,
|
|
)
|
|
|
|
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised:
|
|
await endpoints.image_generation(
|
|
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
|
|
)
|
|
|
|
assert raised.value.headers["x-litellm-call-id"] == call_id
|
|
record = next(r for r in caplog.records if "Exception occured" in r.getMessage())
|
|
assert record.litellm_call_id == call_id
|
|
assert call_id in record.getMessage()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failure_before_the_provider_call_bills_the_callers_litellm_call_id(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""LIT-7836: when the request is rejected while it is still being prepared, the
|
|
failure hook must see the same litellm_call_id the response header answers with,
|
|
otherwise the spend row is stored under a freshly minted id nobody can look up."""
|
|
call_id = "images-early-7836"
|
|
hook_request_data: list[Mapping[str, object]] = []
|
|
|
|
async def rejecting_add_litellm_data_to_request(**_: object) -> object:
|
|
raise HTTPException(status_code=400, detail={"error": "tag not allowed"})
|
|
|
|
async def fake_post_call_failure_hook(*, request_data: Mapping[str, object], **_: object) -> None:
|
|
hook_request_data.append(request_data)
|
|
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", rejecting_add_litellm_data_to_request)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {})
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.proxy_server.proxy_logging_obj",
|
|
SimpleNamespace(post_call_failure_hook=fake_post_call_failure_hook),
|
|
)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version")
|
|
|
|
body = orjson.dumps({"model": "dall-e-3", "prompt": "a lighthouse at dusk", "litellm_call_id": "from-the-body"})
|
|
|
|
async def receive() -> dict[str, object]:
|
|
return {"type": "http.request", "body": body, "more_body": False}
|
|
|
|
request = Request(
|
|
{
|
|
"type": "http",
|
|
"method": "POST",
|
|
"path": "/v1/images/generations",
|
|
"headers": [(b"x-litellm-call-id", call_id.encode())],
|
|
},
|
|
receive,
|
|
)
|
|
|
|
with pytest.raises(ProxyException) as raised:
|
|
await endpoints.image_generation(
|
|
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
|
|
)
|
|
|
|
assert raised.value.headers["x-litellm-call-id"] == call_id
|
|
assert [data["litellm_call_id"] for data in hook_request_data] == [call_id]
|