Add metadata before pushing to Explorer.

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("/"),

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)],
     )
 
 

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)],
     )
 
 

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)],
         )