From 76644778a1d6b470ca66630e69dce696eaebfa3a Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Fri, 31 Jan 2025 16:00:27 +0100
Subject: [PATCH] Remove custom compression and rely on starlette_compress to
 perform the compression as a middleware.

---
 proxy/requirements.txt  |  3 ++-
 proxy/routes/open_ai.py | 28 ++++++----------------------
 proxy/serve.py          |  2 ++
 3 files changed, 10 insertions(+), 23 deletions(-)

diff --git a/proxy/requirements.txt b/proxy/requirements.txt
index af6f909..6be2e1e 100644
--- a/proxy/requirements.txt
+++ b/proxy/requirements.txt
@@ -1,4 +1,5 @@
 fastapi==0.115.7
 httpx==0.28.1
 uvicorn==0.34.0
-invariant-sdk
\ No newline at end of file
+invariant-sdk
+starlette-compress==1.4.0
\ No newline at end of file
diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py
index b670e60..883bdfc 100644
--- a/proxy/routes/open_ai.py
+++ b/proxy/routes/open_ai.py
@@ -1,8 +1,6 @@
 """Proxy service to forward requests to the OpenAI APIs"""
 
-import gzip
 import json
-from io import BytesIO
 
 import httpx
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response
@@ -84,27 +82,13 @@ async def openai_proxy(
                 status_code=500, detail=FAILED_TO_PUSH_TRACE + str(e)
             ) from e
 
-        # Detect if original request expects gzip encoding
-        if "gzip" in request.headers.get("accept-encoding", "").lower():
-            # Compress the response using gzip
-            gzip_buffer = BytesIO()
-            with gzip.GzipFile(mode="wb", fileobj=gzip_buffer) as gz:
-                gz.write(response.content)
-            compressed_response = gzip_buffer.getvalue()
+        response_headers = dict(response.headers)
+        response_headers.pop("Content-Encoding", None)
+        response_headers.pop("Content-Length", None)
 
-            response_headers = dict(response.headers)
-            response_headers.pop("Content-Encoding", None)
-            response_headers.pop("Content-Length", None)
-            response_headers["Content-Encoding"] = "gzip"
-            response_headers["Content-Length"] = str(len(compressed_response))
-
-            return Response(
-                content=compressed_response,
-                status_code=response.status_code,
-                headers=response_headers,
-            )
         return Response(
-            content=response.content,
+            content=json.dumps(response.json()),
             status_code=response.status_code,
-            headers=dict(response.headers),
+            media_type="application/json",
+            headers=response_headers,
         )
diff --git a/proxy/serve.py b/proxy/serve.py
index dd997e2..b665c55 100644
--- a/proxy/serve.py
+++ b/proxy/serve.py
@@ -4,8 +4,10 @@ import fastapi
 import uvicorn
 from routes.anthropic import proxy as anthropic_proxy
 from routes.open_ai import proxy as open_ai_proxy
+from starlette_compress import CompressMiddleware
 
 app = fastapi.FastAPI()
+app.add_middleware(CompressMiddleware)
 
 router = fastapi.APIRouter(prefix="/api/v1")