From e9a7d5d7d91c32c784502072bc0f15962ac7412b Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Thu, 30 Jan 2025 16:04:12 +0100
Subject: [PATCH] Fix failure where messages passed to push_trace API were not
 a list of list of dicts.

---
 proxy/routes/proxy.py   | 10 +++++++---
 proxy/utils/explorer.py |  4 +---
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/proxy/routes/proxy.py b/proxy/routes/proxy.py
index b401dba..dddbce3 100644
--- a/proxy/routes/proxy.py
+++ b/proxy/routes/proxy.py
@@ -25,7 +25,7 @@ proxy = APIRouter()
 MISSING_INVARIANT_AUTH_HEADER = "Missing invariant-authorization header"
 MISSING_AUTH_HEADER = "Missing authorization header"
 NOT_SUPPORTED_ENDPOINT = "Not supported OpenAI endpoint"
-FAILED_TO_PUSH_TRACE = "Failed to push trace to the dataset"
+FAILED_TO_PUSH_TRACE = "Failed to push trace to the dataset: "
 
 
 def validate_headers(
@@ -76,11 +76,13 @@ async def openai_proxy(
             ]
             _ = push_trace(
                 dataset_name=dataset_name,
-                messages=messages,
+                messages=[messages],
                 invariant_authorization=request.headers.get("invariant-authorization"),
             )
         except Exception as e:
-            raise HTTPException(status_code=500, detail=FAILED_TO_PUSH_TRACE) from e
+            raise HTTPException(
+                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():
@@ -91,6 +93,8 @@ async def openai_proxy(
             compressed_response = gzip_buffer.getvalue()
 
             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))
 
diff --git a/proxy/utils/explorer.py b/proxy/utils/explorer.py
index fd275f1..252406d 100644
--- a/proxy/utils/explorer.py
+++ b/proxy/utils/explorer.py
@@ -33,6 +33,4 @@ def push_trace(
         )
         return {"trace_id": push_trace_response.id[0]}
     except Exception as e:
-        raise HTTPException(
-            status_code=500, detail="Failed to push traces to the dataset"
-        ) from e
+        raise HTTPException(status_code=500, detail=str(e)) from e