mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(proxy): preserve typed headers without rebinding messages
Signed-off-by: Jaeyeon Kim <anencore94@gmail.com>
This commit is contained in:
parent
dd7d4eba77
commit
eacacd4446
1 changed files with 15 additions and 10 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue