Merge eacacd4446 into e9b1c11098

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

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)

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

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 {})