From a6e3ac62a69608ae18c01e7539b3160ade4c64dd Mon Sep 17 00:00:00 2001 From: Praising Harris Ratnam Date: Sat, 17 Oct 2026 13:25:16 +0530 Subject: [PATCH 1/6] fix(proxy): initialize litellm_logging_obj for image_generation to dispatch router failure callbacks (#47061) --- litellm/proxy/image_endpoints/endpoints.py | 11 +++++++++++ tests/unit/proxy/image_endpoints/test_endpoints.py | 1 + 2 files changed, 12 insertions(+) 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" From 808cbb04f13bbb40e9a9553d65cb192998eb1863 Mon Sep 17 00:00:00 2001 From: Praising Harris Ratnam Date: Sun, 18 Oct 2026 20:13:59 +0530 Subject: [PATCH 2/6] fix(proxy): sync sanitized prompt to logging_obj, align start_time, and test router failure dispatch --- litellm/proxy/image_endpoints/endpoints.py | 10 +- .../proxy/image_endpoints/test_endpoints.py | 97 ++++++++++++++++++- 2 files changed, 104 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 84670f8c6dd..13bf169f275 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -103,6 +103,7 @@ async def image_generation( version, ) + start_time: Final = datetime.now() litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) data = {"litellm_call_id": litellm_call_id} try: @@ -134,7 +135,6 @@ async def image_generation( 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(), @@ -158,7 +158,13 @@ async def image_generation( messages: Final = data.get("messages") if isinstance(messages, list) and messages: - data["prompt"] = get_str_from_messages(messages) + sanitized_prompt: Final = get_str_from_messages(messages) + data["prompt"] = sanitized_prompt + if logging_obj is not None: + logging_obj.messages = sanitized_prompt + if hasattr(logging_obj, "model_call_details") and isinstance(logging_obj.model_call_details, dict): + logging_obj.model_call_details["prompt"] = sanitized_prompt + logging_obj.model_call_details["messages"] = sanitized_prompt data.pop("messages", None) ## ROUTE TO CORRECT ENDPOINT ## diff --git a/tests/unit/proxy/image_endpoints/test_endpoints.py b/tests/unit/proxy/image_endpoints/test_endpoints.py index 51502cbec03..7ee3db07db0 100644 --- a/tests/unit/proxy/image_endpoints/test_endpoints.py +++ b/tests/unit/proxy/image_endpoints/test_endpoints.py @@ -116,10 +116,105 @@ 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 + 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, MagicMock + + 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(), + ) + + # Mock router whose async_get_available_deployment rejects the request + mock_router = MagicMock() + mock_failure_dispatch = 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: From ee71aa92e828002be3c5e1a6a44b7f6444e0ed9e Mon Sep 17 00:00:00 2001 From: Praising Harris Ratnam Date: Sun, 18 Oct 2026 23:11:28 +0530 Subject: [PATCH 3/6] style: fix ruff formatting and clean unused mock vars in image endpoints --- litellm/proxy/image_endpoints/endpoints.py | 2 +- .../proxy/image_endpoints/test_endpoints.py | 31 ++++++++----------- 2 files changed, 14 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 13bf169f275..c26ee8bcacf 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -1,7 +1,7 @@ import asyncio -from datetime import datetime import io from collections.abc import Sequence +from datetime import datetime from itertools import chain from types import MappingProxyType from typing import Final, get_type_hints diff --git a/tests/unit/proxy/image_endpoints/test_endpoints.py b/tests/unit/proxy/image_endpoints/test_endpoints.py index 7ee3db07db0..2dcf6a8781d 100644 --- a/tests/unit/proxy/image_endpoints/test_endpoints.py +++ b/tests/unit/proxy/image_endpoints/test_endpoints.py @@ -3,7 +3,7 @@ import copy import logging from collections.abc import Iterator, Mapping from types import SimpleNamespace -from typing import Any, Dict +from typing import Any import orjson import pytest @@ -28,7 +28,7 @@ async def test_image_generation_prompt_rerouting(monkeypatch): async def fake_update_request_status(**_: Any) -> None: await asyncio.sleep(0) - proxy_logger_calls: Dict[str, Any] = {} + 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) @@ -60,7 +60,7 @@ async def test_image_generation_prompt_rerouting(monkeypatch): post_call_response_headers_hook=fake_post_call_response_headers_hook, ) - captured_route_request_data: Dict[str, Any] = {} + captured_route_request_data: dict[str, Any] = {} async def fake_route_request(*, data, **kwargs): # type: ignore[override] captured_route_request_data.update(data) @@ -126,7 +126,7 @@ async def test_image_generation_prompt_rerouting(monkeypatch): @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, MagicMock + from unittest.mock import AsyncMock async def fake_add_litellm_data_to_request(**kwargs): return kwargs["data"] @@ -136,10 +136,6 @@ async def test_image_generation_router_rejection_dispatches_failure_callback(mon post_call_failure_hook=AsyncMock(), ) - # Mock router whose async_get_available_deployment rejects the request - mock_router = MagicMock() - mock_failure_dispatch = AsyncMock() - class RejectingRouter: model_names = ["dall-e-3"] team_public_model_names = [] @@ -214,13 +210,12 @@ async def test_image_generation_router_rejection_dispatches_failure_callback(mon assert str(dispatched_exceptions[0]) == "No deployment available" - -def _image_edit_client(monkeypatch, captured: Dict[str, Any]) -> TestClient: +def _image_edit_client(monkeypatch, captured: dict[str, Any]) -> TestClient: class CaptureProcessing: - def __init__(self, data: Dict[str, Any]) -> None: + def __init__(self, data: dict[str, Any]) -> None: captured.update(data) - async def base_process_llm_request(self, **_: Any) -> Dict[str, Any]: + async def base_process_llm_request(self, **_: Any) -> dict[str, Any]: return {"data": [{"b64_json": "aGk="}]} monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", CaptureProcessing) @@ -235,7 +230,7 @@ def _image_edit_client(monkeypatch, captured: Dict[str, Any]) -> TestClient: 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] = {} + captured: dict[str, Any] = {} response = _image_edit_client(monkeypatch, captured).post( "/v1/images/edits", @@ -251,7 +246,7 @@ def test_image_edit_image_array_alias_is_not_forwarded(monkeypatch): 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] = {} + captured: dict[str, Any] = {} response = _image_edit_client(monkeypatch, captured).post( "/v1/images/edits", @@ -270,7 +265,7 @@ def test_image_edit_mask_array_alias_is_not_forwarded(monkeypatch): 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] = {} + captured: dict[str, Any] = {} response = _image_edit_client(monkeypatch, captured).post( "/v1/images/edits", @@ -289,7 +284,7 @@ def test_image_edit_canonical_file_fields_still_reach_the_provider(monkeypatch): 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] = {} + captured: dict[str, Any] = {} response = _image_edit_client(monkeypatch, captured).post( "/v1/images/edits", @@ -306,7 +301,7 @@ def test_image_edit_multipart_n_reaches_the_provider_as_an_int(monkeypatch): 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] = {} + captured: dict[str, Any] = {} response = _image_edit_client(monkeypatch, captured).post( "/v1/images/edits", @@ -332,7 +327,7 @@ def test_image_edit_multipart_n_that_is_not_a_number_is_left_alone(monkeypatch): def test_image_edit_without_an_optional_field_reaches_the_provider_with_it_set_to_none( monkeypatch, files, form, missing ): - captured: Dict[str, Any] = {} + captured: dict[str, Any] = {} response = _image_edit_client(monkeypatch, captured).post("/v1/images/edits", files=files or None, data=form) From 11493298f21c0ff0d66c99133e2c3df59622db0d Mon Sep 17 00:00:00 2001 From: Praising Harris Ratnam Date: Mon, 19 Oct 2026 01:54:14 +0530 Subject: [PATCH 4/6] fix(image_endpoints): add noqa DTZ005 for start_time naive datetime logging --- litellm/proxy/image_endpoints/endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index c26ee8bcacf..dcbc2222d1c 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -103,7 +103,7 @@ async def image_generation( version, ) - start_time: Final = datetime.now() + start_time: Final = datetime.now() # noqa: DTZ005 # proxy standard naive datetime for request duration logging litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) data = {"litellm_call_id": litellm_call_id} try: From 264fb2eb1ab3a83e0be35b6b9b6a795663f61585 Mon Sep 17 00:00:00 2001 From: Pyhroff Date: Fri, 23 Oct 2026 12:25:05 +0530 Subject: [PATCH 5/6] fix(proxy): annotate function_setup rebind and update test_img_gen mock expectation --- litellm/proxy/image_endpoints/endpoints.py | 2 +- tests/unit/proxy/test_proxy_server.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index dcbc2222d1c..0be72dd6f1a 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -135,7 +135,7 @@ async def image_generation( data["model"] = litellm.model_alias_map[data["model"]] ### LOGGING OBJECT SETUP ### - ensures router rejected requests are dispatched to failure callbacks - logging_obj, data = litellm.utils.function_setup( + logging_obj, data = litellm.utils.function_setup( # rebind-ok: function_setup returns updated data dict original_function="aimage_generation", rules_obj=litellm.utils.Rules(), start_time=start_time, diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 2f2563b9ee4..c4d36b68149 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -1004,6 +1004,7 @@ def test_img_gen(mock_aimage_generation, client_no_auth): size="1024x1024", imageConfig={"aspectRatio": "9:16", "imageSize": "1K"}, litellm_call_id=mock.ANY, + litellm_logging_obj=mock.ANY, metadata=mock.ANY, proxy_server_request=mock.ANY, secret_fields=mock.ANY, From e661bb7abc2c90a3374c15e77e8a07d3a59f6dcf Mon Sep 17 00:00:00 2001 From: Pyhroff Date: Fri, 23 Oct 2026 12:34:40 +0530 Subject: [PATCH 6/6] style(proxy): satisfy basedpyright strict checks in image_generation --- litellm/proxy/image_endpoints/endpoints.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 0be72dd6f1a..26bd30560b5 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -139,7 +139,7 @@ async def image_generation( original_function="aimage_generation", rules_obj=litellm.utils.Rules(), start_time=start_time, - **data, + **data, # pyright: ignore[reportUnknownArgumentType] # request payload dict carries dynamic values ) data["litellm_logging_obj"] = logging_obj @@ -160,11 +160,9 @@ async def image_generation( if isinstance(messages, list) and messages: sanitized_prompt: Final = get_str_from_messages(messages) data["prompt"] = sanitized_prompt - if logging_obj is not None: - logging_obj.messages = sanitized_prompt - if hasattr(logging_obj, "model_call_details") and isinstance(logging_obj.model_call_details, dict): - logging_obj.model_call_details["prompt"] = sanitized_prompt - logging_obj.model_call_details["messages"] = sanitized_prompt + logging_obj.messages = sanitized_prompt + logging_obj.model_call_details["prompt"] = sanitized_prompt + logging_obj.model_call_details["messages"] = sanitized_prompt data.pop("messages", None) ## ROUTE TO CORRECT ENDPOINT ##