diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 28fb9baf2aa..84670f8c6dd 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -1,4 +1,5 @@ import asyncio +from datetime import datetime import io from collections.abc import Sequence from itertools import chain @@ -132,6 +133,16 @@ async def image_generation( if data["model"] in litellm.model_alias_map: data["model"] = litellm.model_alias_map[data["model"]] + ### LOGGING OBJECT SETUP ### - ensures router rejected requests are dispatched to failure callbacks + start_time: Final = datetime.now() + logging_obj, data = litellm.utils.function_setup( + original_function="aimage_generation", + rules_obj=litellm.utils.Rules(), + start_time=start_time, + **data, + ) + data["litellm_logging_obj"] = logging_obj + ### CALL HOOKS ### - modify incoming data / reject request before calling the model prompt_value: Final = data.get("prompt") if prompt_value is not None: diff --git a/tests/unit/proxy/image_endpoints/test_endpoints.py b/tests/unit/proxy/image_endpoints/test_endpoints.py index f4aebecc11e..51502cbec03 100644 --- a/tests/unit/proxy/image_endpoints/test_endpoints.py +++ b/tests/unit/proxy/image_endpoints/test_endpoints.py @@ -116,6 +116,7 @@ async def test_image_generation_prompt_rerouting(monkeypatch): 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 captured_route_request_data.get("litellm_logging_obj") is not None assert response.headers.get("x-callback-test") == "value"