scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 267

fix(proxy): limit gzip to dashboard analytics and offload compression (#46712)

Browse source 
Backport of #46712 to stable/1.105.x.
Cherry-picked from 9c62573005 (litellm_dashboard_only_gzip). The line never took #46362, so the OFF_LOOP_SIZE_BYTES hunks resolve to the pick, which removes the constant.
This commit is contained in:
tin-berri 2026-10-27 23:52:05 +00:00 • committed by devin-ai-integration[bot]
parent 3a05120f10
commit a855a6927b
2 changed files with 182 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
litellm/proxy/middleware/gzip_middleware.py
View file

@ -4,11 +4,50 @@ from typing import Final
import anyio.to_thread
from starlette.datastructures import Headers, MutableHeaders
from starlette.routing import get_route_path
from starlette.types import ASGIApp, Message, Receive, Scope, Send
MINIMUM_SIZE_BYTES: Final = 500
OFF_LOOP_SIZE_BYTES: Final = 1024 * 1024
COMPRESS_LEVEL: Final = 6
DASHBOARD_GZIP_PATHS: Final = frozenset(
{
"/user/daily/activity/aggregated",
"/user/daily/activity/aggregated/keys",
"/user/daily/activity/aggregated/search",
"/user/daily/activity/aggregated/model_top_keys",
"/user/daily/activity/aggregated/users",
"/user/daily/activity/aggregated/cache_leakage_keys",
"/team/daily/activity/aggregated",
"/team/daily/activity/aggregated/keys",
"/team/daily/activity/aggregated/search",
"/team/daily/activity/aggregated/model_top_keys",
"/organization/daily/activity/aggregated",
"/organization/daily/activity/aggregated/keys",
"/organization/daily/activity/aggregated/search",
"/organization/daily/activity/aggregated/model_top_keys",
"/tag/daily/activity/aggregated",
"/tag/daily/activity/aggregated/keys",
"/tag/daily/activity/aggregated/search",
"/tag/daily/activity/aggregated/model_top_keys",
"/customer/daily/activity/aggregated",
"/customer/daily/activity/aggregated/keys",
"/customer/daily/activity/aggregated/search",
"/customer/daily/activity/aggregated/model_top_keys",
"/end_user/daily/activity/aggregated",
"/end_user/daily/activity/aggregated/keys",
"/end_user/daily/activity/aggregated/search",
"/end_user/daily/activity/aggregated/model_top_keys",
"/agent/daily/activity/aggregated",
"/agent/daily/activity/aggregated/keys",
"/agent/daily/activity/aggregated/search",
"/agent/daily/activity/aggregated/model_top_keys",
"/gateway/daily/activity",
"/gateway/errors/activity",
"/team/spend/by_user",
"/global/spend/provider",
"/auto_router/benchmarks",
}
)
def _coding_weight(part: str) -> tuple[str, float]:
@ -26,8 +65,6 @@ def accepts_gzip(accept_encoding: str) -> bool:
async def _compress(body: bytes) -> bytes:
if len(body) < OFF_LOOP_SIZE_BYTES:
return gzip.compress(body, compresslevel=COMPRESS_LEVEL)
return await anyio.to_thread.run_sync(gzip.compress, body, COMPRESS_LEVEL)
@ -83,11 +120,12 @@ class _BufferedBodyGzipResponder:
class GZipBufferedResponseMiddleware:
def __init__(self, app: ASGIApp) -> None:
def __init__(self, app: ASGIApp, paths: frozenset[str] = DASHBOARD_GZIP_PATHS) -> None:
self.app = app
self.paths: Final = paths
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
if scope["type"] != "http" or scope["method"] != "GET" or get_route_path(scope) not in self.paths:
await self.app(scope, receive, send)
return
gzip_accepted: Final = accepts_gzip(Headers(scope=scope).get("accept-encoding", ""))

156
tests/unit/proxy/middleware/test_gzip_middleware.py
View file

@ -1,6 +1,7 @@
import asyncio
import gzip
import json
import threading
from typing import Final
import pytest
@ -8,15 +9,19 @@ from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route
from starlette.testclient import TestClient
from starlette.types import ASGIApp, Message, Receive, Scope, Send
from litellm.proxy.middleware import gzip_middleware
from litellm.proxy.middleware.gzip_middleware import (
DASHBOARD_GZIP_PATHS,
MINIMUM_SIZE_BYTES,
OFF_LOOP_SIZE_BYTES,
GZipBufferedResponseMiddleware,
)
LARGE_PAYLOAD = {"rows": [{"date": f"2026-09-{day:02d}", "spend": day * 1.5} for day in range(1, 31)] * 20}
BASE64_EMBEDDINGS_RESPONSE_BYTES: Final = 824_000
HUGE_RESPONSE_BYTES: Final = 2 * 1024 * 1024
STREAM_CHUNKS = tuple(json.dumps({"part": part, "pad": "x" * MINIMUM_SIZE_BYTES}).encode() for part in range(3))
@ -45,7 +50,19 @@ async def _no_transform(request: Request) -> Response:
async def _huge(request: Request) -> Response:
return Response(b"z" * (OFF_LOOP_SIZE_BYTES * 2), media_type="application/json")
return Response(b"z" * HUGE_RESPONSE_BYTES, media_type="application/json")
async def _minimum_eligible(request: Request) -> Response:
return Response(b"u" * MINIMUM_SIZE_BYTES, media_type="application/json")
async def _four_kib(request: Request) -> Response:
return Response(b"a" * 4096, media_type="application/json")
async def _base64_embeddings_sized(request: Request) -> Response:
return Response(b"e" * BASE64_EMBEDDINGS_RESPONSE_BYTES, media_type="application/json")
async def _json_stream(request: Request) -> Response:
@ -64,16 +81,30 @@ APP = Starlette(
Route("/stream", _json_stream),
Route("/etag", _with_etag),
Route("/huge", _huge),
Route("/minimum-eligible", _minimum_eligible),
Route("/four-kib", _four_kib),
Route("/base64-embeddings-sized", _base64_embeddings_sized),
Route("/partial", _partial),
Route("/no-transform", _no_transform),
]
)
APP.add_middleware(GZipBufferedResponseMiddleware)
APP.add_middleware(
GZipBufferedResponseMiddleware, paths=frozenset(route.path for route in APP.routes if isinstance(route, Route))
)
async def _send_messages(path: str, accept_encoding: str | None, app: ASGIApp = APP) -> tuple[Message, ...]:
async def _send_messages(
path: str, accept_encoding: str | None, app: ASGIApp = APP, method: str = "GET", root_path: str = ""
) -> tuple[Message, ...]:
headers = [(b"accept-encoding", accept_encoding.encode())] if accept_encoding is not None else []
scope = {"type": "http", "method": "GET", "path": path, "query_string": b"", "headers": headers}
scope = {
"type": "http",
"method": method,
"path": path,
"root_path": root_path,
"query_string": b"",
"headers": headers,
}
sent: list[Message] = [] # mutable-ok: ASGI send callback collects messages in order
requests: Final = iter(({"type": "http.request", "body": b"", "more_body": False},))
never_disconnects: Final = asyncio.Event()
@ -115,11 +146,31 @@ async def test_large_buffered_json_is_gzipped_and_round_trips(accept_encoding):
@pytest.mark.asyncio
async def test_body_above_off_loop_threshold_round_trips():
async def test_multi_megabyte_body_round_trips():
messages = await _send_messages("/huge", "gzip")
assert _headers(messages)["content-encoding"] == "gzip"
assert gzip.decompress(_body(messages)) == b"z" * (OFF_LOOP_SIZE_BYTES * 2)
assert gzip.decompress(_body(messages)) == b"z" * HUGE_RESPONSE_BYTES
@pytest.mark.parametrize(
"path",
["/minimum-eligible", "/four-kib", "/base64-embeddings-sized", "/huge"],
)
@pytest.mark.asyncio
async def test_all_eligible_bodies_compress_on_a_worker_thread(monkeypatch: pytest.MonkeyPatch, path: str) -> None:
loop_thread: Final = threading.get_ident()
real_compress: Final = gzip.compress
def compress_naming_its_thread(data: bytes, compresslevel: int = 9) -> bytes:
return real_compress(b"loop" if threading.get_ident() == loop_thread else b"worker", compresslevel)
monkeypatch.setattr(gzip_middleware.gzip, "compress", compress_naming_its_thread)
messages = await _send_messages(path, "gzip")
assert _headers(messages)["content-encoding"] == "gzip"
assert gzip.decompress(_body(messages)) == b"worker"
@pytest.mark.parametrize(
@ -183,7 +234,9 @@ async def test_start_message_without_headers_key_is_still_gzipped():
await send({"type": "http.response.start", "status": 200})
await send({"type": "http.response.body", "body": body})
messages = await _send_messages("/", "gzip", GZipBufferedResponseMiddleware(headerless_app))
messages = await _send_messages(
"/user/daily/activity/aggregated", "gzip", GZipBufferedResponseMiddleware(headerless_app)
)
assert _headers(messages)["content-encoding"] == "gzip"
assert gzip.decompress(_body(messages)) == body
@ -194,20 +247,89 @@ async def test_start_without_a_body_message_is_still_forwarded():
async def start_only_app(scope: Scope, receive: Receive, send: Send) -> None:
await send({"type": "http.response.start", "status": 204, "headers": [(b"x-done", b"1")]})
messages = await _send_messages("/", "gzip", GZipBufferedResponseMiddleware(start_only_app))
messages = await _send_messages(
"/user/daily/activity/aggregated", "gzip", GZipBufferedResponseMiddleware(start_only_app)
)
assert messages == ({"type": "http.response.start", "status": 204, "headers": [(b"x-done", b"1")]},)
def test_proxy_app_gzips_large_responses_for_clients_that_accept_it():
from starlette.testclient import TestClient
@pytest.mark.parametrize("path", sorted(DASHBOARD_GZIP_PATHS))
@pytest.mark.parametrize("root_path", ["", "/gateway"])
@pytest.mark.asyncio
async def test_dashboard_reads_are_compressed_under_root_paths(path: str, root_path: str) -> None:
async def dashboard(scope: Scope, receive: Receive, send: Send) -> None:
await JSONResponse(LARGE_PAYLOAD)(scope, receive, send)
messages: Final = await _send_messages(
root_path + path, "gzip", GZipBufferedResponseMiddleware(dashboard), root_path=root_path
)
assert _headers(messages)["content-encoding"] == "gzip"
assert json.loads(gzip.decompress(_body(messages))) == LARGE_PAYLOAD
@pytest.mark.parametrize("suffix", ["", "/keys", "/search", "/model_top_keys"])
@pytest.mark.asyncio
async def test_customer_analytics_aliases_have_equivalent_compressed_responses(suffix: str) -> None:
async def analytics(scope: Scope, receive: Receive, send: Send) -> None:
await JSONResponse(LARGE_PAYLOAD)(scope, receive, send)
app: Final = GZipBufferedResponseMiddleware(analytics)
customer: Final = await _send_messages(f"/customer/daily/activity/aggregated{suffix}", "gzip", app)
end_user: Final = await _send_messages(f"/end_user/daily/activity/aggregated{suffix}", "gzip", app)
assert _headers(customer)["content-encoding"] == _headers(end_user)["content-encoding"] == "gzip"
assert _headers(customer)["vary"] == _headers(end_user)["vary"] == "Accept-Encoding"
assert json.loads(gzip.decompress(_body(customer))) == json.loads(gzip.decompress(_body(end_user))) == LARGE_PAYLOAD
@pytest.mark.parametrize(
"path,method",
[
("/v1/embeddings", "POST"),
("/embeddings", "POST"),
("/v1/chat/completions", "POST"),
("/v1/responses", "POST"),
("/v1/messages", "POST"),
("/openai/deployments/embedding-model/embeddings", "POST"),
("/custom-passthrough", "GET"),
("/openapi.json", "GET"),
("/health/liveliness", "GET"),
("/end_user/daily/activity/export", "GET"),
("/user/daily/activity/aggregated/unknown", "GET"),
("/prefix/user/daily/activity/aggregated", "GET"),
("/user/daily/activity/aggregated", "POST"),
("/user/daily/activity/aggregated", "HEAD"),
],
)
@pytest.mark.asyncio
async def test_non_dashboard_requests_pass_through_without_changing_headers(path: str, method: str) -> None:
body: Final = b"e" * BASE64_EMBEDDINGS_RESPONSE_BYTES
async def inference(scope: Scope, receive: Receive, send: Send) -> None:
await Response(body, media_type="application/json", headers={"vary": "Origin"})(scope, receive, send)
messages: Final = await _send_messages(path, "gzip", GZipBufferedResponseMiddleware(inference), method=method)
assert _headers(messages) == {
"content-type": "application/json",
"content-length": str(len(body)),
"vary": "Origin",
}
assert _body(messages) == body
def test_proxy_app_does_not_compress_openapi() -> None:
from litellm.proxy.proxy_server import app
client = TestClient(app)
compressed = client.get("/openapi.json", headers={"accept-encoding": "gzip"})
identity = client.get("/openapi.json", headers={"accept-encoding": "identity"})
client: Final = TestClient(app)
gzip_accepted: Final = client.get("/openapi.json", headers={"accept-encoding": "gzip"})
identity: Final = client.get("/openapi.json", headers={"accept-encoding": "identity"})
assert compressed.headers["content-encoding"] == "gzip"
assert int(compressed.headers["content-length"]) < int(identity.headers["content-length"])
assert compressed.json() == identity.json()
assert gzip_accepted.status_code == identity.status_code == 200, (
f"gzip={gzip_accepted.status_code} {gzip_accepted.text[:MINIMUM_SIZE_BYTES]}, "
f"identity={identity.status_code} {identity.text[:MINIMUM_SIZE_BYTES]}"
)
assert "content-encoding" not in gzip_accepted.headers
assert gzip_accepted.content == identity.content