diff --git a/litellm/__init__.py b/litellm/__init__.py index 8e5dd68a186..5f882f7df84 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -425,6 +425,7 @@ autorouter_presets_url: str = os.getenv( "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/autorouter_presets.json", ) suppress_debug_info: bool = False +include_litellm_headers: bool = True dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None s3_audit_callback_params: Optional[Dict] = None diff --git a/litellm/proxy/middleware/litellm_response_headers_middleware.py b/litellm/proxy/middleware/litellm_response_headers_middleware.py new file mode 100644 index 00000000000..879e5e58c5f --- /dev/null +++ b/litellm/proxy/middleware/litellm_response_headers_middleware.py @@ -0,0 +1,34 @@ +"""Apply the proxy's opt-out to all HTTP response headers, including streams.""" + +from typing import Final + +from starlette.datastructures import Headers +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +import litellm + + +class LiteLLMResponseHeadersMiddleware: + def __init__(self, app: ASGIApp) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + async def send_response(message: Message) -> None: + if message["type"] != "http.response.start" or litellm.include_litellm_headers is not False: + await send(message) + return + response: Final[Message] = { + **message, + "headers": [ + (name, value) + for name, value in Headers(scope=message).raw + if not name.lower().startswith(b"x-litellm-") + ], + } + await send(response) + + await self.app(scope, receive, send_response) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3297a298af4..a82ea036e65 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -875,6 +875,7 @@ from litellm.proxy.middleware.billable_request_metrics_middleware import ( from litellm.proxy.middleware.budget_reservation_release_middleware import ( BudgetReservationReleaseMiddleware, ) +from litellm.proxy.middleware.litellm_response_headers_middleware import LiteLLMResponseHeadersMiddleware from litellm.proxy.middleware.redis_request_batch_middleware import RedisRequestBatchMiddleware from litellm.proxy.plugin_routes import ( register_plugins_from_config, @@ -22531,6 +22532,8 @@ app.add_middleware( get_settings=lambda: get_admission_control_settings(general_settings), state=admission_control_state, ) +app.add_middleware(LiteLLMResponseHeadersMiddleware) + # Added last on purpose - last-added is outermost, and the client-visible URL # prefix must be resolved into scope["root_path"] before any inner middleware # or the router inspects the path. Only added when SERVER_ROOT_PATHS is diff --git a/tests/unit/proxy/middleware/test_litellm_response_headers_middleware.py b/tests/unit/proxy/middleware/test_litellm_response_headers_middleware.py new file mode 100644 index 00000000000..dbc37849ad5 --- /dev/null +++ b/tests/unit/proxy/middleware/test_litellm_response_headers_middleware.py @@ -0,0 +1,33 @@ +import pytest +from fastapi import FastAPI +from fastapi.responses import StreamingResponse +from fastapi.testclient import TestClient + +import litellm +from litellm.proxy.middleware.litellm_response_headers_middleware import LiteLLMResponseHeadersMiddleware + + +@pytest.mark.parametrize("include", [True, False]) +@pytest.mark.parametrize("status_code", [200, 400]) +@pytest.mark.parametrize("streaming", [True, False]) +def test_header_opt_out_preserves_body_and_other_headers(monkeypatch, include, status_code, streaming): + from starlette.responses import Response + + monkeypatch.setattr(litellm, "include_litellm_headers", include) + app = FastAPI() + app.add_middleware(LiteLLMResponseHeadersMiddleware) + + @app.get("/reply") + def reply(): + headers = {"X-LiteLLM-Call-ID": "call", "x-litellm-priority": "1", "retry-after": "2"} + if streaming: + return StreamingResponse(iter([b"hello", b" world"]), headers=headers, status_code=status_code) + return Response(b"hello world", headers=headers, status_code=status_code) + + with TestClient(app) as client: + response = client.get("/reply") + assert response.status_code == status_code + assert response.content == b"hello world" + assert response.headers["retry-after"] == "2" + actual = {k: v for k, v in response.headers.items() if k.startswith("x-litellm-")} + assert actual == ({"x-litellm-call-id": "call", "x-litellm-priority": "1"} if include else {})