Add explorer.py to push traces to explorer and cleanup code a bit.

.env

@@ -0,0 +1 @@
+INVARIANT_API_URL=https://explorer.invariantlabs.ai

proxy/routes/proxy.py

@@ -22,17 +22,20 @@ IGNORED_HEADERS = [
 ]
 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"
+
 
 def validate_headers(
     invariant_authorization: str = Header(None), authorization: str = Header(None)
 ):
     """Require the invariant-authorization and authorization headers to be present"""
     if invariant_authorization is None:
-        raise HTTPException(
-            status_code=400, detail="Missing invariant-authorization header"
-        )
+        raise HTTPException(status_code=400, detail=MISSING_INVARIANT_AUTH_HEADER)
     if authorization is None:
-        raise HTTPException(status_code=400, detail="Missing authorization header")
+        raise HTTPException(status_code=400, detail=MISSING_AUTH_HEADER)
 
 
 @proxy.post(
@@ -46,13 +49,11 @@ async def openai_proxy(
 ):
     """Proxy call to a language model provider"""
     if endpoint not in ALLOWED_OPEN_AI_ENDPOINTS:
-        raise HTTPException(status_code=404, detail="Not supported OpenAI endpoint")
+        raise HTTPException(status_code=404, detail=NOT_SUPPORTED_ENDPOINT)
 
-    headers = dict(request.headers)
-
-    # Remove extra headers
-    for h in IGNORED_HEADERS:
-        headers.pop(h, None)
+    headers = {
+        k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS
+    }
     headers["accept-encoding"] = "identity"
 
     request_body = await request.body()
@@ -67,28 +68,22 @@ async def openai_proxy(
         response = await client.send(open_ai_request)
         try:
             json_response = response.json()
+            # push messages to the Invariant Explorer
+            # use both the request and response messages
+            messages = json.loads(request_body).get("messages", [])
+            messages += [
+                choice["message"] for choice in json_response.get("choices", [])
+            ]
             _ = push_trace(
                 dataset_name=dataset_name,
-                # Extract messages from the request body and response
-                messages=[
-                    json.loads(request_body)["messages"]
-                    + [
-                        dict(choice["message"])
-                        for choice in json_response.get("choices", [])
-                    ]
-                ],
+                messages=messages,
                 invariant_authorization=request.headers.get("invariant-authorization"),
             )
         except Exception as e:
-            raise HTTPException(
-                status_code=500, detail="Failed to push trace to the dataset"
-            ) from e
+            raise HTTPException(status_code=500, detail=FAILED_TO_PUSH_TRACE) from e
 
         # Detect if original request expects gzip encoding
-        original_accept_encoding = request.headers.get("accept-encoding", "")
-        should_gzip = "gzip" in original_accept_encoding.lower()
-
-        if should_gzip:
+        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:

proxy/utils/explorer.py

@@ -0,0 +1,40 @@
+"""Utility functions for the Invariant explorer."""
+
+import os
+from typing import Any, Dict, List
+
+from fastapi import HTTPException
+from invariant_sdk.client import Client
+
+DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
+
+
+def push_trace(
+    messages: List[Dict[str, Any]],
+    dataset_name: str,
+    invariant_authorization: str,
+    api_url: str = DEFAULT_API_URL,
+) -> Dict[str, str]:
+    """Pushes traces to the dataset on the Invariant Explorer.
+
+    Args:
+        messages (List[Dict[str, Any]]): List of messages to push.
+        dataset_name (str): Name of the dataset.
+        invariant_authorization (str): Authorization token.
+        api_url (str): URL of the Invariant Explorer API.
+
+    Returns:
+        Dict[str, str]: Response containing the trace ID.
+    """
+    api_url = os.getenv("INVARIANT_API_URL", DEFAULT_API_URL)
+    api_key = invariant_authorization.split("Bearer ")[1]
+    client = Client(api_url=api_url, api_key=api_key)
+    try:
+        push_trace_response = client.create_request_and_push_trace(
+            messages=messages, dataset=dataset_name
+        )
+        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