From a4ee5f0c4737e3d91fba42845deadc403ac26378 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Thu, 20 Mar 2025 09:25:28 +0100
Subject: [PATCH] Add metadata before pushing to Explorer.

---
 gateway/integrations/explorer.py |  6 +++++-
 gateway/routes/anthropic.py      | 12 ++++++++++++
 gateway/routes/gemini.py         | 21 +++++++++++++++++++++
 gateway/routes/open_ai.py        | 21 +++++++++++++++++++++
 4 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/gateway/integrations/explorer.py b/gateway/integrations/explorer.py
index b61e610..ed4098d 100644
--- a/gateway/integrations/explorer.py
+++ b/gateway/integrations/explorer.py
@@ -33,6 +33,7 @@ async def push_trace(
     dataset_name: str,
     invariant_authorization: str,
     annotations: List[List[AnnotationCreate]] = None,
+    metadata: List[Dict[str, Any]] = None,
 ) -> PushTracesResponse:
     """Pushes traces to the dataset on the Invariant Explorer.
 
@@ -53,7 +54,10 @@ async def push_trace(
         for msg_list in messages
     ]
     request = PushTracesRequest(
-        messages=update_messages, annotations=annotations, dataset=dataset_name
+        messages=update_messages,
+        annotations=annotations,
+        dataset=dataset_name,
+        metadata=metadata,
     )
     client = AsyncClient(
         api_url=os.getenv("INVARIANT_API_URL", DEFAULT_API_URL).rstrip("/"),
diff --git a/gateway/routes/anthropic.py b/gateway/routes/anthropic.py
index 2b04a6b..c6716a4 100644
--- a/gateway/routes/anthropic.py
+++ b/gateway/routes/anthropic.py
@@ -90,6 +90,17 @@ async def anthropic_v1_messages_gateway(
     return await handle_non_streaming_response(context, response)
 
 
+def create_metadata(
+    context: RequestContextData, response_json: dict[str, Any]
+) -> dict[str, Any]:
+    """Creates metadata for the trace"""
+    metadata = {k: v for k, v in context.request_json.items() if k != "messages"}
+    metadata["via_gateway"] = True
+    if response_json.get("usage"):
+        metadata["usage"] = response_json.get("usage")
+    return metadata
+
+
 async def push_to_explorer(
     context: RequestContextData,
     merged_response: dict[str, Any],
@@ -104,6 +115,7 @@ async def push_to_explorer(
         dataset_name=context.dataset_name,
         messages=[converted_messages],
         invariant_authorization=context.invariant_authorization,
+        metadata=[create_metadata(context, merged_response)],
     )
 
 
diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py
index 7ff67e8..0905146 100644
--- a/gateway/routes/gemini.py
+++ b/gateway/routes/gemini.py
@@ -184,6 +184,26 @@ def update_merged_response(merged_response: dict[str, Any], chunk_json: dict) ->
             merged_response["candidates"][0]["finishReason"] = candidate["finishReason"]
 
 
+def create_metadata(
+    context: RequestContextData, response_json: dict[str, Any]
+) -> dict[str, Any]:
+    """Creates metadata for the trace"""
+    metadata = {
+        k: v
+        for k, v in context.request_json.items()
+        if k not in ("systemInstruction", "contents")
+    }
+    metadata["via_gateway"] = True
+    metadata.update(
+        {
+            key: value
+            for key, value in response_json.items()
+            if key in ("usageMetadata", "modelVersion")
+        }
+    )
+    return metadata
+
+
 async def push_to_explorer(
     context: RequestContextData,
     response_json: dict[str, Any],
@@ -195,6 +215,7 @@ async def push_to_explorer(
         dataset_name=context.dataset_name,
         messages=[converted_requests + converted_responses],
         invariant_authorization=context.invariant_authorization,
+        metadata=[create_metadata(context, response_json)],
     )
 
 
diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py
index ce04c95..e464ec3 100644
--- a/gateway/routes/open_ai.py
+++ b/gateway/routes/open_ai.py
@@ -315,6 +315,26 @@ def update_existing_choice_with_delta(
         existing_choice["finish_reason"] = finish_reason
 
 
+def create_metadata(
+    context: RequestContextData, merged_response: dict[str, Any]
+) -> dict[str, Any]:
+    """Creates metadata for the trace"""
+    metadata = {
+        k: v
+        for k, v in context.request_json.items()
+        if k != "messages" and v is not None
+    }
+    metadata["via_gateway"] = True
+    metadata.update(
+        {
+            key: value
+            for key, value in merged_response.items()
+            if key in ("usage", "model")
+        }
+    )
+    return metadata
+
+
 async def push_to_explorer(
     context: RequestContextData,
     merged_response: dict[str, Any],
@@ -339,6 +359,7 @@ async def push_to_explorer(
             invariant_authorization=context.invariant_authorization,
             messages=[messages],
             annotations=[annotations],
+            metadata=[create_metadata(context, merged_response)],
         )