mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
8b89c909a9
error_status_code only read status_code, so a ProxyException raised
before routing (which stores its status as the string code) answered
500 with its 4xx type through the rerank, images, realtime, files, and
pass-through tails. It now falls back to a decimal code. A 408 maps to
timeout_error instead of invalid_request_error.
Tail regressions for rerank, images, realtime calls, and the chat
pass-through fail at the merge base with ('None', 'None'); the new
files-test helpers are fully typed.
213 lines
8 KiB
Python
213 lines
8 KiB
Python
import asyncio
|
|
import copy
|
|
from types import SimpleNamespace
|
|
from typing import Any, Dict
|
|
|
|
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.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
|
|
assert response.headers.get("x-callback-test") == "value"
|
|
|
|
|
|
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_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.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")
|