Merge branch 'main' into anthropic-implement

proxy/utils/explorer.py

@@ -3,35 +3,53 @@
 import os
 from typing import Any, Dict, List
 
-from fastapi import HTTPException
-from invariant_sdk.client import Client
+import httpx
+from invariant_sdk.types.push_traces import PushTracesRequest
 
 DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
+PUSH_ENDPOINT = "/api/v1/push/trace"
 
 
-def push_trace(
-    messages: List[Dict[str, Any]],
+async def push_trace(
+    messages: List[List[Dict[str, Any]]],
     dataset_name: str,
     invariant_authorization: str,
 ) -> Dict[str, str]:
     """Pushes traces to the dataset on the Invariant Explorer.
 
     Args:
-        messages (List[Dict[str, Any]]): List of messages to push.
+        messages (List[List[Dict[str, Any]]]): List of messages to push.
         dataset_name (str): Name of the dataset.
-        invariant_authorization (str): Authorization token.
+        invariant_authorization (str): Authorization token from the
+                                       invariant-authorization header.
 
     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:
-        # TODO: Change this to the async version once that is available
-        push_trace_response = client.create_request_and_push_trace(
-            messages=messages, dataset=dataset_name
+    api_url = os.getenv("INVARIANT_API_URL", DEFAULT_API_URL).rstrip("/")
+    # Remove any None values from the messages
+    update_messages = [
+        [{k: v for k, v in msg.items() if v is not None} for msg in msg_list]
+        for msg_list in messages
+    ]
+    request = PushTracesRequest(messages=update_messages, dataset=dataset_name)
+    async with httpx.AsyncClient() as client:
+        explorer_push_request = client.build_request(
+            "POST",
+            f"{api_url}{PUSH_ENDPOINT}",
+            json=request.to_json(),
+            headers={
+                "Authorization": f"{invariant_authorization}",
+                "Accept": "application/json",
+            },
         )
-        return {"trace_id": push_trace_response.id[0]}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e)) from e
+        try:
+            response = await client.send(explorer_push_request)
+            response.raise_for_status()
+            return response.json()
+        except httpx.HTTPStatusError as e:
+            print(f"Failed to push trace: {e.response.text}")
+            return {"error": str(e)}
+        except Exception as e:
+            print(f"Unexpected error pushing trace: {str(e)}")
+            return {"error": str(e)}