diff --git a/litellm/proxy/middleware/litellm_response_headers_middleware.py b/litellm/proxy/middleware/litellm_response_headers_middleware.py index 620c63f8715..879e5e58c5f 100644 --- a/litellm/proxy/middleware/litellm_response_headers_middleware.py +++ b/litellm/proxy/middleware/litellm_response_headers_middleware.py @@ -1,5 +1,8 @@ """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 @@ -15,15 +18,17 @@ class LiteLLMResponseHeadersMiddleware: return async def send_response(message: Message) -> None: - if message["type"] == "http.response.start" and litellm.include_litellm_headers is False: - message = { - **message, - "headers": [ - (name, value) - for name, value in message.get("headers", []) - if not name.lower().startswith(b"x-litellm-") - ], - } - await send(message) + 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)