mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
b84493d9a7
* fix(proxy): record the route's call type on failed speech spend rows A failed /v1/audio/speech request never puts a logging object on the request data, so the failure lift returned nothing and the spend row's call_type came out empty (or whatever the client sent in the body). Without a logging object, lift the call type the route already resolved, so a failure row records the same call type a success row on that route records * fix(proxy): keep the call type of a failure logged under its own route An MCP tool listing that fails inside a Responses, chat, or messages call is logged under /mcp/list_tools, so its spend row keeps list_mcp_tools instead of taking the outer request's call type * fix(proxy): take a failed row's call type from its endpoint, not a route path guess * fix(proxy): log the endpoint call type on failed image generation and rerank rows * fix(proxy): estimate failed-call ledger usage from the route's call type again The failure usage snapshot now gets the same call type the spend row uses for its estimate, so a failed Responses call counts input from the newest compaction item again, as on main. The endpoint call type tests now go through real routers whose provider answers 400, instead of asserting on a mocked failure hook. * test(proxy): type the failed spend row tests and fixture --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
743 lines
31 KiB
Python
743 lines
31 KiB
Python
import asyncio
|
|
import copy
|
|
import logging
|
|
from collections.abc import Iterator, Mapping
|
|
from datetime import datetime
|
|
from queue import Queue
|
|
from types import SimpleNamespace
|
|
from typing import Any, Dict, Final, cast
|
|
|
|
import orjson
|
|
import pytest
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.testclient import TestClient
|
|
from starlette.requests import Request
|
|
from starlette.responses import Response
|
|
|
|
import litellm
|
|
from litellm._logging import session_id_var, trace_id_var, verbose_proxy_logger
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.models.team import LiteLLM_TeamTableCachedObj
|
|
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
|
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
|
from litellm.proxy.image_endpoints import endpoints
|
|
from litellm.proxy.proxy_server import PrismaClient, ProxyConfig
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_failures_for_current_minute
|
|
from litellm.types.utils import ImageObject, ImageResponse
|
|
from litellm.utils import custom_llm_setup
|
|
|
|
|
|
@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_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]
|
|
|
|
|
|
class _ImageRouterHandler(litellm.CustomLLM):
|
|
def __init__(self) -> None:
|
|
self.calls: list[str] = []
|
|
|
|
async def aimage_generation(
|
|
self,
|
|
model: str,
|
|
prompt: str,
|
|
model_response: ImageResponse,
|
|
api_key: str | None,
|
|
api_base: str | None,
|
|
optional_params: Mapping[str, object],
|
|
logging_obj: object,
|
|
timeout: object | None = None,
|
|
client: object | None = None,
|
|
) -> ImageResponse:
|
|
self.calls.append(model)
|
|
if model.endswith("failing"):
|
|
raise litellm.RateLimitError(message="primary failed", llm_provider="image-regression", model=model)
|
|
return ImageResponse(created=1, data=[ImageObject(url="https://images.example/generated.png")])
|
|
|
|
|
|
def _image_generation_client(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
team_settings: Mapping[str, object] | None = None,
|
|
working_params: Mapping[str, object] | None = None,
|
|
) -> tuple[TestClient, _ImageRouterHandler]:
|
|
from litellm.proxy import proxy_server
|
|
|
|
handler: Final = _ImageRouterHandler()
|
|
cache: Final = UserApiKeyCache()
|
|
monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list])
|
|
monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers])
|
|
monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "image-regression", "custom_handler": handler}])
|
|
custom_llm_setup()
|
|
if team_settings is not None:
|
|
team_id: Final = user_api_key_dict.team_id
|
|
assert team_id is not None
|
|
cache.set_cache(
|
|
key=f"team_id:{team_id}",
|
|
value=LiteLLM_TeamTableCachedObj(team_id=team_id, router_settings=dict(team_settings)),
|
|
local_only=True,
|
|
model_type=LiteLLM_TeamTableCachedObj,
|
|
)
|
|
monkeypatch.setattr(proxy_server, "prisma_client", cast(PrismaClient, object()))
|
|
|
|
router: Final = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "A",
|
|
"litellm_params": {"model": "image-regression/failing"},
|
|
"model_info": {"id": "image-regression-failing"},
|
|
},
|
|
{
|
|
"model_name": "B",
|
|
"litellm_params": {"model": "image-regression/working", **(working_params or {})},
|
|
"model_info": {"id": "image-regression-working"},
|
|
},
|
|
],
|
|
num_retries=0,
|
|
)
|
|
monkeypatch.setattr(litellm, "model_alias_map", {})
|
|
monkeypatch.setattr(proxy_server, "general_settings", {})
|
|
monkeypatch.setattr(proxy_server, "llm_router", router)
|
|
monkeypatch.setattr(proxy_server, "proxy_config", ProxyConfig())
|
|
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache))
|
|
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
|
|
monkeypatch.setattr(proxy_server, "user_model", None)
|
|
monkeypatch.setattr(proxy_server, "version", "test-version")
|
|
|
|
app: Final = FastAPI()
|
|
app.include_router(endpoints.router)
|
|
app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler)
|
|
app.dependency_overrides[user_api_key_auth] = lambda: user_api_key_dict
|
|
return TestClient(app, raise_server_exceptions=False), handler
|
|
|
|
|
|
class _ImageOutcomeLogger(CustomLogger):
|
|
def __init__(self) -> None:
|
|
self.events: Queue[tuple[object, object, object]] = Queue()
|
|
self.messages: Queue[object] = Queue()
|
|
self.correlation: Queue[tuple[str, str]] = Queue()
|
|
|
|
async def async_log_failure_event(
|
|
self,
|
|
kwargs: dict[str, object],
|
|
response_obj: object,
|
|
start_time: datetime,
|
|
end_time: datetime,
|
|
) -> None:
|
|
self.messages.put(kwargs.get("messages"))
|
|
self.correlation.put((trace_id_var.get(), session_id_var.get()))
|
|
self.events.put((kwargs.get("call_type"), kwargs.get("litellm_call_id"), kwargs.get("exception")))
|
|
|
|
async def async_log_success_event(
|
|
self,
|
|
kwargs: dict[str, object],
|
|
response_obj: object,
|
|
start_time: datetime,
|
|
end_time: datetime,
|
|
) -> None:
|
|
self.messages.put(kwargs.get("messages"))
|
|
self.events.put((kwargs.get("call_type"), kwargs.get("litellm_call_id"), response_obj))
|
|
|
|
|
|
class _RewritingImageLogger(_ImageOutcomeLogger):
|
|
async def async_pre_call_hook(
|
|
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str
|
|
) -> dict[str, object]:
|
|
return {**data, "messages": [{"role": "user", "content": "sanitized prompt"}]}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"path", ("/v1/images/generations", "/images/generations", "/openai/deployments/B/images/generations")
|
|
)
|
|
async def test_router_refused_image_generation_reaches_failure_callbacks(
|
|
monkeypatch: pytest.MonkeyPatch, path: str
|
|
) -> None:
|
|
from litellm.proxy import proxy_server
|
|
|
|
logger: Final = _ImageOutcomeLogger()
|
|
monkeypatch.setattr(litellm, "callbacks", [logger])
|
|
client, handler = _image_generation_client(
|
|
monkeypatch, UserAPIKeyAuth(), working_params={"max_budget": 1, "budget_duration": "1d"}
|
|
)
|
|
router: Final = proxy_server.llm_router
|
|
assert router is not None and router.router_budget_logger is not None
|
|
router.router_budget_logger.dual_cache.set_cache("deployment_spend:image-regression-working:1d", 2)
|
|
|
|
with client:
|
|
response: Final = client.post(
|
|
path, json={"model": "B", "prompt": "draw a tree"}, headers={"x-litellm-call-id": "refused-image"}
|
|
)
|
|
assert response.status_code == 429, response.text
|
|
call_type, call_id, exception = logger.events.get(timeout=5)
|
|
|
|
assert (call_type, call_id) == ("aimage_generation", "refused-image")
|
|
assert isinstance(exception, ValueError) and "crossed budget" in str(exception)
|
|
assert handler.calls == []
|
|
assert logger.events.empty()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("supply_logging", (False, True))
|
|
async def test_image_router_refusal_without_or_with_caller_logging(
|
|
monkeypatch: pytest.MonkeyPatch, supply_logging: bool
|
|
) -> None:
|
|
from litellm.proxy import proxy_server
|
|
|
|
logger: Final = _ImageOutcomeLogger()
|
|
monkeypatch.setattr(litellm, "callbacks", [logger])
|
|
monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
|
|
_image_generation_client(monkeypatch, UserAPIKeyAuth(), working_params={"max_budget": 1, "budget_duration": "1d"})
|
|
router: Final = proxy_server.llm_router
|
|
assert router is not None and router.router_budget_logger is not None
|
|
router.router_budget_logger.dual_cache.set_cache("deployment_spend:image-regression-working:1d", 2)
|
|
kwargs: Final[dict[str, object]] = {"litellm_trace_id": "refused-trace", "litellm_session_id": "refused-session"}
|
|
if supply_logging:
|
|
logging_obj, _ = litellm.utils.function_setup(
|
|
"aimage_generation", litellm.utils.Rules(), datetime.now(),
|
|
model="B", prompt="draw a tree", litellm_call_id="supplied-image", **kwargs,
|
|
)
|
|
kwargs.update(litellm_logging_obj=logging_obj, litellm_call_id="supplied-image")
|
|
|
|
context_before: Final = (trace_id_var.get(), session_id_var.get())
|
|
with pytest.raises(ValueError, match="crossed budget"):
|
|
await router.aimage_generation(model="B", prompt="draw a tree", **kwargs)
|
|
assert (trace_id_var.get(), session_id_var.get()) == context_before
|
|
call_type, call_id, exception = await asyncio.to_thread(logger.events.get, timeout=5)
|
|
assert call_type == "aimage_generation" and isinstance(exception, ValueError)
|
|
assert call_id == "supplied-image" if supply_logging else isinstance(call_id, str) and call_id
|
|
assert logger.events.empty()
|
|
assert logger.correlation.get_nowait() == ("refused-trace", "refused-session")
|
|
|
|
|
|
@pytest.mark.parametrize(("model", "status_code"), (("A", 429), ("B", 200)))
|
|
@pytest.mark.parametrize("rewrite_prompt", (False, True))
|
|
def test_image_generation_provider_outcome_is_logged_once(
|
|
monkeypatch: pytest.MonkeyPatch, model: str, status_code: int, rewrite_prompt: bool
|
|
) -> None:
|
|
logger: Final = _RewritingImageLogger() if rewrite_prompt else _ImageOutcomeLogger()
|
|
monkeypatch.setattr(litellm, "callbacks", [logger])
|
|
client, handler = _image_generation_client(monkeypatch, UserAPIKeyAuth())
|
|
with client:
|
|
response: Final = client.post(
|
|
"/v1/images/generations",
|
|
json={"model": model, "prompt": "draw a tree"},
|
|
headers={"x-litellm-call-id": "provider-image"},
|
|
)
|
|
assert response.status_code == status_code, response.text
|
|
call_type, call_id, outcome = logger.events.get(timeout=5)
|
|
assert (call_type, call_id) == ("aimage_generation", "provider-image")
|
|
assert isinstance(outcome, litellm.RateLimitError if model == "A" else ImageResponse)
|
|
assert logger.messages.get_nowait() == [
|
|
{"role": "user", "content": "sanitized prompt" if rewrite_prompt else "draw a tree"}
|
|
]
|
|
assert len(handler.calls) == 1
|
|
assert logger.events.empty()
|
|
|
|
|
|
def test_image_generation_missing_prompt_keeps_its_validation_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
client, handler = _image_generation_client(monkeypatch, UserAPIKeyAuth())
|
|
with client:
|
|
response: Final = client.post("/v1/images/generations", json={"model": "B"})
|
|
assert response.status_code == 400, response.text
|
|
assert response.json()["error"]["param"] == "prompt"
|
|
assert handler.calls == []
|
|
|
|
|
|
def test_image_generation_uses_key_router_settings_fallbacks(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
client, handler = _image_generation_client(
|
|
monkeypatch,
|
|
UserAPIKeyAuth(router_settings={"fallbacks": [{"A": ["B"]}]}),
|
|
)
|
|
|
|
response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"})
|
|
|
|
assert response.status_code == 200, response.text
|
|
assert response.json() == ImageResponse(
|
|
created=1,
|
|
data=[ImageObject(url="https://images.example/generated.png")],
|
|
).model_dump(), response.text
|
|
assert handler.calls == ["failing", "working"], response.text
|
|
|
|
|
|
def test_image_generation_uses_team_router_settings_fallbacks(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
team_id: Final = "image-regression-team"
|
|
client, handler = _image_generation_client(
|
|
monkeypatch,
|
|
UserAPIKeyAuth(team_id=team_id),
|
|
{"fallbacks": [{"A": ["B"]}]},
|
|
)
|
|
|
|
response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"})
|
|
|
|
assert response.status_code == 200, response.text
|
|
assert response.json() == ImageResponse(
|
|
created=1,
|
|
data=[ImageObject(url="https://images.example/generated.png")],
|
|
).model_dump(), response.text
|
|
assert handler.calls == ["failing", "working"], response.text
|
|
|
|
|
|
def test_image_generation_without_router_settings_does_not_use_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
client, handler = _image_generation_client(monkeypatch, UserAPIKeyAuth())
|
|
|
|
response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"})
|
|
|
|
assert response.status_code == 429, response.text
|
|
assert handler.calls == ["failing"], response.text
|
|
|
|
|
|
def test_image_generation_uses_key_model_group_alias(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
client, handler = _image_generation_client(
|
|
monkeypatch,
|
|
UserAPIKeyAuth(router_settings={"model_group_alias": {"alias": "B"}}),
|
|
)
|
|
|
|
response: Final = client.post("/v1/images/generations", json={"model": "alias", "prompt": "draw a tree"})
|
|
|
|
assert response.status_code == 200, response.text
|
|
assert response.json() == ImageResponse(
|
|
created=1,
|
|
data=[ImageObject(url="https://images.example/generated.png")],
|
|
).model_dump(), response.text
|
|
assert handler.calls == ["working"], response.text
|
|
|
|
|
|
def test_image_generation_uses_key_num_retries(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from litellm.proxy import proxy_server
|
|
|
|
logger: Final = _ImageOutcomeLogger()
|
|
monkeypatch.setattr(litellm, "callbacks", [logger])
|
|
client, handler = _image_generation_client(
|
|
monkeypatch,
|
|
UserAPIKeyAuth(router_settings={"num_retries": 2}),
|
|
)
|
|
|
|
with client:
|
|
response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"})
|
|
events: Final = [logger.events.get(timeout=5) for _ in range(3)]
|
|
|
|
assert response.status_code == 429, response.text
|
|
assert handler.calls == ["failing", "failing", "failing"], response.text
|
|
assert all(isinstance(event[2], litellm.RateLimitError) for event in events)
|
|
assert logger.events.empty()
|
|
router: Final = proxy_server.llm_router
|
|
assert router is not None
|
|
assert get_deployment_failures_for_current_minute(router, "image-regression-failing") == 3
|
|
|
|
|
|
def test_failed_image_generation_spend_row_records_aimage_generation(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
|
|
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
|
|
|
spend_writer: Final = MagicMock(update_database=AsyncMock())
|
|
monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger(spend_writer=lambda: spend_writer)])
|
|
client, handler = _image_generation_client(monkeypatch, UserAPIKeyAuth(api_key="sk-test"))
|
|
try:
|
|
with client:
|
|
response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"})
|
|
finally:
|
|
ProxyLogging._callback_capabilities_cache.clear()
|
|
|
|
assert (response.status_code, handler.calls) == (429, ["failing"]), response.text
|
|
db_call: Final = spend_writer.update_database.call_args.kwargs
|
|
payload: Final = get_logging_payload(
|
|
kwargs=db_call["kwargs"],
|
|
response_obj=db_call["completion_response"],
|
|
start_time=db_call["start_time"],
|
|
end_time=db_call["end_time"],
|
|
)
|
|
assert (payload["status"], payload["call_type"]) == ("failure", "aimage_generation")
|